From 7d8e1dfdd7cf0f50cd43296b1d0e57c78f6fd541 Mon Sep 17 00:00:00 2001
From: Andrea Cracco <guilucand@gmail.com>
Date: Thu, 10 Oct 2024 10:43:32 +0200
Subject: [PATCH] Fix colors encoding when kmer-counters flag is enabled, fixes
 #54

---
 crates/colors/src/managers/multiple.rs | 24 ++++++++++++++++++++----
 1 file changed, 20 insertions(+), 4 deletions(-)

diff --git a/crates/colors/src/managers/multiple.rs b/crates/colors/src/managers/multiple.rs
index 2317c61..bdaf6a6 100644
--- a/crates/colors/src/managers/multiple.rs
+++ b/crates/colors/src/managers/multiple.rs
@@ -120,6 +120,22 @@ const VISITED_BIT: usize = 1 << (COUNTER_BITS - 1);
 const TEMP_BUFFER_START_SIZE: usize = 1024 * 64;
 const READS_BUFFERS_MAX_CAPACITY: usize = 1024 * 32;
 
+#[cfg(feature = "support_kmer_counters")]
+type HashMapTempColorIndex = usize;
+
+#[cfg(not(feature = "support_kmer_counters"))]
+type HashMapTempColorIndex = ();
+
+#[inline]
+fn get_entry_color(entry: &MapEntry<HashMapTempColorIndex>) -> ColorIndexType {
+    (match () {
+        #[cfg(not(feature = "support_kmer_counters"))]
+        () => entry.get_counter() & !VISITED_BIT,
+        #[cfg(feature = "support_kmer_counters")]
+        () => entry.color_index & !VISITED_BIT,
+    }) as ColorIndexType
+}
+
 impl<H: MinimizerHashFunctionFactory, MH: HashFunctionFactory> ColorsMergeManager<H, MH>
     for MultipleColorsManager<H, MH>
 {
@@ -395,7 +411,7 @@ impl<H: MinimizerHashFunctionFactory, MH: HashFunctionFactory> ColorsMergeManage
         ts: &mut Self::TempUnitigColorStructure,
         entry: &MapEntry<Self::HashMapTempColorIndex>,
     ) {
-        let kmer_color = (entry.get_counter() & !VISITED_BIT) as ColorIndexType;
+        let kmer_color = get_entry_color(entry);
 
         if let Some(back_ts) = ts.colors.back_mut() {
             if back_ts.color == kmer_color {
@@ -414,7 +430,7 @@ impl<H: MinimizerHashFunctionFactory, MH: HashFunctionFactory> ColorsMergeManage
         ts: &mut Self::TempUnitigColorStructure,
         entry: &MapEntry<Self::HashMapTempColorIndex>,
     ) {
-        let kmer_color = (entry.get_counter() & !VISITED_BIT) as ColorIndexType;
+        let kmer_color = get_entry_color(entry);
 
         if let Some(front_ts) = ts.colors.front_mut() {
             if front_ts.color == kmer_color {
@@ -528,7 +544,7 @@ impl<H: MinimizerHashFunctionFactory, MH: HashFunctionFactory> ColorsMergeManage
                 .flatten(),
         ) {
             let entry = hmap.get(&hash.to_unextendable()).unwrap();
-            let kmer_color = (entry.get_counter() & !VISITED_BIT) as ColorIndexType;
+            let kmer_color = get_entry_color(entry);
             if kmer_color != color {
                 let hashes = MH::new(read, 31);
                 println!(
@@ -537,7 +553,7 @@ impl<H: MinimizerHashFunctionFactory, MH: HashFunctionFactory> ColorsMergeManage
                         .iter()
                         .map(|h| {
                             let entry = hmap.get(&h.to_unextendable()).unwrap();
-                            let kmer_color = (entry.get_counter() & !VISITED_BIT) as ColorIndexType;
+                            let kmer_color = get_entry_color(entry);
                             kmer_color
                         })
                         .zip(