Prepare for compactor rewrite

Author: Guilucand

Cargo.lock

@@ -70,7 +70,7 @@ impl<H: Serialize + DeserializeOwned + Copy> BucketItemSerializer
     type ExtraDataBuffer = ();
     type ReadType<'a> = MaximalHashEntry<H>;
 
-    type ChunkData = ();
+    type CheckpointData = ();
 
     #[inline(always)]
     fn new() -> Self {

crates/assembler/src/pipeline/maximal_unitig_links/maximal_hash_entry.rs

@@ -123,7 +123,7 @@ impl BucketItemSerializer for MaximalUnitigLinkSerializer {
     type ExtraDataBuffer = ();
     type ReadType<'a> = MaximalUnitigLink;
 
-    type ChunkData = ();
+    type CheckpointData = ();
 
     #[inline(always)]
     fn new() -> Self {

crates/assembler/src/pipeline/maximal_unitig_links/maximal_unitig_index.rs

@@ -19,7 +19,7 @@ impl BucketItemSerializer for LinkMappingSerializer {
     type ExtraDataBuffer = ();
     type ReadType<'a> = LinkMapping;
 
-    type ChunkData = ();
+    type CheckpointData = ();
 
     #[inline(always)]
     fn new() -> Self {

crates/assembler/src/structs/link_mapping.rs

@@ -1,7 +1,7 @@
 use crate::final_executor::ParallelKmersMergeFinalExecutor;
 use crate::map_processor::{ParallelKmersMergeMapProcessor, KMERGE_TEMP_DIR};
-use crate::preprocessor::ParallelKmersMergePreprocessor;
 use crate::structs::{ResultsBucket, RetType};
+use assembler_minimizer_bucketing::rewrite_bucket::RewriteBucketComputeAssembler;
 use assembler_minimizer_bucketing::AssemblerMinimizerBucketingExecutorFactory;
 use colors::colors_manager::color_types::{
     GlobalColorsTableWriter, MinimizerBucketingSeqColorDataType,
@@ -17,7 +17,9 @@ use hashes::HashFunctionFactory;
 use io::structs::hash_entry::HashEntry;
 use io::structs::hash_entry::{Direction, HashEntrySerializer};
 use kmers_transform::processor::KmersTransformProcessor;
-use kmers_transform::{KmersTransform, KmersTransformExecutorFactory};
+use kmers_transform::{
+    KmersTransform, KmersTransformExecutorFactory, KmersTransformGlobalExtraData,
+};
 use minimizer_bucketing::{MinimizerBucketingCommonData, MinimizerBucketingExecutorFactory};
 use parallel_processor::buckets::bucket_writer::BucketItemSerializer;
 use parallel_processor::buckets::concurrent::BucketsThreadDispatcher;
@@ -37,7 +39,6 @@ use utils::owned_drop::OwnedDrop;
 
 mod final_executor;
 mod map_processor;
-mod preprocessor;
 pub mod structs;
 
 pub struct GlobalMergeData<CX: ColorsManager> {
@@ -55,6 +56,17 @@ pub struct GlobalMergeData<CX: ColorsManager> {
     kmer_batches_count: AtomicU64,
 }
 
+impl<CX: ColorsManager> KmersTransformGlobalExtraData for GlobalMergeData<CX> {
+    #[inline(always)]
+    fn get_k(&self) -> usize {
+        self.k
+    }
+    #[inline(always)]
+    fn get_m(&self) -> usize {
+        self.m
+    }
+}
+
 pub struct ParallelKmersMergeFactory<
     MH: HashFunctionFactory,
     CX: ColorsManager,
@@ -68,7 +80,7 @@ impl<MH: HashFunctionFactory, CX: ColorsManager, const COMPUTE_SIMPLITIGS: bool>
     type GlobalExtraData = GlobalMergeData<CX>;
     type AssociatedExtraData = MinimizerBucketingSeqColorDataType<CX>;
 
-    type PreprocessorType = ParallelKmersMergePreprocessor<MH, CX, COMPUTE_SIMPLITIGS>;
+    type PreprocessorType = RewriteBucketComputeAssembler;
     type MapProcessorType = ParallelKmersMergeMapProcessor<MH, CX, COMPUTE_SIMPLITIGS>;
     type FinalExecutorType = ParallelKmersMergeFinalExecutor<MH, CX, COMPUTE_SIMPLITIGS>;
 
@@ -82,10 +94,6 @@ impl<MH: HashFunctionFactory, CX: ColorsManager, const COMPUTE_SIMPLITIGS: bool>
         AssemblerMinimizerBucketingExecutorFactory::new(&global_data.global_resplit_data)
     }
 
-    fn new_preprocessor(_global_data: &Arc<Self::GlobalExtraData>) -> Self::PreprocessorType {
-        ParallelKmersMergePreprocessor::new()
-    }
-
     fn new_map_processor(
         _global_data: &Arc<Self::GlobalExtraData>,
         mem_tracker: MemoryTracker<KmersTransformProcessor<Self>>,

crates/assembler_kmers_merge/src/lib.rs

@@ -1,3 +1,5 @@
+pub mod rewrite_bucket;
+
 use ::dynamic_dispatch::dynamic_dispatch;
 use colors::colors_manager::color_types::MinimizerBucketingSeqColorDataType;
 use colors::colors_manager::{ColorsManager, MinimizerBucketingSeqColorData};
@@ -66,6 +68,7 @@ impl<CX: ColorsManager> MinimizerBucketingExecutorFactory
     type StreamInfo = InputFileInfo;
 
     type ColorsManager = CX;
+    type RewriteBucketCompute = rewrite_bucket::RewriteBucketComputeAssembler;
 
     #[allow(non_camel_case_types)]
     type FLAGS_COUNT = typenum::U2;

crates/assembler_kmers_merge/src/preprocessor.rs

@@ -0,0 +1,38 @@
+use config::BucketIndexType;
+use config::MultiplicityCounterType;
+use config::READ_FLAG_INCL_END;
+use hashes::default::MNHFactory;
+use hashes::ExtendableHashTraitType;
+use hashes::HashFunction;
+use hashes::{HashFunctionFactory, HashableSequence, MinimizerHashFunctionFactory};
+use io::compressed_read::CompressedRead;
+use minimizer_bucketing::resplit_bucket::RewriteBucketCompute;
+
+pub struct RewriteBucketComputeAssembler;
+
+impl RewriteBucketCompute for RewriteBucketComputeAssembler {
+    fn get_rewrite_bucket<C>(
+        k: usize,
+        m: usize,
+        seq_data: &(u8, u8, C, CompressedRead, MultiplicityCounterType),
+        used_hash_bits: usize,
+        bucket_bits_count: usize,
+    ) -> BucketIndexType {
+        let read = &seq_data.3;
+        let flags = seq_data.0;
+        let decr_val = ((read.bases_count() == k) && (flags & READ_FLAG_INCL_END) == 0) as usize;
+
+        let hashes = MNHFactory::new(read.sub_slice((1 - decr_val)..(k - decr_val)), m);
+
+        let minimizer = hashes
+            .iter()
+            .min_by_key(|k| MNHFactory::get_full_minimizer(k.to_unextendable()))
+            .unwrap();
+
+        MNHFactory::get_bucket(
+            used_hash_bits,
+            bucket_bits_count,
+            minimizer.to_unextendable(),
+        )
+    }
+}

crates/assembler_minimizer_bucketing/src/lib.rs

@@ -11,6 +11,7 @@ use io::concurrent::temp_reads::extra_data::{
 use io::sequences_reader::{DnaSequence, DnaSequencesFileType};
 use io::sequences_stream::fasta::FastaFileSequencesStream;
 use io::sequences_stream::SequenceInfo;
+use minimizer_bucketing::resplit_bucket::RewriteBucketCompute;
 use minimizer_bucketing::{
     GenericMinimizerBucketing, MinimizerBucketingCommonData, MinimizerBucketingExecutor,
     MinimizerBucketingExecutorFactory, MinimizerInputSequence,
@@ -101,6 +102,26 @@ pub struct DumperMinimizerBucketingExecutor<CX: ColorsManager> {
     _phantom: PhantomData<CX>,
 }
 
+pub struct RewriteBucketComputeDumper;
+
+impl RewriteBucketCompute for RewriteBucketComputeDumper {
+    fn get_rewrite_bucket<C>(
+        _k: usize,
+        _m: usize,
+        _seq_data: &(
+            u8,
+            u8,
+            C,
+            io::compressed_read::CompressedRead,
+            config::MultiplicityCounterType,
+        ),
+        _used_hash_bits: usize,
+        _bucket_bits_count: usize,
+    ) -> BucketIndexType {
+        unimplemented!()
+    }
+}
+
 pub struct DumperMinimizerBucketingExecutorFactory<CX: ColorsManager>(PhantomData<CX>);
 
 impl<CX: ColorsManager> MinimizerBucketingExecutorFactory
@@ -112,6 +133,7 @@ impl<CX: ColorsManager> MinimizerBucketingExecutorFactory
     type StreamInfo = ();
 
     type ColorsManager = CX;
+    type RewriteBucketCompute = RewriteBucketComputeDumper;
 
     #[allow(non_camel_case_types)]
     type FLAGS_COUNT = typenum::U0;

crates/assembler_minimizer_bucketing/src/rewrite_bucket.rs

@@ -4,7 +4,7 @@ use crate::varint::{
     VARINT_MAX_SIZE,
 };
 use byteorder::ReadBytesExt;
-use config::MultiplicityCounterType;
+use config::{BucketIndexType, MultiplicityCounterType};
 use parallel_processor::buckets::bucket_writer::BucketItemSerializer;
 use serde::{Deserialize, Serialize};
 use std::io::Read;
@@ -123,8 +123,8 @@ pub struct CompressedReadsBucketDataSerializer<
 }
 
 #[derive(Serialize, Deserialize)]
-pub struct ReadsChunkData {
-    pub target_subbucket: u8,
+pub struct ReadsCheckpointData {
+    pub target_subbucket: BucketIndexType,
 }
 
 impl<
@@ -142,7 +142,7 @@ impl<
     type ExtraDataBuffer = E::TempBuffer;
     type ReadType<'b> = (u8, u8, E, CompressedRead<'b>, MultiplicityCounterType);
 
-    type ChunkData = ReadsChunkData;
+    type CheckpointData = ReadsCheckpointData;
 
     #[inline(always)]
     fn new() -> Self {