1mod resource;
22
23pub(crate) use resource::{current_open_file_count, open_file_limit};
24pub use resource::{raise_open_file_limit, report_process_memory, trim_process_allocations};
25
26use crate::DEFAULT_VERTEX_PARTITIONS;
27use crate::Side;
28use crate::buckets::{BucketError, BucketLocation, BucketManifestEntry, BucketStore};
29use crate::color::{
30 ColorError, ColorRepositoryManifest, ColorRunSidecar, ColorRunSidecarWriter,
31 ConcurrentColorRepository, ConcurrentColorRunSidecarWriter, append_color_runs,
32 read_unitig_color_runs, reverse_color_runs, reverse_color_runs_in_place,
33 write_unitig_color_runs,
34};
35use crate::dna::{Base, complement_ascii};
36use crate::hash::{FastBuildHasher, hash_bytes, hash_u64, wyhash_u64};
37use crate::kmer::Kmer;
38use crate::state::{UnitigColor, VertexState};
39use crate::subgraph::{LocalSubgraph, LocalSubgraphError, LocalUnitig, LocalVertexMap};
40use rayon::prelude::*;
41use rayon::{ThreadPool, ThreadPoolBuilder};
42use std::cell::UnsafeCell;
43use std::collections::{BTreeMap, BTreeSet, HashMap, HashSet};
44use std::fs::{self, File, OpenOptions};
45use std::io::{BufReader, BufWriter, Read, Seek, SeekFrom, Write};
46use std::marker::PhantomData;
47use std::mem::MaybeUninit;
48use std::os::unix::fs::FileExt;
49use std::path::{Path, PathBuf};
50use std::sync::{
51 Arc, Mutex,
52 atomic::{AtomicU8, AtomicU64, AtomicUsize, Ordering},
53 mpsc,
54};
55use std::time::{Duration, Instant};
56use xxhash_rust::xxh3::xxh3_64;
57
58type FastHashMap<K, V> = HashMap<K, V, FastBuildHasher>;
59type FastHashSet<T> = HashSet<T, FastBuildHasher>;
60
61fn keep_intermediates() -> bool {
62 std::env::var_os("CF3_RS_KEEP_INTERMEDIATES").is_some()
63}
64
65fn remove_serial_file(path: &Path) -> Result<(), SerialCollationError> {
66 if keep_intermediates() {
67 return Ok(());
68 }
69 match fs::remove_file(path) {
70 Ok(()) => Ok(()),
71 Err(source) if source.kind() == std::io::ErrorKind::NotFound => Ok(()),
72 Err(source) => Err(SerialCollationError::Io {
73 path: path.to_path_buf(),
74 source,
75 }),
76 }
77}
78
79#[derive(Debug, Clone, PartialEq, Eq)]
80pub struct DiscontinuityInputs<const K: usize> {
84 pub unitigs: Vec<DiscontinuityUnitig<K>>,
85 labels: Vec<u8>,
86 pub stats: DiscontinuityInputStats,
87}
88
89#[derive(Debug)]
90pub struct ExternalDiscontinuityInputs<const K: usize> {
95 unitig_path: PathBuf,
96 label_path: PathBuf,
97 unitigs: usize,
98 compact_unitigs: bool,
99 ranges: Vec<ExternalLocalUnitigRange>,
100 edge_matrix: Option<BlockedEdgeMatrix<K>>,
101 color_runs: Option<ColorRunSidecar>,
102 local_unitig_buckets: Option<Vec<LocalUnitigBucketEntry>>,
103 local_unitig_bucket_dir: Option<PathBuf>,
107 trivial_fasta: Option<(PathBuf, u64, u64)>,
108 trivial_is_output: bool,
112 color_repository: Option<ColorRepositoryManifest>,
113 pub stats: DiscontinuityInputStats,
114}
115
116impl<const K: usize> ExternalDiscontinuityInputs<K> {
117 pub fn unitig_count(&self) -> usize {
118 self.unitigs
119 }
120
121 pub fn color_runs(&self) -> Option<&ColorRunSidecar> {
122 self.color_runs.as_ref()
123 }
124
125 pub fn color_repository(&self) -> Option<&ColorRepositoryManifest> {
126 self.color_repository.as_ref()
127 }
128}
129
130#[derive(Debug, Clone, Copy, PartialEq, Eq)]
131struct ExternalLocalUnitigRange {
132 start_unitig: usize,
133 unitigs: usize,
134 label_start: u64,
135 label_len: u64,
136 color_start: u64,
137}
138
139#[derive(Debug)]
140struct LocalUnitigBucketEntry {
141 bucket_id: u16,
142 unitig_path: PathBuf,
143 label_path: PathBuf,
144 unitigs: usize,
145 colored: bool,
146}
147
148struct LocalUnitigBucketWriter {
149 bucket_id: u16,
150 unitig_path: PathBuf,
151 label_path: PathBuf,
152 unitigs: BufWriter<File>,
153 labels: BufWriter<File>,
154 colored: bool,
155 unitig_count: usize,
156}
157
158impl LocalUnitigBucketWriter {
159 fn create(dir: &Path, bucket_id: u16, colored: bool) -> Result<Self, DiscontinuityInputError> {
160 let unitig_path = dir.join(format!("{bucket_id:03}.unitigs"));
161 let label_path = dir.join(format!("{bucket_id:03}.labels"));
162 let unitig_file =
163 File::create(&unitig_path).map_err(|source| DiscontinuityInputError::Io {
164 path: unitig_path.clone(),
165 source,
166 })?;
167 let label_file =
168 File::create(&label_path).map_err(|source| DiscontinuityInputError::Io {
169 path: label_path.clone(),
170 source,
171 })?;
172 Ok(Self {
173 bucket_id,
174 unitig_path,
175 label_path,
176 unitigs: BufWriter::with_capacity(1024 * 1024, unitig_file),
177 labels: BufWriter::with_capacity(4 * 1024 * 1024, label_file),
178 colored,
179 unitig_count: 0,
180 })
181 }
182
183 fn write<const K: usize>(
184 &mut self,
185 labels: &[u8],
186 unitigs: &[DiscontinuityUnitig<K>],
187 colors: Option<&[Vec<UnitigColor>]>,
188 ) -> Result<usize, DiscontinuityInputError> {
189 if colors.is_some_and(|runs| runs.len() != unitigs.len()) {
190 return Err(DiscontinuityInputError::MissingColorRuns);
191 }
192 if self.colored != colors.is_some() {
193 return Err(DiscontinuityInputError::MissingColorRuns);
194 }
195 let base = self.unitig_count;
196 self.labels
197 .write_all(labels)
198 .map_err(|source| DiscontinuityInputError::Io {
199 path: self.label_path.clone(),
200 source,
201 })?;
202 for (index, unitig) in unitigs.iter().enumerate() {
203 write_discontinuity_unitig_record(&mut self.unitigs, &self.unitig_path, unitig, true)?;
204 if let Some(colors) = colors {
205 write_unitig_color_runs(&mut self.unitigs, &colors[index]).map_err(|source| {
206 DiscontinuityInputError::Io {
207 path: self.unitig_path.clone(),
208 source,
209 }
210 })?;
211 }
212 }
213 self.unitig_count += unitigs.len();
214 Ok(base)
215 }
216
217 fn finish(mut self) -> Result<LocalUnitigBucketEntry, DiscontinuityInputError> {
218 self.unitigs
219 .flush()
220 .map_err(|source| DiscontinuityInputError::Io {
221 path: self.unitig_path.clone(),
222 source,
223 })?;
224 self.labels
225 .flush()
226 .map_err(|source| DiscontinuityInputError::Io {
227 path: self.label_path.clone(),
228 source,
229 })?;
230 Ok(LocalUnitigBucketEntry {
231 bucket_id: self.bucket_id,
232 unitig_path: self.unitig_path,
233 label_path: self.label_path,
234 unitigs: self.unitig_count,
235 colored: self.colored,
236 })
237 }
238}
239
240fn finish_local_unitig_writers(
241 writers: Vec<Mutex<LocalUnitigBucketWriter>>,
242 threads: usize,
243) -> Result<Vec<LocalUnitigBucketEntry>, DiscontinuityInputError> {
244 let worker_count = threads.max(1).min(writers.len().max(1));
245 let mut work = (0..worker_count).map(|_| Vec::new()).collect::<Vec<_>>();
246 for (index, writer) in writers.into_iter().enumerate() {
247 work[index % worker_count].push(writer);
248 }
249 let mut entries = std::thread::scope(|scope| {
250 let mut handles = Vec::with_capacity(worker_count);
251 for worker_writers in work {
252 handles.push(scope.spawn(move || {
253 let mut entries = Vec::with_capacity(worker_writers.len());
254 for writer in worker_writers {
255 entries.push(
256 writer
257 .into_inner()
258 .map_err(|_| DiscontinuityInputError::WorkerPanic)?
259 .finish()?,
260 );
261 }
262 Ok::<_, DiscontinuityInputError>(entries)
263 }));
264 }
265 let mut entries = Vec::new();
266 for handle in handles {
267 entries.extend(
268 handle
269 .join()
270 .map_err(|_| DiscontinuityInputError::WorkerPanic)??,
271 );
272 }
273 Ok::<_, DiscontinuityInputError>(entries)
274 })?;
275 entries.sort_unstable_by_key(|entry| entry.bucket_id);
276 Ok(entries)
277}
278
279#[derive(Debug, Clone, Copy, PartialEq, Eq)]
280struct ExternalLabelRef {
281 label_start: u64,
282 label_len: u32,
283}
284
285#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
286pub struct DiscontinuityInputStats {
288 pub input_buckets: usize,
289 pub weak_superkmers: u64,
290 pub local_unitigs: u64,
291 pub discontinuity_exits: u64,
292 pub unitig_bases: u64,
293}
294
295#[repr(C)]
296#[derive(Debug, Clone, Copy, PartialEq, Eq)]
297pub struct DiscontinuityUnitig<const K: usize> {
298 pub label_start: u64,
299 left_vertex: Kmer<K>,
300 right_vertex: Kmer<K>,
301 pub label_len: u32,
302 flags: u8,
303}
304
305const DISCONTINUITY_LEFT_EXIT: u8 = 1 << 0;
306const DISCONTINUITY_LEFT_BACK: u8 = 1 << 1;
307const DISCONTINUITY_RIGHT_EXIT: u8 = 1 << 2;
308const DISCONTINUITY_RIGHT_BACK: u8 = 1 << 3;
309const DISCONTINUITY_CYCLE: u8 = 1 << 4;
310
311impl<const K: usize> DiscontinuityInputs<K> {
312 pub fn empty(stats: DiscontinuityInputStats) -> Self {
313 Self {
314 unitigs: Vec::new(),
315 labels: Vec::new(),
316 stats,
317 }
318 }
319
320 pub fn from_unitigs(
321 unitigs: impl IntoIterator<Item = OwnedDiscontinuityUnitig<K>>,
322 stats: DiscontinuityInputStats,
323 ) -> Self {
324 let mut inputs = Self::empty(stats);
325 for unitig in unitigs {
326 inputs.push_unitig(unitig);
327 }
328 inputs
329 }
330
331 pub fn push_unitig(&mut self, unitig: OwnedDiscontinuityUnitig<K>) -> usize {
332 let index = self.unitigs.len();
333 let label_start = self.labels.len();
334 let label_len = unitig.label.len();
335 assert!(u64::try_from(label_start).is_ok());
336 assert!(u32::try_from(label_len).is_ok());
337 self.labels.extend_from_slice(&unitig.label);
338 self.unitigs.push(DiscontinuityUnitig {
339 label_start: label_start as u64,
340 left_vertex: unitig
341 .left_exit
342 .map(|endpoint| endpoint.vertex)
343 .unwrap_or_else(Kmer::zero),
344 right_vertex: unitig
345 .right_exit
346 .map(|endpoint| endpoint.vertex)
347 .unwrap_or_else(Kmer::zero),
348 label_len: label_len as u32,
349 flags: discontinuity_unitig_flags(unitig.left_exit, unitig.right_exit, unitig.is_cycle),
350 });
351 index
352 }
353
354 #[inline]
355 pub fn label(&self, unitig_index: usize) -> &[u8] {
356 self.unitigs[unitig_index].label(self)
357 }
358
359 #[inline]
360 pub fn try_label(&self, unitig_index: usize) -> Option<&[u8]> {
361 self.unitigs
362 .get(unitig_index)
363 .map(|unitig| unitig.label(self))
364 }
365}
366
367impl<const K: usize> DiscontinuityUnitig<K> {
368 #[inline]
369 pub fn label<'a>(&self, inputs: &'a DiscontinuityInputs<K>) -> &'a [u8] {
370 let start = self.label_start as usize;
371 let end = start + self.label_len as usize;
372 &inputs.labels[start..end]
373 }
374
375 #[inline]
376 pub fn left_exit(&self) -> Option<DiscontinuityEndpoint<K>> {
377 (self.flags & DISCONTINUITY_LEFT_EXIT != 0).then_some(DiscontinuityEndpoint {
378 vertex: self.left_vertex,
379 side: if self.flags & DISCONTINUITY_LEFT_BACK != 0 {
380 Side::Back
381 } else {
382 Side::Front
383 },
384 })
385 }
386
387 #[inline]
388 pub fn right_exit(&self) -> Option<DiscontinuityEndpoint<K>> {
389 (self.flags & DISCONTINUITY_RIGHT_EXIT != 0).then_some(DiscontinuityEndpoint {
390 vertex: self.right_vertex,
391 side: if self.flags & DISCONTINUITY_RIGHT_BACK != 0 {
392 Side::Back
393 } else {
394 Side::Front
395 },
396 })
397 }
398
399 #[inline]
400 pub fn is_cycle(&self) -> bool {
401 self.flags & DISCONTINUITY_CYCLE != 0
402 }
403}
404
405struct ExternalDiscontinuityReader<const K: usize> {
406 label_file: File,
407 unitig_path: PathBuf,
408 label_path: PathBuf,
409 unitigs: usize,
410 compact_unitigs: bool,
411}
412
413impl<const K: usize> ExternalDiscontinuityReader<K> {
414 fn open(inputs: &ExternalDiscontinuityInputs<K>) -> Result<Self, SerialCollationError> {
415 let label_file =
416 File::open(&inputs.label_path).map_err(|source| SerialCollationError::Io {
417 path: inputs.label_path.clone(),
418 source,
419 })?;
420 Ok(Self {
421 label_file,
422 unitig_path: inputs.unitig_path.clone(),
423 label_path: inputs.label_path.clone(),
424 unitigs: inputs.unitigs,
425 compact_unitigs: inputs.compact_unitigs,
426 })
427 }
428
429 fn read_label(
430 &self,
431 unitig: &DiscontinuityUnitig<K>,
432 scratch: &mut Vec<u8>,
433 ) -> Result<(), SerialCollationError> {
434 scratch.resize(unitig.label_len as usize, 0);
435 self.label_file
436 .read_exact_at(scratch, unitig.label_start)
437 .map_err(|source| SerialCollationError::Io {
438 path: self.label_path.clone(),
439 source,
440 })
441 }
442
443 fn iter(&self) -> Result<ExternalDiscontinuityIter<K>, SerialCollationError> {
444 let file = File::open(&self.unitig_path).map_err(|source| SerialCollationError::Io {
445 path: self.unitig_path.clone(),
446 source,
447 })?;
448 Ok(ExternalDiscontinuityIter {
449 input: BufReader::with_capacity(1024 * 1024, file),
450 path: self.unitig_path.clone(),
451 remaining: self.unitigs,
452 compact_unitigs: self.compact_unitigs,
453 })
454 }
455}
456
457struct ExternalDiscontinuityIter<const K: usize> {
458 input: BufReader<File>,
459 path: PathBuf,
460 remaining: usize,
461 compact_unitigs: bool,
462}
463
464impl<const K: usize> ExternalDiscontinuityIter<K> {
465 fn next_unitig(&mut self) -> Result<DiscontinuityUnitig<K>, SerialCollationError> {
466 if self.remaining == 0 {
467 return Err(SerialCollationError::MalformedCoordBucket(
468 self.path.clone(),
469 ));
470 }
471 self.remaining -= 1;
472 read_discontinuity_unitig_from_reader(&mut self.input, &self.path, self.compact_unitigs)
473 }
474}
475
476impl<const K: usize> Iterator for ExternalDiscontinuityIter<K> {
477 type Item = Result<DiscontinuityUnitig<K>, SerialCollationError>;
478
479 fn next(&mut self) -> Option<Self::Item> {
480 if self.remaining == 0 {
481 return None;
482 }
483 Some(self.next_unitig())
484 }
485}
486
487fn read_discontinuity_unitig_from_reader<const K: usize>(
488 input: &mut BufReader<File>,
489 path: &Path,
490 compact: bool,
491) -> Result<DiscontinuityUnitig<K>, SerialCollationError> {
492 if compact {
493 let mut bytes = [0u8; 8];
494 input
495 .read_exact(&mut bytes)
496 .map_err(|source| SerialCollationError::Io {
497 path: path.to_path_buf(),
498 source,
499 })?;
500 return Ok(DiscontinuityUnitig {
501 label_start: 0,
502 left_vertex: Kmer::zero(),
503 right_vertex: Kmer::zero(),
504 label_len: u32::from_le_bytes(bytes[..4].try_into().expect("label length")),
505 flags: bytes[4],
506 });
507 }
508 let mut out = MaybeUninit::<DiscontinuityUnitig<K>>::zeroed();
509 let bytes = unsafe {
510 std::slice::from_raw_parts_mut(
511 out.as_mut_ptr().cast::<u8>(),
512 std::mem::size_of::<DiscontinuityUnitig<K>>(),
513 )
514 };
515 input
516 .read_exact(bytes)
517 .map_err(|source| SerialCollationError::Io {
518 path: path.to_path_buf(),
519 source,
520 })?;
521 Ok(unsafe { out.assume_init() })
522}
523
524const fn external_unitig_record_len<const K: usize>(compact: bool) -> usize {
525 if compact {
526 8
527 } else {
528 std::mem::size_of::<DiscontinuityUnitig<K>>()
529 }
530}
531
532fn read_discontinuity_unitig_from_reader_for_input<const K: usize>(
533 input: &mut BufReader<File>,
534 path: &Path,
535) -> Result<DiscontinuityUnitig<K>, DiscontinuityInputError> {
536 let mut out = MaybeUninit::<DiscontinuityUnitig<K>>::zeroed();
537 let bytes = unsafe {
538 std::slice::from_raw_parts_mut(
539 out.as_mut_ptr().cast::<u8>(),
540 std::mem::size_of::<DiscontinuityUnitig<K>>(),
541 )
542 };
543 input
544 .read_exact(bytes)
545 .map_err(|source| DiscontinuityInputError::Io {
546 path: path.to_path_buf(),
547 source,
548 })?;
549 Ok(unsafe { out.assume_init() })
550}
551
552#[inline]
553fn discontinuity_unitig_flags<const K: usize>(
554 left_exit: Option<DiscontinuityEndpoint<K>>,
555 right_exit: Option<DiscontinuityEndpoint<K>>,
556 is_cycle: bool,
557) -> u8 {
558 let mut flags = 0u8;
559 if let Some(endpoint) = left_exit {
560 flags |= DISCONTINUITY_LEFT_EXIT;
561 if endpoint.side == Side::Back {
562 flags |= DISCONTINUITY_LEFT_BACK;
563 }
564 }
565 if let Some(endpoint) = right_exit {
566 flags |= DISCONTINUITY_RIGHT_EXIT;
567 if endpoint.side == Side::Back {
568 flags |= DISCONTINUITY_RIGHT_BACK;
569 }
570 }
571 if is_cycle {
572 flags |= DISCONTINUITY_CYCLE;
573 }
574 flags
575}
576
577#[derive(Debug, Clone, PartialEq, Eq)]
578pub struct OwnedDiscontinuityUnitig<const K: usize> {
579 pub graph_id: usize,
580 pub label: Vec<u8>,
581 pub left_exit: Option<DiscontinuityEndpoint<K>>,
582 pub right_exit: Option<DiscontinuityEndpoint<K>>,
583 pub is_cycle: bool,
584}
585
586#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
587pub struct DiscontinuityEndpoint<const K: usize> {
588 pub vertex: Kmer<K>,
589 pub side: Side,
590}
591
592#[derive(Debug, Clone, PartialEq, Eq)]
593pub struct SerialEdgeMatrix<const K: usize> {
594 vertex_partitions: usize,
595 blocks: Vec<Vec<Vec<DiscontinuityEdge<K>>>>,
596 stats: SerialEdgeMatrixStats,
597}
598
599#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
600pub struct SerialEdgeMatrixStats {
601 pub edges: u64,
602 pub phi_edges: u64,
603 pub diagonal_edges: u64,
604}
605
606#[derive(Debug, Clone, PartialEq, Eq)]
607pub struct DiscontinuityEdge<const K: usize> {
608 pub first: MatrixEndpoint<K>,
609 pub second: MatrixEndpoint<K>,
610 pub weight: u64,
611 pub unitig_bucket: u16,
612 pub unitig_index: usize,
613 pub unitig_exit_side: Side,
614 pub phantom_unitig: Option<DiscontinuityEndpoint<K>>,
615 pub swapped: bool,
616}
617
618#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
619pub enum MatrixEndpoint<const K: usize> {
620 Phi,
621 Vertex(DiscontinuityEndpoint<K>),
622}
623
624impl<const K: usize> SerialEdgeMatrix<K> {
625 pub fn new(vertex_partitions: usize) -> Result<Self, SerialEdgeMatrixError> {
626 if vertex_partitions == 0 || !vertex_partitions.is_power_of_two() {
627 return Err(SerialEdgeMatrixError::InvalidPartitionCount(
628 vertex_partitions,
629 ));
630 }
631
632 let partition_count = vertex_partitions + 1;
633 Ok(Self {
634 vertex_partitions,
635 blocks: vec![vec![Vec::new(); partition_count]; partition_count],
636 stats: SerialEdgeMatrixStats::default(),
637 })
638 }
639
640 pub fn from_inputs(
641 inputs: &DiscontinuityInputs<K>,
642 vertex_partitions: usize,
643 ) -> Result<Self, SerialEdgeMatrixError> {
644 let mut matrix = Self::new(vertex_partitions)?;
645 for (unitig_index, unitig) in inputs.unitigs.iter().enumerate() {
646 match (unitig.left_exit(), unitig.right_exit()) {
647 (Some(left), Some(right)) => matrix.add_edge(
648 MatrixEndpoint::Vertex(left),
649 MatrixEndpoint::Vertex(right),
650 1,
651 unitig_index,
652 ),
653 (Some(endpoint), None) => matrix.add_edge_with_orientation(
654 MatrixEndpoint::Phi,
655 MatrixEndpoint::Vertex(endpoint),
656 1,
657 unitig_index,
658 Side::Front,
659 ),
660 (None, Some(endpoint)) => matrix.add_edge_with_orientation(
661 MatrixEndpoint::Phi,
662 MatrixEndpoint::Vertex(endpoint),
663 1,
664 unitig_index,
665 Side::Back,
666 ),
667 (None, None) => {}
668 }
669 }
670 Ok(matrix)
671 }
672
673 pub fn add_edge(
674 &mut self,
675 first: MatrixEndpoint<K>,
676 second: MatrixEndpoint<K>,
677 weight: u64,
678 unitig_index: usize,
679 ) {
680 self.add_edge_with_orientation(first, second, weight, unitig_index, Side::Back);
681 }
682
683 pub fn add_edge_with_orientation(
684 &mut self,
685 first: MatrixEndpoint<K>,
686 second: MatrixEndpoint<K>,
687 weight: u64,
688 unitig_index: usize,
689 unitig_exit_side: Side,
690 ) {
691 self.add_edge_with_orientation_and_phantom(
692 first,
693 second,
694 weight,
695 unitig_index,
696 unitig_exit_side,
697 None,
698 );
699 }
700
701 pub fn add_edge_with_orientation_and_phantom(
702 &mut self,
703 first: MatrixEndpoint<K>,
704 second: MatrixEndpoint<K>,
705 weight: u64,
706 unitig_index: usize,
707 unitig_exit_side: Side,
708 phantom_unitig: Option<DiscontinuityEndpoint<K>>,
709 ) {
710 let first_partition = self.partition(first);
711 let second_partition = self.partition(second);
712 let swapped = first_partition > second_partition;
713 let unitig_exit_side = if swapped {
714 unitig_exit_side.inverse()
715 } else {
716 unitig_exit_side
717 };
718 let (row, col, first, second) = if swapped {
719 (second_partition, first_partition, second, first)
720 } else {
721 (first_partition, second_partition, first, second)
722 };
723
724 self.stats.edges += 1;
725 if first.is_phi() || second.is_phi() {
726 self.stats.phi_edges += 1;
727 }
728 if row == col {
729 self.stats.diagonal_edges += 1;
730 }
731
732 self.blocks[row][col].push(DiscontinuityEdge {
733 first,
734 second,
735 weight,
736 unitig_bucket: 0,
737 unitig_index,
738 unitig_exit_side,
739 phantom_unitig,
740 swapped,
741 });
742 }
743
744 #[inline]
745 pub const fn vertex_partitions(&self) -> usize {
746 self.vertex_partitions
747 }
748
749 #[inline]
750 pub const fn partition_count(&self) -> usize {
751 self.vertex_partitions + 1
752 }
753
754 #[inline]
755 pub const fn stats(&self) -> SerialEdgeMatrixStats {
756 self.stats
757 }
758
759 #[inline]
760 pub fn partition(&self, endpoint: MatrixEndpoint<K>) -> usize {
761 match endpoint {
762 MatrixEndpoint::Phi => 0,
763 MatrixEndpoint::Vertex(endpoint) => {
764 ((endpoint.vertex.hash64(0) as usize) & (self.vertex_partitions - 1)) + 1
765 }
766 }
767 }
768
769 #[inline]
770 pub fn block(&self, row: usize, col: usize) -> &[DiscontinuityEdge<K>] {
771 assert!(row <= col);
772 &self.blocks[row][col]
773 }
774
775 pub fn edges(&self) -> impl Iterator<Item = &DiscontinuityEdge<K>> {
776 self.blocks
777 .iter()
778 .enumerate()
779 .flat_map(|(row, blocks)| blocks[row..].iter())
780 .flat_map(|block| block.iter())
781 }
782}
783
784impl<const K: usize> MatrixEndpoint<K> {
785 #[inline]
786 pub const fn is_phi(self) -> bool {
787 matches!(self, Self::Phi)
788 }
789}
790
791const BLOCKED_EDGE_WRITE_BUFFER_BYTES: usize = 256 * 1024;
792
793#[derive(Debug, Clone, Copy)]
801struct BlockExtent {
802 offset: u64,
803 len: u64,
804}
805
806#[derive(Debug, Default)]
807struct BlockedEdgeBlock {
808 extents: Vec<BlockExtent>,
810 edges: usize,
811 record_len: usize,
812 buffer: Vec<u8>,
813}
814
815#[derive(Debug, Clone, Copy, PartialEq, Eq)]
822enum EdgeContainerAxis {
823 Row,
824 Column,
825}
826
827fn edge_container_axis() -> EdgeContainerAxis {
828 static AXIS: std::sync::OnceLock<EdgeContainerAxis> = std::sync::OnceLock::new();
829 *AXIS.get_or_init(
830 || match std::env::var("CF3_RS_EDGE_CONTAINER_AXIS").as_deref() {
831 Ok("column") => EdgeContainerAxis::Column,
832 _ => EdgeContainerAxis::Row,
833 },
834 )
835}
836
837#[derive(Debug)]
844struct EdgeContainers {
845 files: Vec<EdgeContainerFile>,
846 axis: EdgeContainerAxis,
847 partition_count: usize,
848}
849
850#[derive(Debug)]
851struct EdgeContainerFile {
852 path: PathBuf,
853 file: File,
854 cursor: AtomicU64,
855}
856
857impl EdgeContainers {
858 fn create(dir: &Path, partition_count: usize) -> Result<Self, SerialCollationError> {
859 let budget = open_file_limit()
863 .saturating_sub(current_open_file_count())
864 .saturating_sub(RESERVED_NON_MATRIX_DESCRIPTORS)
865 / 2;
866 let container_count = partition_count.min(budget.max(1));
867 if container_count < partition_count {
868 eprintln!(
869 "cuttlefish: descriptor budget allows {container_count} edge-matrix container(s) rather than {partition_count}"
870 );
871 }
872 let mut files = Vec::with_capacity(container_count);
873 for index in 0..container_count {
874 let path = dir.join(format!("{index:05}.edge"));
875 let file = OpenOptions::new()
876 .create(true)
877 .truncate(true)
878 .read(true)
879 .write(true)
880 .open(&path)
881 .map_err(|source| SerialCollationError::Io {
882 path: path.clone(),
883 source,
884 })?;
885 files.push(EdgeContainerFile {
886 path,
887 file,
888 cursor: AtomicU64::new(0),
889 });
890 }
891 Ok(Self {
892 files,
893 axis: edge_container_axis(),
894 partition_count,
895 })
896 }
897
898 #[inline]
899 fn container_index(&self, block_index: usize) -> usize {
900 let axis_index = match self.axis {
901 EdgeContainerAxis::Row => block_index / self.partition_count,
902 EdgeContainerAxis::Column => block_index % self.partition_count,
903 };
904 axis_index % self.files.len()
910 }
911
912 #[inline]
913 fn container_for(&self, block_index: usize) -> &EdgeContainerFile {
914 &self.files[self.container_index(block_index)]
915 }
916
917 fn append(
919 &self,
920 block_index: usize,
921 bytes: &[u8],
922 ) -> Result<BlockExtent, SerialCollationError> {
923 let container = self.container_for(block_index);
924 let offset = container
925 .cursor
926 .fetch_add(bytes.len() as u64, Ordering::Relaxed);
927 container
928 .file
929 .write_all_at(bytes, offset)
930 .map_err(|source| SerialCollationError::Io {
931 path: container.path.clone(),
932 source,
933 })?;
934 Ok(BlockExtent {
935 offset,
936 len: bytes.len() as u64,
937 })
938 }
939
940 fn read_run(&self, run: &ContainerRun) -> Result<Vec<u8>, SerialCollationError> {
942 let container = &self.files[run.container];
943 let mut bytes = vec![0u8; run.len];
944 container
945 .file
946 .read_exact_at(&mut bytes, run.offset)
947 .map_err(|source| SerialCollationError::Io {
948 path: container.path.clone(),
949 source,
950 })?;
951 Ok(bytes)
952 }
953
954 fn read_pass(
966 &self,
967 blocks: &[(usize, &[BlockExtent])],
968 ) -> Result<ContainerPass, SerialCollationError> {
969 let mut by_container: FastHashMap<usize, Vec<(usize, &[BlockExtent])>> =
970 FastHashMap::with_hasher(FastBuildHasher::default());
971 for (slot, (block_index, extents)) in blocks.iter().enumerate() {
972 by_container
973 .entry(self.container_index(*block_index))
974 .or_default()
975 .push((slot, extents));
976 }
977 let mut runs = Vec::new();
978 for (container, group) in by_container {
979 runs.extend(plan_container_runs(container, &group));
980 }
981 let buffers = runs
982 .par_iter()
983 .map(|run| self.read_run(run))
984 .collect::<Result<Vec<_>, _>>()?;
985 Ok(ContainerPass { runs, buffers })
986 }
987
988 fn read_block(
995 &self,
996 block_index: usize,
997 extents: &[BlockExtent],
998 ) -> Result<Vec<u8>, SerialCollationError> {
999 let container = self.container_for(block_index);
1000 let total: usize = extents.iter().map(|extent| extent.len as usize).sum();
1001 let mut bytes = vec![0u8; total];
1002 for run in plan_container_runs(self.container_index(block_index), &[(0, extents)]) {
1003 if let [only] = run.extents.as_slice()
1006 && only.len == run.len
1007 {
1008 container
1009 .file
1010 .read_exact_at(
1011 &mut bytes[only.dest_offset..only.dest_offset + only.len],
1012 run.offset,
1013 )
1014 .map_err(|source| SerialCollationError::Io {
1015 path: container.path.clone(),
1016 source,
1017 })?;
1018 continue;
1019 }
1020 let buffer = self.read_run(&run)?;
1021 for extent in &run.extents {
1022 bytes[extent.dest_offset..extent.dest_offset + extent.len]
1023 .copy_from_slice(&buffer[extent.run_offset..extent.run_offset + extent.len]);
1024 }
1025 }
1026 Ok(bytes)
1027 }
1028
1029 fn path_for(&self, block_index: usize) -> PathBuf {
1030 self.container_for(block_index).path.clone()
1031 }
1032}
1033
1034type ExpandedPartition<const K: usize> = (
1038 Vec<(StitchedCoordRecord, DiscontinuityEndpoint<K>)>,
1039 u64,
1040 u64,
1041 u64,
1042);
1043
1044const RESERVED_NON_MATRIX_DESCRIPTORS: usize = 192;
1050
1051const CONTAINER_READ_GAP_BYTES: u64 = 1024 * 1024;
1052
1053const CONTAINER_RUN_BYTES: u64 = 16 * 1024 * 1024;
1056
1057#[derive(Debug, Clone, Copy)]
1059struct RunExtent {
1060 slot: usize,
1062 run_offset: usize,
1064 dest_offset: usize,
1066 len: usize,
1067}
1068
1069#[derive(Debug)]
1071struct ContainerRun {
1072 container: usize,
1073 offset: u64,
1074 len: usize,
1075 extents: Vec<RunExtent>,
1076}
1077
1078#[derive(Debug)]
1080struct ContainerPass {
1081 runs: Vec<ContainerRun>,
1082 buffers: Vec<Vec<u8>>,
1083}
1084
1085impl ContainerPass {
1086 fn extents(&self) -> impl Iterator<Item = (usize, &[u8])> {
1091 self.runs
1092 .iter()
1093 .zip(&self.buffers)
1094 .flat_map(|(run, buffer)| {
1095 run.extents.iter().map(move |extent| {
1096 (
1097 extent.slot,
1098 &buffer[extent.run_offset..extent.run_offset + extent.len],
1099 )
1100 })
1101 })
1102 }
1103}
1104
1105fn plan_container_runs(container: usize, blocks: &[(usize, &[BlockExtent])]) -> Vec<ContainerRun> {
1113 let mut placed = Vec::new();
1114 for (slot, extents) in blocks {
1115 let mut dest_offset = 0usize;
1116 for extent in *extents {
1117 if extent.len != 0 {
1118 placed.push((extent.offset, extent.len, *slot, dest_offset));
1119 }
1120 dest_offset += extent.len as usize;
1121 }
1122 }
1123 placed.sort_unstable_by_key(|(offset, ..)| *offset);
1124
1125 let mut runs: Vec<ContainerRun> = Vec::new();
1126 for (offset, len, slot, dest_offset) in placed {
1127 let end = offset + len;
1128 let merge = runs.last().is_some_and(|run| {
1131 let run_end = run.offset + run.len as u64;
1132 offset >= run_end
1133 && offset - run_end <= CONTAINER_READ_GAP_BYTES
1134 && end - run.offset <= CONTAINER_RUN_BYTES
1135 });
1136 if merge && let Some(run) = runs.last_mut() {
1137 let run_offset = (offset - run.offset) as usize;
1138 run.len = (end - run.offset) as usize;
1139 run.extents.push(RunExtent {
1140 slot,
1141 run_offset,
1142 dest_offset,
1143 len: len as usize,
1144 });
1145 } else {
1146 runs.push(ContainerRun {
1147 container,
1148 offset,
1149 len: len as usize,
1150 extents: vec![RunExtent {
1151 slot,
1152 run_offset: 0,
1153 dest_offset,
1154 len: len as usize,
1155 }],
1156 });
1157 }
1158 }
1159 runs
1160}
1161
1162#[derive(Debug)]
1163struct BlockedEdgeMatrix<const K: usize> {
1164 dir: PathBuf,
1165 vertex_partitions: usize,
1166 containers: Arc<EdgeContainers>,
1167 blocks: Vec<BlockedEdgeBlock>,
1168 stats: SerialEdgeMatrixStats,
1169 phantom: PhantomData<[(); K]>,
1170}
1171
1172struct PreparedBlockedEdge {
1173 block: usize,
1174 bytes: [u8; 72],
1175 phi: bool,
1176 diagonal: bool,
1177}
1178
1179enum BlockedReadTask<'a> {
1180 File {
1181 file: &'a File,
1182 path: &'a Path,
1183 offset: u64,
1184 len: usize,
1185 },
1186 Memory(&'a [u8]),
1187}
1188
1189struct ConcurrentBlockedAppend {
1190 extents: Mutex<Vec<BlockExtent>>,
1193 record_len: usize,
1194 buffer: Mutex<Vec<u8>>,
1195 edges: AtomicUsize,
1196}
1197
1198struct ConcurrentBlockedEdgeWriters {
1199 containers: Arc<EdgeContainers>,
1200 partition_count: usize,
1201 blocks: Vec<ConcurrentBlockedAppend>,
1202 phi_edges: AtomicU64,
1203 diagonal_edges: AtomicU64,
1204}
1205
1206impl ConcurrentBlockedEdgeWriters {
1207 fn new<const K: usize>(matrix: &BlockedEdgeMatrix<K>) -> Self {
1208 Self {
1209 partition_count: matrix.partition_count(),
1210 containers: Arc::clone(&matrix.containers),
1211 blocks: matrix
1212 .blocks
1213 .iter()
1214 .map(|block| ConcurrentBlockedAppend {
1215 extents: Mutex::new(Vec::new()),
1216 record_len: block.record_len,
1217 buffer: Mutex::new(Vec::new()),
1218 edges: AtomicUsize::new(0),
1219 })
1220 .collect(),
1221 phi_edges: AtomicU64::new(0),
1222 diagonal_edges: AtomicU64::new(0),
1223 }
1224 }
1225
1226 #[allow(dead_code)]
1228 fn add(&self, edge: &PreparedBlockedEdge) -> Result<(), SerialCollationError> {
1229 let block = &self.blocks[edge.block];
1230 let mut buffer = block
1231 .buffer
1232 .lock()
1233 .unwrap_or_else(|poison| poison.into_inner());
1234 if buffer.len() + block.record_len > BLOCKED_EDGE_WRITE_BUFFER_BYTES {
1235 self.spill(edge.block, &mut buffer)?;
1236 }
1237 buffer.extend_from_slice(&edge.bytes[..block.record_len]);
1238 block.edges.fetch_add(1, Ordering::Relaxed);
1239 Ok(())
1240 }
1241
1242 fn add_batch(&self, edges: &[PreparedBlockedEdge]) -> Result<(), SerialCollationError> {
1243 if edges.is_empty() {
1244 return Ok(());
1245 }
1246 let block = &self.blocks[edges[0].block];
1247 let mut buffer = block
1248 .buffer
1249 .lock()
1250 .unwrap_or_else(|poison| poison.into_inner());
1251 for edge in edges {
1252 debug_assert_eq!(edge.block, edges[0].block);
1253 if buffer.len() + block.record_len > BLOCKED_EDGE_WRITE_BUFFER_BYTES {
1254 self.spill(edges[0].block, &mut buffer)?;
1255 }
1256 buffer.extend_from_slice(&edge.bytes[..block.record_len]);
1257 }
1258 block.edges.fetch_add(edges.len(), Ordering::Relaxed);
1259 Ok(())
1260 }
1261
1262 fn add_prepared_edges<const K: usize>(
1263 &self,
1264 edges: &mut [PreparedBlockedEdge],
1265 unitig_base: usize,
1266 unitig_bucket: u16,
1267 ) -> Result<(), SerialCollationError> {
1268 if edges.is_empty() {
1269 return Ok(());
1270 }
1271 edges.sort_unstable_by_key(|edge| edge.block);
1272 let unitig_off = blocked_edge_unitig_offset::<K>();
1273 let mut phi = 0u64;
1274 let mut diagonal = 0u64;
1275 let mut start = 0;
1276 while start < edges.len() {
1277 let block_id = edges[start].block;
1278 let mut end = start + 1;
1279 while end < edges.len() && edges[end].block == block_id {
1280 end += 1;
1281 }
1282 let block = &self.blocks[block_id];
1283 let mut buffer = block
1284 .buffer
1285 .lock()
1286 .unwrap_or_else(|poison| poison.into_inner());
1287 for edge in &edges[start..end] {
1288 let mut bytes = edge.bytes;
1289 add_unitig_base_to_encoded_edge::<K>(&mut bytes, unitig_off, unitig_base);
1290 set_encoded_edge_unitig_bucket::<K>(&mut bytes, unitig_bucket);
1291 if buffer.len() + block.record_len > BLOCKED_EDGE_WRITE_BUFFER_BYTES {
1292 self.spill(block_id, &mut buffer)?;
1293 }
1294 buffer.extend_from_slice(&bytes[..block.record_len]);
1295 phi += u64::from(edge.phi);
1297 diagonal += u64::from(edge.diagonal);
1298 }
1299 block.edges.fetch_add(end - start, Ordering::Relaxed);
1300 start = end;
1301 }
1302 self.phi_edges.fetch_add(phi, Ordering::Relaxed);
1303 self.diagonal_edges.fetch_add(diagonal, Ordering::Relaxed);
1304 Ok(())
1305 }
1306
1307 fn finish_into<const K: usize>(
1308 &self,
1309 matrix: &mut BlockedEdgeMatrix<K>,
1310 ) -> Result<(), SerialCollationError> {
1311 let mut edges = 0u64;
1312 for (block_index, (block, target)) in self.blocks.iter().zip(&mut matrix.blocks).enumerate()
1313 {
1314 let mut buffer = block
1315 .buffer
1316 .lock()
1317 .unwrap_or_else(|poison| poison.into_inner());
1318 if !buffer.is_empty() {
1319 self.spill(block_index, &mut buffer)?;
1320 }
1321 for extent in std::mem::take(
1325 &mut *block
1326 .extents
1327 .lock()
1328 .unwrap_or_else(|poison| poison.into_inner()),
1329 ) {
1330 push_coalesced_extent(&mut target.extents, extent);
1331 }
1332 target.edges = block.edges.load(Ordering::Relaxed);
1333 edges += target.edges as u64;
1334 }
1335 matrix.stats.edges = edges;
1336 matrix.stats.phi_edges = self.phi_edges.load(Ordering::Relaxed);
1337 matrix.stats.diagonal_edges = self.diagonal_edges.load(Ordering::Relaxed);
1338 Ok(())
1339 }
1340
1341 fn merge_block_into<const K: usize>(
1353 &self,
1354 matrix: &mut BlockedEdgeMatrix<K>,
1355 row: usize,
1356 col: usize,
1357 ) -> Result<usize, SerialCollationError> {
1358 let block_index = row * self.partition_count + col;
1359 self.flush_block(block_index)?;
1360 let block = &self.blocks[block_index];
1361 let target = &mut matrix.blocks[block_index];
1362 for extent in std::mem::take(
1363 &mut *block
1364 .extents
1365 .lock()
1366 .unwrap_or_else(|poison| poison.into_inner()),
1367 ) {
1368 push_coalesced_extent(&mut target.extents, extent);
1369 }
1370 let added = block.edges.swap(0, Ordering::Relaxed);
1371 target.edges += added;
1372 debug_assert_eq!(
1373 target.extents.iter().map(|e| e.len as usize).sum::<usize>() + target.buffer.len(),
1374 target.edges * target.record_len,
1375 "block {block_index} extents and edge count disagree after merge",
1376 );
1377 Ok(added)
1378 }
1379
1380 fn spill(&self, block_index: usize, buffer: &mut Vec<u8>) -> Result<(), SerialCollationError> {
1386 if buffer.is_empty() {
1387 return Ok(());
1388 }
1389 let extent = self.containers.append(block_index, buffer)?;
1390 let mut extents = self.blocks[block_index]
1391 .extents
1392 .lock()
1393 .unwrap_or_else(|poison| poison.into_inner());
1394 push_coalesced_extent(&mut extents, extent);
1395 buffer.clear();
1396 Ok(())
1397 }
1398
1399 fn flush_block(&self, block_id: usize) -> Result<(), SerialCollationError> {
1400 let block = &self.blocks[block_id];
1401 let mut buffer = block
1402 .buffer
1403 .lock()
1404 .unwrap_or_else(|poison| poison.into_inner());
1405 if !buffer.is_empty() {
1406 self.spill(block_id, &mut buffer)?;
1407 }
1408 Ok(())
1409 }
1410}
1411
1412fn push_coalesced_extent(extents: &mut Vec<BlockExtent>, extent: BlockExtent) {
1414 match extents.last_mut() {
1415 Some(last) if last.offset + last.len == extent.offset => last.len += extent.len,
1416 _ => extents.push(extent),
1417 }
1418}
1419
1420impl<const K: usize> BlockedEdgeMatrix<K> {
1421 fn create(dir: &Path, vertex_partitions: usize) -> Result<Self, SerialCollationError> {
1422 if dir.exists() {
1423 fs::remove_dir_all(dir).map_err(|source| SerialCollationError::Io {
1424 path: dir.to_path_buf(),
1425 source,
1426 })?;
1427 }
1428 fs::create_dir_all(dir).map_err(|source| SerialCollationError::Io {
1429 path: dir.to_path_buf(),
1430 source,
1431 })?;
1432 let partition_count = vertex_partitions + 1;
1433 let record_len = discontinuity_edge_record_len::<K>();
1434 let containers = Arc::new(EdgeContainers::create(dir, partition_count)?);
1435 let mut blocks = Vec::with_capacity(partition_count * partition_count);
1436 blocks.resize_with(partition_count * partition_count, || BlockedEdgeBlock {
1437 record_len,
1438 ..BlockedEdgeBlock::default()
1439 });
1440 Ok(Self {
1441 dir: dir.to_path_buf(),
1442 vertex_partitions,
1443 containers,
1444 blocks,
1445 stats: SerialEdgeMatrixStats::default(),
1446 phantom: PhantomData,
1447 })
1448 }
1449
1450 #[inline]
1451 fn partition_count(&self) -> usize {
1452 self.vertex_partitions + 1
1453 }
1454
1455 #[inline]
1456 #[allow(dead_code)]
1458 fn partition(&self, endpoint: MatrixEndpoint<K>) -> usize {
1459 edge_matrix_partition(self.vertex_partitions, endpoint)
1460 }
1461
1462 #[inline]
1463 fn block_index(&self, row: usize, col: usize) -> usize {
1464 row * self.partition_count() + col
1465 }
1466
1467 #[allow(dead_code)]
1469 fn add_edge_record(
1470 &mut self,
1471 mut edge: DiscontinuityEdge<K>,
1472 ) -> Result<(), SerialCollationError> {
1473 let first_partition = self.partition(edge.first);
1474 let second_partition = self.partition(edge.second);
1475 if first_partition > second_partition {
1476 std::mem::swap(&mut edge.first, &mut edge.second);
1477 edge.unitig_exit_side = edge.unitig_exit_side.inverse();
1478 edge.swapped = !edge.swapped;
1479 }
1480 let row = first_partition.min(second_partition);
1481 let col = first_partition.max(second_partition);
1482 self.stats.edges += 1;
1483 self.stats.phi_edges += u64::from(edge.first.is_phi() || edge.second.is_phi());
1484 self.stats.diagonal_edges += u64::from(row == col);
1485 let idx = self.block_index(row, col);
1486 let containers = Arc::clone(&self.containers);
1487 let record = encode_discontinuity_edge(&edge);
1488 let block = &mut self.blocks[idx];
1489 if block.buffer.len() + block.record_len > BLOCKED_EDGE_WRITE_BUFFER_BYTES {
1490 flush_blocked_edge_block(&containers, idx, block)?;
1491 }
1492 block.buffer.extend_from_slice(&record[..block.record_len]);
1493 block.edges += 1;
1494 if block.buffer.len() + block.record_len > BLOCKED_EDGE_WRITE_BUFFER_BYTES {
1495 flush_blocked_edge_block(&containers, idx, block)?;
1496 }
1497 Ok(())
1498 }
1499
1500 fn add_prepared_edges(
1501 &mut self,
1502 edges: &[PreparedBlockedEdge],
1503 unitig_base: usize,
1504 ) -> Result<(), SerialCollationError> {
1505 let unitig_off = blocked_edge_unitig_offset::<K>();
1506 let containers = Arc::clone(&self.containers);
1507 for edge in edges {
1508 let mut bytes = edge.bytes;
1509 add_unitig_base_to_encoded_edge::<K>(&mut bytes, unitig_off, unitig_base);
1510 let block = &mut self.blocks[edge.block];
1511 if block.buffer.len() + block.record_len > BLOCKED_EDGE_WRITE_BUFFER_BYTES {
1512 flush_blocked_edge_block(&containers, edge.block, block)?;
1513 }
1514 block.buffer.extend_from_slice(&bytes[..block.record_len]);
1515 block.edges += 1;
1516 self.stats.edges += 1;
1517 self.stats.phi_edges += u64::from(edge.phi);
1518 self.stats.diagonal_edges += u64::from(edge.diagonal);
1519 }
1520 Ok(())
1521 }
1522
1523 #[allow(dead_code)]
1525 fn add_prepared_edges_absolute(
1526 &mut self,
1527 edges: &[PreparedBlockedEdge],
1528 ) -> Result<(), SerialCollationError> {
1529 let containers = Arc::clone(&self.containers);
1530 for edge in edges {
1531 let block = &mut self.blocks[edge.block];
1532 if block.buffer.len() + block.record_len > BLOCKED_EDGE_WRITE_BUFFER_BYTES {
1533 flush_blocked_edge_block(&containers, edge.block, block)?;
1534 }
1535 block
1536 .buffer
1537 .extend_from_slice(&edge.bytes[..block.record_len]);
1538 block.edges += 1;
1539 self.stats.edges += 1;
1540 self.stats.phi_edges += u64::from(edge.phi);
1541 self.stats.diagonal_edges += u64::from(edge.diagonal);
1542 }
1543 Ok(())
1544 }
1545
1546 fn flush_block(&mut self, row: usize, col: usize) -> Result<(), SerialCollationError> {
1547 let idx = self.block_index(row, col);
1548 let containers = Arc::clone(&self.containers);
1549 flush_blocked_edge_block(&containers, idx, &mut self.blocks[idx])
1550 }
1551
1552 fn flush_all(&mut self) -> Result<(), SerialCollationError> {
1553 let containers = Arc::clone(&self.containers);
1554 for (index, block) in self.blocks.iter_mut().enumerate() {
1555 flush_blocked_edge_block(&containers, index, block)?;
1556 }
1557 Ok(())
1558 }
1559
1560 fn flush_all_with_threads(&mut self, threads: usize) -> Result<(), SerialCollationError> {
1561 let workers = threads.max(1).min(self.blocks.len());
1562 if workers == 1 {
1563 return self.flush_all();
1564 }
1565 let chunk = self.blocks.len().div_ceil(workers);
1566 let containers = Arc::clone(&self.containers);
1567 std::thread::scope(|scope| {
1568 let containers = &containers;
1569 let mut handles = Vec::with_capacity(workers);
1570 for (group, blocks) in self.blocks.chunks_mut(chunk).enumerate() {
1571 let base = group * chunk;
1572 handles.push(scope.spawn(move || {
1573 for (offset, block) in blocks.iter_mut().enumerate() {
1574 flush_blocked_edge_block(containers, base + offset, block)?;
1575 }
1576 Ok::<_, SerialCollationError>(())
1577 }));
1578 }
1579 for handle in handles {
1580 handle
1581 .join()
1582 .map_err(|_| SerialCollationError::WorkerPanic)??;
1583 }
1584 Ok(())
1585 })
1586 }
1587
1588 fn read_flushed_block(
1589 &self,
1590 row: usize,
1591 col: usize,
1592 ) -> Result<Vec<DiscontinuityEdge<K>>, SerialCollationError> {
1593 let idx = self.block_index(row, col);
1594 let block = &self.blocks[idx];
1595 if block.edges == 0 {
1596 return Ok(Vec::new());
1597 }
1598 let bytes = self.containers.read_block(idx, &block.extents)?;
1599 if bytes.len() + block.buffer.len() != block.edges * block.record_len {
1600 return Err(SerialCollationError::MalformedCoordBucket(
1601 self.containers.path_for(idx),
1602 ));
1603 }
1604 let mut records = Vec::with_capacity(block.edges);
1605 for encoded in bytes.chunks_exact(block.record_len) {
1606 records.push(decode_discontinuity_edge::<K>(encoded));
1607 }
1608 for encoded in block.buffer.chunks_exact(block.record_len) {
1609 records.push(decode_discontinuity_edge::<K>(encoded));
1610 }
1611 Ok(records)
1612 }
1613
1614 #[allow(dead_code)]
1616 fn load_column(
1617 &mut self,
1618 col: usize,
1619 threads: usize,
1620 ) -> Result<SerialEdgeMatrix<K>, SerialCollationError> {
1621 let mut column = SerialEdgeMatrix::new(self.vertex_partitions)
1622 .map_err(|_| SerialCollationError::MalformedCoordBucket(self.dir.clone()))?;
1623 self.load_column_into(col, threads, &mut column)?;
1624 Ok(column)
1625 }
1626
1627 fn load_column_into(
1628 &mut self,
1629 col: usize,
1630 threads: usize,
1631 column: &mut SerialEdgeMatrix<K>,
1632 ) -> Result<(), SerialCollationError> {
1633 for row in 0..=col {
1634 self.flush_block(row, col)?;
1635 }
1636 let workers = threads.max(1).min(col + 1);
1637 let chunk = (col + 1).div_ceil(workers);
1638 let block_groups = std::thread::scope(|scope| {
1639 let mut handles = Vec::with_capacity(workers);
1640 let matrix = &*self;
1641 for row_start in (0..=col).step_by(chunk) {
1642 let row_end = (row_start + chunk).min(col + 1);
1643 handles.push(scope.spawn(move || {
1644 let mut blocks = Vec::with_capacity(row_end - row_start);
1645 for row in row_start..row_end {
1646 blocks.push((row, matrix.read_flushed_block(row, col)?));
1647 }
1648 Ok::<_, SerialCollationError>(blocks)
1649 }));
1650 }
1651 let mut groups = Vec::with_capacity(handles.len());
1652 for handle in handles {
1653 groups.push(
1654 handle
1655 .join()
1656 .map_err(|_| SerialCollationError::WorkerPanic)??,
1657 );
1658 }
1659 Ok::<_, SerialCollationError>(groups)
1660 })?;
1661 column.stats = SerialEdgeMatrixStats::default();
1662 for group in block_groups {
1663 for (row, edges) in group {
1664 column.stats.edges += edges.len() as u64;
1665 column.stats.phi_edges += edges
1666 .iter()
1667 .filter(|edge| edge.first.is_phi() || edge.second.is_phi())
1668 .count() as u64;
1669 column.stats.diagonal_edges += u64::from(row == col) * edges.len() as u64;
1670 let old = std::mem::replace(&mut column.blocks[row][col], edges);
1671 drop(old);
1672 }
1673 }
1674 Ok(())
1675 }
1676
1677 fn read_flushed_row(
1678 &self,
1679 row: usize,
1680 threads: usize,
1681 ) -> Result<Vec<Vec<DiscontinuityEdge<K>>>, SerialCollationError> {
1682 let partition_count = self.partition_count();
1683 let block_count = partition_count - row;
1684 let workers = threads.max(1).min(block_count);
1685 let chunk = block_count.div_ceil(workers);
1686 let groups = std::thread::scope(|scope| {
1687 let mut handles = Vec::with_capacity(workers);
1688 for col_start in (row..partition_count).step_by(chunk) {
1689 let col_end = (col_start + chunk).min(partition_count);
1690 handles.push(scope.spawn(move || {
1691 let mut blocks = Vec::with_capacity(col_end - col_start);
1692 for col in col_start..col_end {
1693 blocks.push((col, self.read_flushed_block(row, col)?));
1694 }
1695 Ok::<_, SerialCollationError>(blocks)
1696 }));
1697 }
1698 let mut groups = Vec::with_capacity(handles.len());
1699 for handle in handles {
1700 groups.push(
1701 handle
1702 .join()
1703 .map_err(|_| SerialCollationError::WorkerPanic)??,
1704 );
1705 }
1706 Ok::<_, SerialCollationError>(groups)
1707 })?;
1708 let mut row_blocks = (0..block_count).map(|_| Vec::new()).collect::<Vec<_>>();
1709 for group in groups {
1710 for (col, edges) in group {
1711 row_blocks[col - row] = edges;
1712 }
1713 }
1714 Ok(row_blocks)
1715 }
1716}
1717
1718#[derive(Debug)]
1719struct ExternalBlockedContraction<const K: usize> {
1720 vertex_partitions: usize,
1721 expansion_matrix: BlockedEdgeMatrix<K>,
1722 pub compressed_diagonal_edges: Vec<Vec<DiscontinuityEdge<K>>>,
1723 meta_vertex_dir: PathBuf,
1724 meta_vertex_count: u64,
1725 meta_vertices_per_partition: Vec<usize>,
1726 stats: FullSerialContractionStats,
1727}
1728
1729#[derive(Default)]
1730struct BlockedContractTimings {
1731 flush: Duration,
1732 clear: Duration,
1733 diagonal: Duration,
1734 read: Duration,
1735 scan: Duration,
1736 join: Duration,
1740 tasks: Duration,
1743 gather: Duration,
1745 finish: Duration,
1746}
1747
1748fn flush_blocked_edge_block(
1749 containers: &EdgeContainers,
1750 block_index: usize,
1751 block: &mut BlockedEdgeBlock,
1752) -> Result<(), SerialCollationError> {
1753 if block.buffer.is_empty() {
1754 return Ok(());
1755 }
1756 let extent = containers.append(block_index, &block.buffer)?;
1757 push_coalesced_extent(&mut block.extents, extent);
1758 block.buffer.clear();
1759 Ok(())
1760}
1761
1762fn encode_discontinuity_edge<const K: usize>(edge: &DiscontinuityEdge<K>) -> [u8; 72] {
1763 let mut bytes = [0u8; 72];
1764 if K <= 31 {
1765 let endpoint_bits = |endpoint: MatrixEndpoint<K>| match endpoint {
1766 MatrixEndpoint::Phi => u64::MAX,
1767 MatrixEndpoint::Vertex(endpoint) => endpoint.vertex.as_u128() as u64,
1768 };
1769 bytes[..8].copy_from_slice(&endpoint_bits(edge.first).to_le_bytes());
1770 bytes[8..16].copy_from_slice(&endpoint_bits(edge.second).to_le_bytes());
1771 let weight = u16::try_from(edge.weight).expect("discontinuity edge weight fits u16");
1772 bytes[16..18].copy_from_slice(&weight.to_le_bytes());
1773 let unitig_index = if edge.unitig_index == usize::MAX {
1774 u32::MAX
1775 } else {
1776 u32::try_from(edge.unitig_index).expect("local unitig index fits compact edge")
1777 };
1778 bytes[18..22].copy_from_slice(&unitig_index.to_le_bytes());
1779 bytes[22] = u8::from(
1780 matches!(edge.first, MatrixEndpoint::Vertex(endpoint) if endpoint.side == Side::Back),
1781 ) | (u8::from(
1782 matches!(edge.second, MatrixEndpoint::Vertex(endpoint) if endpoint.side == Side::Back),
1783 ) << 1)
1784 | (u8::from(edge.unitig_exit_side == Side::Back) << 2)
1785 | (u8::from(edge.swapped) << 3)
1786 | (u8::from(edge.phantom_unitig.is_some()) << 4);
1787 bytes[23..25].copy_from_slice(&edge.unitig_bucket.to_le_bytes());
1788 return bytes;
1789 }
1790 let kmer_bytes = discontinuity_edge_kmer_bytes::<K>();
1791 let endpoint_len = kmer_bytes + 2;
1792 encode_matrix_endpoint(edge.first, &mut bytes[0..endpoint_len], kmer_bytes);
1793 encode_matrix_endpoint(
1794 edge.second,
1795 &mut bytes[endpoint_len..2 * endpoint_len],
1796 kmer_bytes,
1797 );
1798 let weight_off = 2 * endpoint_len;
1799 let unitig_off = weight_off + 8;
1800 let flags_off = unitig_off + 8;
1801 bytes[weight_off..unitig_off].copy_from_slice(&edge.weight.to_le_bytes());
1802 bytes[unitig_off..flags_off].copy_from_slice(&(edge.unitig_index as u64).to_le_bytes());
1803 bytes[flags_off] = edge.unitig_exit_side as u8;
1804 bytes[flags_off + 1] = u8::from(edge.swapped);
1805 if let Some(endpoint) = edge.phantom_unitig {
1806 bytes[flags_off + 2] = 1;
1807 let phantom_off = flags_off + 3;
1808 bytes[phantom_off..phantom_off + kmer_bytes]
1809 .copy_from_slice(&endpoint.vertex.as_u128().to_le_bytes()[..kmer_bytes]);
1810 bytes[phantom_off + kmer_bytes] = endpoint.side as u8;
1811 }
1812 let bucket_off = discontinuity_edge_record_len::<K>() - 2;
1813 bytes[bucket_off..bucket_off + 2].copy_from_slice(&edge.unitig_bucket.to_le_bytes());
1814 bytes
1815}
1816
1817fn encode_matrix_endpoint<const K: usize>(
1818 endpoint: MatrixEndpoint<K>,
1819 dst: &mut [u8],
1820 kmer_bytes: usize,
1821) {
1822 if let MatrixEndpoint::Vertex(endpoint) = endpoint {
1823 dst[0] = 1;
1824 dst[1..1 + kmer_bytes]
1825 .copy_from_slice(&endpoint.vertex.as_u128().to_le_bytes()[..kmer_bytes]);
1826 dst[1 + kmer_bytes] = endpoint.side as u8;
1827 }
1828}
1829
1830fn decode_discontinuity_edge<const K: usize>(bytes: &[u8]) -> DiscontinuityEdge<K> {
1831 if K <= 31 {
1832 let read_u64 = |range: std::ops::Range<usize>| {
1833 u64::from_le_bytes(bytes[range].try_into().expect("eight-byte edge field"))
1834 };
1835 let flags = bytes[22];
1836 let endpoint = |bits: u64, side_bit: u8| {
1837 if bits == u64::MAX {
1838 MatrixEndpoint::Phi
1839 } else {
1840 MatrixEndpoint::Vertex(DiscontinuityEndpoint {
1841 vertex: Kmer::from_bits(bits as u128),
1842 side: if flags & side_bit == 0 {
1843 Side::Front
1844 } else {
1845 Side::Back
1846 },
1847 })
1848 }
1849 };
1850 let first = endpoint(read_u64(0..8), 1);
1851 let second = endpoint(read_u64(8..16), 2);
1852 let raw_unitig = u32::from_le_bytes(bytes[18..22].try_into().unwrap());
1853 let phantom_unitig = if flags & (1 << 4) == 0 {
1854 None
1855 } else {
1856 match (first, second) {
1857 (MatrixEndpoint::Vertex(endpoint), MatrixEndpoint::Phi)
1858 | (MatrixEndpoint::Phi, MatrixEndpoint::Vertex(endpoint)) => Some(endpoint),
1859 _ => None,
1860 }
1861 };
1862 return DiscontinuityEdge {
1863 first,
1864 second,
1865 weight: u64::from(u16::from_le_bytes(bytes[16..18].try_into().unwrap())),
1866 unitig_bucket: u16::from_le_bytes(bytes[23..25].try_into().unwrap()),
1867 unitig_index: if raw_unitig == u32::MAX {
1868 usize::MAX
1869 } else {
1870 raw_unitig as usize
1871 },
1872 unitig_exit_side: if flags & (1 << 2) == 0 {
1873 Side::Front
1874 } else {
1875 Side::Back
1876 },
1877 phantom_unitig,
1878 swapped: flags & (1 << 3) != 0,
1879 };
1880 }
1881 let kmer_bytes = discontinuity_edge_kmer_bytes::<K>();
1882 let endpoint_len = kmer_bytes + 2;
1883 let endpoint = |src: &[u8]| {
1884 if src[0] == 0 {
1885 MatrixEndpoint::Phi
1886 } else {
1887 let mut bits = [0u8; 16];
1888 bits[..kmer_bytes].copy_from_slice(&src[1..1 + kmer_bytes]);
1889 MatrixEndpoint::Vertex(DiscontinuityEndpoint {
1890 vertex: Kmer::from_bits(u128::from_le_bytes(bits)),
1891 side: decode_side(src[1 + kmer_bytes]),
1892 })
1893 }
1894 };
1895 let weight_off = 2 * endpoint_len;
1896 let unitig_off = weight_off + 8;
1897 let flags_off = unitig_off + 8;
1898 let mut weight = [0u8; 8];
1899 weight.copy_from_slice(&bytes[weight_off..unitig_off]);
1900 let mut unitig = [0u8; 8];
1901 unitig.copy_from_slice(&bytes[unitig_off..flags_off]);
1902 let phantom_unitig = (bytes[flags_off + 2] != 0).then(|| {
1903 let phantom_off = flags_off + 3;
1904 let mut bits = [0u8; 16];
1905 bits[..kmer_bytes].copy_from_slice(&bytes[phantom_off..phantom_off + kmer_bytes]);
1906 DiscontinuityEndpoint {
1907 vertex: Kmer::from_bits(u128::from_le_bytes(bits)),
1908 side: decode_side(bytes[phantom_off + kmer_bytes]),
1909 }
1910 });
1911 DiscontinuityEdge {
1912 first: endpoint(&bytes[0..endpoint_len]),
1913 second: endpoint(&bytes[endpoint_len..2 * endpoint_len]),
1914 weight: u64::from_le_bytes(weight),
1915 unitig_bucket: u16::from_le_bytes(
1916 bytes[discontinuity_edge_record_len::<K>() - 2..discontinuity_edge_record_len::<K>()]
1917 .try_into()
1918 .unwrap(),
1919 ),
1920 unitig_index: u64::from_le_bytes(unitig) as usize,
1921 unitig_exit_side: decode_side(bytes[flags_off]),
1922 phantom_unitig,
1923 swapped: bytes[flags_off + 1] != 0,
1924 }
1925}
1926
1927#[inline]
1928fn decode_partition_incoming<const K: usize>(
1929 bytes: &[u8],
1930) -> Option<(Kmer<K>, PartitionOtherEnd<K>)> {
1931 if K <= 31 {
1932 let first = u64::from_le_bytes(bytes[..8].try_into().unwrap());
1933 let second = u64::from_le_bytes(bytes[8..16].try_into().unwrap());
1934 if second == u64::MAX {
1935 return None;
1936 }
1937 let flags = bytes[22];
1938 let endpoint = if first == u64::MAX {
1939 MatrixEndpoint::Phi
1940 } else {
1941 MatrixEndpoint::Vertex(DiscontinuityEndpoint {
1942 vertex: Kmer::from_bits(first as u128),
1943 side: if flags & 1 == 0 {
1944 Side::Front
1945 } else {
1946 Side::Back
1947 },
1948 })
1949 };
1950 return Some((
1951 Kmer::from_bits(second as u128),
1952 PartitionOtherEnd {
1953 endpoint,
1954 side_at_current: if flags & 2 == 0 {
1955 Side::Front
1956 } else {
1957 Side::Back
1958 },
1959 weight: u64::from(u16::from_le_bytes(bytes[16..18].try_into().unwrap())),
1960 in_same_part: false,
1961 processed: false,
1962 },
1963 ));
1964 }
1965 let kmer_bytes = discontinuity_edge_kmer_bytes::<K>();
1966 let endpoint_len = kmer_bytes + 2;
1967 let decode_endpoint = |src: &[u8]| {
1968 if src[0] == 0 {
1969 MatrixEndpoint::Phi
1970 } else {
1971 let mut bits = [0u8; 16];
1972 bits[..kmer_bytes].copy_from_slice(&src[1..1 + kmer_bytes]);
1973 MatrixEndpoint::Vertex(DiscontinuityEndpoint {
1974 vertex: Kmer::from_bits(u128::from_le_bytes(bits)),
1975 side: decode_side(src[1 + kmer_bytes]),
1976 })
1977 }
1978 };
1979 let lower = decode_endpoint(&bytes[..endpoint_len]);
1980 let MatrixEndpoint::Vertex(current) = decode_endpoint(&bytes[endpoint_len..2 * endpoint_len])
1981 else {
1982 return None;
1983 };
1984 let weight_off = 2 * endpoint_len;
1985 let mut weight = [0u8; 8];
1986 weight.copy_from_slice(&bytes[weight_off..weight_off + 8]);
1987 Some((
1988 current.vertex,
1989 PartitionOtherEnd {
1990 endpoint: lower,
1991 side_at_current: current.side,
1992 weight: u64::from_le_bytes(weight),
1993 in_same_part: false,
1994 processed: false,
1995 },
1996 ))
1997}
1998
1999#[inline]
2000const fn discontinuity_edge_kmer_bytes<const K: usize>() -> usize {
2001 (2 * K).div_ceil(8)
2002}
2003
2004#[inline]
2005const fn discontinuity_edge_record_len<const K: usize>() -> usize {
2006 if K <= 31 {
2007 25
2008 } else {
2009 3 * discontinuity_edge_kmer_bytes::<K>() + 24
2010 }
2011}
2012
2013#[inline]
2014const fn blocked_edge_unitig_offset<const K: usize>() -> usize {
2015 if K <= 31 {
2016 18
2017 } else {
2018 2 * (discontinuity_edge_kmer_bytes::<K>() + 2) + 8
2019 }
2020}
2021
2022fn add_unitig_base_to_encoded_edge<const K: usize>(
2023 bytes: &mut [u8; 72],
2024 unitig_off: usize,
2025 unitig_base: usize,
2026) {
2027 if K <= 31 {
2028 let local = u32::from_le_bytes(bytes[unitig_off..unitig_off + 4].try_into().unwrap());
2029 if local != u32::MAX {
2030 let index = unitig_base
2031 .checked_add(local as usize)
2032 .and_then(|index| u32::try_from(index).ok())
2033 .expect("global local-unitig index fits compact edge");
2034 bytes[unitig_off..unitig_off + 4].copy_from_slice(&index.to_le_bytes());
2035 }
2036 } else {
2037 let local = u64::from_le_bytes(bytes[unitig_off..unitig_off + 8].try_into().unwrap());
2038 if local != u64::MAX {
2039 let index = unitig_base + local as usize;
2040 bytes[unitig_off..unitig_off + 8].copy_from_slice(&(index as u64).to_le_bytes());
2041 }
2042 }
2043}
2044
2045#[inline]
2046fn set_encoded_edge_unitig_bucket<const K: usize>(bytes: &mut [u8; 72], bucket: u16) {
2047 let offset = discontinuity_edge_record_len::<K>() - 2;
2048 bytes[offset..offset + 2].copy_from_slice(&bucket.to_le_bytes());
2049}
2050
2051#[inline]
2052fn decode_side(value: u8) -> Side {
2053 if value == Side::Front as u8 {
2054 Side::Front
2055 } else {
2056 Side::Back
2057 }
2058}
2059
2060#[derive(Debug, Clone, PartialEq, Eq)]
2061pub struct SerialContraction {
2062 pub components: Vec<SerialComponent>,
2063 pub stats: SerialContractionStats,
2064}
2065
2066pub const DISCONTINUITY_PARALLELIZATION_OPPORTUNITIES: &[&str] = &[
2067 "diagonal block compression per partition can run concurrently with non-diagonal column scans",
2068 "non-diagonal blocks in a partition column can be scanned independently with a shared vertex table",
2069 "compressed diagonal-chain edges can be processed independently after the column scan",
2070 "false-phantom filtering is a parallel scan over unprocessed partition-table entries",
2071 "different output edge/meta-vertex buckets can use batched thread-local buffers before external-memory flush",
2072];
2073
2074#[derive(Debug, Clone, PartialEq, Eq)]
2075pub struct DiagonalCompression<const K: usize> {
2076 pub partition: usize,
2077 pub edges: Vec<DiscontinuityEdge<K>>,
2078 pub expansion_edges: Vec<DiscontinuityEdge<K>>,
2079 pub meta_vertices: Vec<SerialMetaVertex<K>>,
2080 pub stats: DiagonalCompressionStats,
2081}
2082
2083#[derive(Debug, Clone, PartialEq, Eq)]
2084pub struct SerialMetaVertex<const K: usize> {
2085 pub vertex: Kmer<K>,
2086 pub partition: usize,
2087 pub entry_side: Side,
2088 pub weight: u64,
2089 pub is_cycle: bool,
2090}
2091
2092#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
2093pub struct DiagonalCompressionStats {
2094 pub input_edges: u64,
2095 pub compressed_edges: u64,
2096 pub meta_vertices: u64,
2097 pub isolated_cordless_cycles: u64,
2098}
2099
2100#[derive(Debug, Clone, PartialEq, Eq)]
2101pub struct PartitionContraction<const K: usize> {
2102 pub partition: usize,
2103 pub edges: Vec<DiscontinuityEdge<K>>,
2104 pub expansion_edges: Vec<DiscontinuityEdge<K>>,
2105 pub compressed_diagonal_edges: Vec<DiscontinuityEdge<K>>,
2106 pub meta_vertices: Vec<SerialMetaVertex<K>>,
2107 pub stats: PartitionContractionStats,
2108}
2109
2110#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
2111pub struct PartitionContractionStats {
2112 pub input_non_diagonal_edges: u64,
2113 pub compressed_diagonal_edges: u64,
2114 pub output_edges: u64,
2115 pub meta_vertices: u64,
2116 pub phantom_edges: u64,
2117 pub isolated_cordless_cycles: u64,
2118}
2119
2120#[derive(Debug, Clone, PartialEq, Eq)]
2121pub struct FullSerialDiscontinuityContraction<const K: usize> {
2122 pub vertex_partitions: usize,
2123 pub final_edges: Vec<DiscontinuityEdge<K>>,
2124 pub expansion_edges: Vec<DiscontinuityEdge<K>>,
2125 pub expansion_matrix: SerialEdgeMatrix<K>,
2126 pub compressed_diagonal_edges: Vec<Vec<DiscontinuityEdge<K>>>,
2127 pub meta_vertices: Vec<SerialMetaVertex<K>>,
2128 pub partitions: Vec<PartitionContraction<K>>,
2129 pub stats: FullSerialContractionStats,
2130}
2131
2132#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
2133pub struct FullSerialContractionStats {
2134 pub partitions: u64,
2135 pub input_edges: u64,
2136 pub partition_output_edges: u64,
2137 pub reinserted_edges: u64,
2138 pub final_edges: u64,
2139 pub meta_vertices: u64,
2140 pub phantom_edges: u64,
2141 pub isolated_cordless_cycles: u64,
2142}
2143
2144#[derive(Debug, Clone, Copy, PartialEq, Eq)]
2145pub struct PathInfo<const K: usize> {
2146 pub path_id: Kmer<K>,
2147 pub rank: u64,
2148 pub exit_side: Side,
2149 pub is_cycle: bool,
2150}
2151
2152#[derive(Debug, Clone, PartialEq, Eq)]
2153pub struct VertexPathInfo<const K: usize> {
2154 pub vertex: Kmer<K>,
2155 pub info: PathInfo<K>,
2156}
2157
2158#[derive(Debug, Clone, PartialEq, Eq)]
2159pub struct EdgePathInfo<const K: usize> {
2160 pub unitig_index: usize,
2161 pub phantom_unitig: Option<DiscontinuityEndpoint<K>>,
2162 pub info: PathInfo<K>,
2163}
2164
2165#[derive(Debug, Clone, PartialEq, Eq)]
2166pub struct SerialExpansion<const K: usize> {
2167 pub vertices: Vec<VertexPathInfo<K>>,
2168 pub edges: Vec<EdgePathInfo<K>>,
2169 pub stats: SerialExpansionStats,
2170}
2171
2172#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
2173pub struct SerialExpansionStats {
2174 pub seed_vertices: u64,
2175 pub inferred_vertices: u64,
2176 pub edge_path_infos: u64,
2177 pub unresolved_edges: u64,
2178}
2179
2180#[derive(Debug, Clone, PartialEq, Eq)]
2181pub struct SerialCollation {
2182 pub unitigs: Vec<Vec<u8>>,
2183 pub stats: SerialCollationStats,
2184}
2185
2186#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
2187pub struct SerialCollationStats {
2188 pub input_path_infos: u64,
2189 pub emitted_unitigs: u64,
2190 pub emitted_bases: u64,
2191 pub missing_unitig_labels: u64,
2192 pub direct_local_unitigs: u64,
2193 pub stitched_discontinuity_unitigs: u64,
2194}
2195
2196#[derive(Debug)]
2197pub enum SerialCollationError {
2198 Io {
2199 path: std::path::PathBuf,
2200 source: std::io::Error,
2201 },
2202 MalformedCoordBucket(std::path::PathBuf),
2203 WorkerPanic,
2204 Color(ColorError),
2205}
2206
2207impl std::fmt::Display for SerialCollationError {
2208 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
2209 match self {
2210 Self::Io { path, source } => write!(f, "{}: {source}", path.display()),
2211 Self::MalformedCoordBucket(path) => {
2212 write!(
2213 f,
2214 "malformed stitched coordinate bucket: {}",
2215 path.display()
2216 )
2217 }
2218 Self::WorkerPanic => write!(f, "stitched coordinate worker thread panicked"),
2219 Self::Color(err) => write!(f, "{err}"),
2220 }
2221 }
2222}
2223
2224impl std::error::Error for SerialCollationError {}
2225
2226impl From<ColorError> for SerialCollationError {
2227 fn from(value: ColorError) -> Self {
2228 Self::Color(value)
2229 }
2230}
2231
2232#[derive(Debug, Clone, PartialEq, Eq)]
2233pub struct SerialComponent {
2234 pub endpoints: usize,
2235 pub edges: u64,
2236 pub weight: u64,
2237 pub phi_edges: u64,
2238 pub cyclic: bool,
2239}
2240
2241#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
2242pub struct SerialContractionStats {
2243 pub input_edges: u64,
2244 pub components: u64,
2245 pub phi_edges: u64,
2246 pub cyclic_components: u64,
2247}
2248
2249#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
2250struct EndpointKey<const K: usize>(MatrixEndpoint<K>);
2251
2252pub struct SerialDiscontinuityContractor;
2257
2258struct PartitionColumnScan<const K: usize> {
2259 table: FastHashMap<Kmer<K>, PartitionOtherEnd<K>>,
2260 edges: Vec<DiscontinuityEdge<K>>,
2261 meta_vertices: Vec<SerialMetaVertex<K>>,
2262 input_non_diagonal_edges: u64,
2263}
2264
2265struct PartitionColumnScanOutput<const K: usize> {
2266 edges: Vec<DiscontinuityEdge<K>>,
2267 meta_vertices: Vec<SerialMetaVertex<K>>,
2268 input_non_diagonal_edges: u64,
2269}
2270
2271#[derive(Debug, Clone, Copy)]
2272struct PartitionColumnIncoming<const K: usize> {
2273 vertex: Kmer<K>,
2274 end: PartitionOtherEnd<K>,
2275}
2276
2277fn partition_column_vertex_shard<const K: usize>(vertex: Kmer<K>, shard_mask: usize) -> usize {
2278 let bits = vertex.as_u128();
2279 let mixed = hash_u64((bits as u64) ^ ((bits >> 64) as u64), 0);
2280 (mixed as usize) & shard_mask
2281}
2282
2283impl SerialDiscontinuityContractor {
2284 pub fn compress_diagonal_block<const K: usize>(
2285 matrix: &SerialEdgeMatrix<K>,
2286 partition: usize,
2287 ) -> DiagonalCompression<K> {
2288 let mut ends = FastHashMap::with_hasher(FastBuildHasher::default());
2289 Self::compress_diagonal_block_with_ends(matrix, partition, &mut ends)
2290 }
2291
2292 fn compress_diagonal_block_with_ends<const K: usize>(
2293 matrix: &SerialEdgeMatrix<K>,
2294 partition: usize,
2295 ends: &mut FastHashMap<Kmer<K>, DiagonalOtherEnd<K>>,
2296 ) -> DiagonalCompression<K> {
2297 assert!(partition > 0 && partition < matrix.partition_count());
2298
2299 let diagonal_edges = matrix.block(partition, partition);
2300 ends.clear();
2301 ends.reserve(diagonal_edges.len().saturating_mul(2));
2302 let mut expansion_edges = Vec::with_capacity(diagonal_edges.len());
2303 let mut meta_vertices = Vec::new();
2304 let mut isolated_cordless_cycles = 0u64;
2305
2306 for edge in diagonal_edges {
2307 let (x, y) = match (edge.first, edge.second) {
2308 (MatrixEndpoint::Vertex(x), MatrixEndpoint::Vertex(y)) => (x, y),
2309 _ => continue,
2310 };
2311
2312 let end_x = ends.get(&x.vertex).copied();
2313 let end_y = ends.get(&y.vertex).copied();
2314
2315 let u = end_x
2316 .map(|end| DiscontinuityEndpoint {
2317 vertex: end.vertex,
2318 side: end.side_at_vertex,
2319 })
2320 .unwrap_or(x);
2321 let v = end_y
2322 .map(|end| DiscontinuityEndpoint {
2323 vertex: end.vertex,
2324 side: end.side_at_vertex,
2325 })
2326 .unwrap_or(y);
2327 let weight =
2328 end_x.map_or(0, |end| end.weight) + edge.weight + end_y.map_or(0, |end| end.weight);
2329
2330 assert_eq!(matrix.partition(MatrixEndpoint::Vertex(u)), partition);
2331 assert_eq!(matrix.partition(MatrixEndpoint::Vertex(v)), partition);
2332
2333 if u.vertex == v.vertex {
2334 ends.insert(
2335 u.vertex,
2336 DiagonalOtherEnd {
2337 vertex: Kmer::zero(),
2338 side_at_vertex: Side::Back,
2339 side_at_current: Side::Front,
2340 weight: 1,
2341 unitig_index: 0,
2342 unitig_exit_side: Side::Back,
2343 is_phi: true,
2344 },
2345 );
2346 meta_vertices.push(SerialMetaVertex {
2347 vertex: u.vertex,
2348 partition,
2349 entry_side: Side::Front.inverse(),
2350 weight: 1,
2351 is_cycle: true,
2352 });
2353 isolated_cordless_cycles += 1;
2354 } else if u.vertex == y.vertex && v.vertex == x.vertex {
2355 ends.insert(
2356 u.vertex,
2357 DiagonalOtherEnd {
2358 vertex: Kmer::zero(),
2359 side_at_vertex: Side::Back,
2360 side_at_current: u.side.inverse(),
2361 weight: 1,
2362 unitig_index: 0,
2363 unitig_exit_side: Side::Back,
2364 is_phi: true,
2365 },
2366 );
2367 ends.insert(
2368 v.vertex,
2369 DiagonalOtherEnd {
2370 vertex: Kmer::zero(),
2371 side_at_vertex: Side::Back,
2372 side_at_current: v.side.inverse(),
2373 weight: 1,
2374 unitig_index: 0,
2375 unitig_exit_side: Side::Back,
2376 is_phi: true,
2377 },
2378 );
2379 meta_vertices.push(SerialMetaVertex {
2380 vertex: u.vertex,
2381 partition,
2382 entry_side: u.side,
2383 weight: 1,
2384 is_cycle: true,
2385 });
2386 isolated_cordless_cycles += 1;
2387 } else {
2388 expansion_edges.push(DiscontinuityEdge {
2389 first: MatrixEndpoint::Vertex(DiscontinuityEndpoint {
2390 vertex: u.vertex,
2391 side: u.side,
2392 }),
2393 second: MatrixEndpoint::Vertex(DiscontinuityEndpoint {
2394 vertex: v.vertex,
2395 side: v.side,
2396 }),
2397 weight,
2398 unitig_bucket: if weight == 1 { edge.unitig_bucket } else { 0 },
2399 unitig_index: if weight == 1 { edge.unitig_index } else { 0 },
2400 unitig_exit_side: if weight == 1 {
2401 edge.unitig_exit_side
2402 } else {
2403 Side::Back
2404 },
2405 phantom_unitig: None,
2406 swapped: false,
2407 });
2408 ends.insert(
2409 u.vertex,
2410 DiagonalOtherEnd {
2411 vertex: v.vertex,
2412 side_at_vertex: v.side,
2413 side_at_current: u.side,
2414 weight,
2415 unitig_index: 0,
2416 unitig_exit_side: Side::Back,
2417 is_phi: false,
2418 },
2419 );
2420 ends.insert(
2421 v.vertex,
2422 DiagonalOtherEnd {
2423 vertex: u.vertex,
2424 side_at_vertex: u.side,
2425 side_at_current: v.side,
2426 weight,
2427 unitig_index: 0,
2428 unitig_exit_side: Side::Back,
2429 is_phi: false,
2430 },
2431 );
2432 }
2433 }
2434
2435 let mut compressed_edges = Vec::new();
2436 for (&vertex, &end) in ends.iter() {
2437 if end.is_phi {
2438 continue;
2439 }
2440
2441 if vertex < end.vertex {
2442 let Some(reverse) = ends.get(&end.vertex) else {
2443 continue;
2444 };
2445 if reverse.vertex == vertex {
2446 compressed_edges.push(DiscontinuityEdge {
2447 first: MatrixEndpoint::Vertex(DiscontinuityEndpoint {
2448 vertex,
2449 side: end.side_at_current,
2450 }),
2451 second: MatrixEndpoint::Vertex(DiscontinuityEndpoint {
2452 vertex: end.vertex,
2453 side: end.side_at_vertex,
2454 }),
2455 weight: end.weight,
2456 unitig_bucket: 0,
2457 unitig_index: end.unitig_index,
2458 unitig_exit_side: end.unitig_exit_side,
2459 phantom_unitig: None,
2460 swapped: false,
2461 });
2462 }
2463 }
2464 }
2465 DiagonalCompression {
2466 partition,
2467 stats: DiagonalCompressionStats {
2468 input_edges: diagonal_edges.len() as u64,
2469 compressed_edges: compressed_edges.len() as u64,
2470 meta_vertices: meta_vertices.len() as u64,
2471 isolated_cordless_cycles,
2472 },
2473 edges: compressed_edges,
2474 expansion_edges,
2475 meta_vertices,
2476 }
2477 }
2478
2479 pub fn contract_partition<const K: usize>(
2480 matrix: &SerialEdgeMatrix<K>,
2481 partition: usize,
2482 ) -> PartitionContraction<K> {
2483 Self::contract_partition_with_threads(matrix, partition, 1)
2484 }
2485
2486 pub fn contract_partition_with_threads<const K: usize>(
2487 matrix: &SerialEdgeMatrix<K>,
2488 partition: usize,
2489 threads: usize,
2490 ) -> PartitionContraction<K> {
2491 assert!(partition > 0 && partition < matrix.partition_count());
2492
2493 let (diagonal, scan) = if threads > 1 {
2494 std::thread::scope(|scope| {
2495 let diagonal = scope.spawn(|| Self::compress_diagonal_block(matrix, partition));
2496 let scan = Self::scan_partition_column(matrix, partition);
2497 (
2498 diagonal
2499 .join()
2500 .expect("diagonal compression worker panicked"),
2501 scan,
2502 )
2503 })
2504 } else {
2505 (
2506 Self::compress_diagonal_block(matrix, partition),
2507 Self::scan_partition_column(matrix, partition),
2508 )
2509 };
2510
2511 let mut table = scan.table;
2512 let mut edges = scan.edges;
2513 let mut expansion_edges = Vec::new();
2514 let mut meta_vertices = scan.meta_vertices;
2515 meta_vertices.extend(diagonal.meta_vertices);
2516 let mut phantom_edges = 0u64;
2517
2518 for edge in &diagonal.edges {
2522 let (MatrixEndpoint::Vertex(x), MatrixEndpoint::Vertex(y)) = (edge.first, edge.second)
2523 else {
2524 continue;
2525 };
2526
2527 table.entry(x.vertex).or_insert_with(|| {
2528 phantom_edges += 1;
2529 let phantom = DiscontinuityEndpoint {
2530 vertex: x.vertex,
2531 side: x.side.inverse(),
2532 };
2533 expansion_edges.push(join_other_ends_with_phantom(
2534 MatrixEndpoint::Phi,
2535 MatrixEndpoint::Vertex(phantom),
2536 1,
2537 Some(phantom),
2538 ));
2539 PartitionOtherEnd {
2540 endpoint: MatrixEndpoint::Phi,
2541 side_at_current: x.side.inverse(),
2542 weight: 1,
2543 in_same_part: false,
2544 processed: false,
2545 }
2546 });
2547
2548 table.entry(y.vertex).or_insert_with(|| {
2549 phantom_edges += 1;
2550 let phantom = DiscontinuityEndpoint {
2551 vertex: y.vertex,
2552 side: y.side.inverse(),
2553 };
2554 expansion_edges.push(join_other_ends_with_phantom(
2555 MatrixEndpoint::Phi,
2556 MatrixEndpoint::Vertex(phantom),
2557 1,
2558 Some(phantom),
2559 ));
2560 PartitionOtherEnd {
2561 endpoint: MatrixEndpoint::Phi,
2562 side_at_current: y.side.inverse(),
2563 weight: 1,
2564 in_same_part: false,
2565 processed: false,
2566 }
2567 });
2568
2569 let x_end = table.get(&x.vertex).copied().unwrap();
2570 let y_end = table.get(&y.vertex).copied().unwrap();
2571 if x_end.endpoint.is_phi() && y_end.endpoint.is_phi() {
2572 meta_vertices.push(two_weight_meta_vertex(
2573 x.vertex,
2574 partition,
2575 x.side.inverse(),
2576 x_end.weight,
2577 edge.weight + y_end.weight,
2578 false,
2579 ));
2580 } else {
2581 edges.push(join_other_ends(
2582 x_end.endpoint,
2583 y_end.endpoint,
2584 x_end.weight + edge.weight + y_end.weight,
2585 ));
2586 }
2587
2588 if let Some(end) = table.get_mut(&x.vertex) {
2589 end.processed = true;
2590 }
2591 if let Some(end) = table.get_mut(&y.vertex) {
2592 end.processed = true;
2593 }
2594 }
2595
2596 let unprocessed = table
2598 .iter()
2599 .filter_map(|(&vertex, &end)| (!end.processed).then_some((vertex, end)))
2600 .collect::<Vec<_>>();
2601 for (vertex, end) in unprocessed {
2602 phantom_edges += 1;
2603 let phantom = DiscontinuityEndpoint {
2604 vertex,
2605 side: end.side_at_current.inverse(),
2606 };
2607 expansion_edges.push(join_other_ends_with_phantom(
2608 MatrixEndpoint::Phi,
2609 MatrixEndpoint::Vertex(phantom),
2610 1,
2611 Some(phantom),
2612 ));
2613
2614 if end.endpoint.is_phi() {
2615 meta_vertices.push(two_weight_meta_vertex(
2616 vertex,
2617 partition,
2618 end.side_at_current,
2619 end.weight,
2620 1,
2621 false,
2622 ));
2623 } else {
2624 edges.push(join_other_ends(
2625 MatrixEndpoint::Phi,
2626 end.endpoint,
2627 1 + end.weight,
2628 ));
2629 }
2630 }
2631
2632 edges.sort_by_key(|edge| {
2633 (
2634 endpoint_sort_key(edge.first),
2635 endpoint_sort_key(edge.second),
2636 edge.weight,
2637 )
2638 });
2639 meta_vertices.sort_by_key(|meta| {
2640 (
2641 meta.partition,
2642 meta.vertex.as_u128(),
2643 meta.entry_side as u8,
2644 meta.weight,
2645 meta.is_cycle,
2646 )
2647 });
2648
2649 PartitionContraction {
2650 partition,
2651 compressed_diagonal_edges: diagonal.expansion_edges,
2652 expansion_edges,
2653 stats: PartitionContractionStats {
2654 input_non_diagonal_edges: scan.input_non_diagonal_edges,
2655 compressed_diagonal_edges: diagonal.stats.compressed_edges,
2656 output_edges: edges.len() as u64,
2657 meta_vertices: meta_vertices.len() as u64,
2658 phantom_edges,
2659 isolated_cordless_cycles: diagonal.stats.isolated_cordless_cycles,
2660 },
2661 edges,
2662 meta_vertices,
2663 }
2664 }
2665
2666 fn scan_partition_column<const K: usize>(
2667 matrix: &SerialEdgeMatrix<K>,
2668 partition: usize,
2669 ) -> PartitionColumnScan<K> {
2670 let column_edges = (0..partition)
2671 .map(|row| matrix.block(row, partition).len())
2672 .sum::<usize>();
2673 let mut table = FastHashMap::<Kmer<K>, PartitionOtherEnd<K>>::with_capacity_and_hasher(
2674 column_edges,
2675 FastBuildHasher::default(),
2676 );
2677 let mut edges = Vec::with_capacity(column_edges / 2);
2678 let mut meta_vertices = Vec::new();
2679 let mut input_non_diagonal_edges = 0u64;
2680
2681 for row in 0..partition {
2685 for edge in matrix.block(row, partition) {
2686 let (lower, current) = match endpoint_in_partition(matrix, edge, partition) {
2687 Some(pair) => pair,
2688 None => continue,
2689 };
2690 input_non_diagonal_edges += 1;
2691
2692 let incoming = PartitionOtherEnd {
2693 endpoint: lower,
2694 side_at_current: current.side,
2695 weight: edge.weight,
2696 in_same_part: false,
2697 processed: false,
2698 };
2699
2700 absorb_partition_other_end(
2701 current.vertex,
2702 incoming,
2703 partition,
2704 &mut table,
2705 &mut edges,
2706 &mut meta_vertices,
2707 );
2708 }
2709 }
2710
2711 PartitionColumnScan {
2712 table,
2713 edges,
2714 meta_vertices,
2715 input_non_diagonal_edges,
2716 }
2717 }
2718
2719 fn scan_partition_column_into<const K: usize>(
2720 matrix: &SerialEdgeMatrix<K>,
2721 partition: usize,
2722 table: &mut FastHashMap<Kmer<K>, PartitionOtherEnd<K>>,
2723 ) -> PartitionColumnScanOutput<K> {
2724 table.clear();
2725 let column_edges = (0..partition)
2726 .map(|row| matrix.block(row, partition).len())
2727 .sum::<usize>();
2728 table.reserve(column_edges.saturating_sub(table.capacity()));
2729 let mut edges = Vec::with_capacity(column_edges / 2);
2730 let mut meta_vertices = Vec::new();
2731 let mut input_non_diagonal_edges = 0u64;
2732
2733 for row in 0..partition {
2734 for edge in matrix.block(row, partition) {
2735 let (lower, current) = match endpoint_in_partition(matrix, edge, partition) {
2736 Some(pair) => pair,
2737 None => continue,
2738 };
2739 input_non_diagonal_edges += 1;
2740
2741 let incoming = PartitionOtherEnd {
2742 endpoint: lower,
2743 side_at_current: current.side,
2744 weight: edge.weight,
2745 in_same_part: false,
2746 processed: false,
2747 };
2748
2749 absorb_partition_other_end(
2750 current.vertex,
2751 incoming,
2752 partition,
2753 table,
2754 &mut edges,
2755 &mut meta_vertices,
2756 );
2757 }
2758 }
2759
2760 PartitionColumnScanOutput {
2761 edges,
2762 meta_vertices,
2763 input_non_diagonal_edges,
2764 }
2765 }
2766
2767 fn contract_partition_with_reusable_table<const K: usize>(
2768 matrix: &SerialEdgeMatrix<K>,
2769 partition: usize,
2770 threads: usize,
2771 table: &mut FastHashMap<Kmer<K>, PartitionOtherEnd<K>>,
2772 ) -> PartitionContraction<K> {
2773 assert!(partition > 0 && partition < matrix.partition_count());
2774
2775 let (diagonal, scan) = if threads > 1 {
2776 std::thread::scope(|scope| {
2777 let diagonal = scope.spawn(|| Self::compress_diagonal_block(matrix, partition));
2778 let scan = Self::scan_partition_column_into(matrix, partition, table);
2779 (
2780 diagonal
2781 .join()
2782 .expect("diagonal compression worker panicked"),
2783 scan,
2784 )
2785 })
2786 } else {
2787 (
2788 Self::compress_diagonal_block(matrix, partition),
2789 Self::scan_partition_column_into(matrix, partition, table),
2790 )
2791 };
2792
2793 let mut edges = scan.edges;
2794 let mut expansion_edges = Vec::new();
2795 let mut meta_vertices = scan.meta_vertices;
2796 meta_vertices.extend(diagonal.meta_vertices);
2797 let mut phantom_edges = 0u64;
2798
2799 for edge in &diagonal.edges {
2800 let (MatrixEndpoint::Vertex(x), MatrixEndpoint::Vertex(y)) = (edge.first, edge.second)
2801 else {
2802 continue;
2803 };
2804
2805 table.entry(x.vertex).or_insert_with(|| {
2806 phantom_edges += 1;
2807 let phantom = DiscontinuityEndpoint {
2808 vertex: x.vertex,
2809 side: x.side.inverse(),
2810 };
2811 expansion_edges.push(join_other_ends_with_phantom(
2812 MatrixEndpoint::Phi,
2813 MatrixEndpoint::Vertex(phantom),
2814 1,
2815 Some(phantom),
2816 ));
2817 PartitionOtherEnd {
2818 endpoint: MatrixEndpoint::Phi,
2819 side_at_current: x.side.inverse(),
2820 weight: 1,
2821 in_same_part: false,
2822 processed: false,
2823 }
2824 });
2825
2826 table.entry(y.vertex).or_insert_with(|| {
2827 phantom_edges += 1;
2828 let phantom = DiscontinuityEndpoint {
2829 vertex: y.vertex,
2830 side: y.side.inverse(),
2831 };
2832 expansion_edges.push(join_other_ends_with_phantom(
2833 MatrixEndpoint::Phi,
2834 MatrixEndpoint::Vertex(phantom),
2835 1,
2836 Some(phantom),
2837 ));
2838 PartitionOtherEnd {
2839 endpoint: MatrixEndpoint::Phi,
2840 side_at_current: y.side.inverse(),
2841 weight: 1,
2842 in_same_part: false,
2843 processed: false,
2844 }
2845 });
2846
2847 let x_end = table.get(&x.vertex).copied().unwrap();
2848 let y_end = table.get(&y.vertex).copied().unwrap();
2849 if x_end.endpoint.is_phi() && y_end.endpoint.is_phi() {
2850 meta_vertices.push(two_weight_meta_vertex(
2851 x.vertex,
2852 partition,
2853 x.side.inverse(),
2854 x_end.weight,
2855 edge.weight + y_end.weight,
2856 false,
2857 ));
2858 } else {
2859 edges.push(join_other_ends(
2860 x_end.endpoint,
2861 y_end.endpoint,
2862 x_end.weight + edge.weight + y_end.weight,
2863 ));
2864 }
2865
2866 if let Some(end) = table.get_mut(&x.vertex) {
2867 end.processed = true;
2868 }
2869 if let Some(end) = table.get_mut(&y.vertex) {
2870 end.processed = true;
2871 }
2872 }
2873
2874 for (&vertex, &end) in table.iter() {
2875 if end.processed {
2876 continue;
2877 }
2878 phantom_edges += 1;
2879 let phantom = DiscontinuityEndpoint {
2880 vertex,
2881 side: end.side_at_current.inverse(),
2882 };
2883 expansion_edges.push(join_other_ends_with_phantom(
2884 MatrixEndpoint::Phi,
2885 MatrixEndpoint::Vertex(phantom),
2886 1,
2887 Some(phantom),
2888 ));
2889
2890 if end.endpoint.is_phi() {
2891 meta_vertices.push(two_weight_meta_vertex(
2892 vertex,
2893 partition,
2894 end.side_at_current,
2895 end.weight,
2896 1,
2897 false,
2898 ));
2899 } else {
2900 edges.push(join_other_ends(
2901 MatrixEndpoint::Phi,
2902 end.endpoint,
2903 1 + end.weight,
2904 ));
2905 }
2906 }
2907
2908 PartitionContraction {
2909 partition,
2910 compressed_diagonal_edges: diagonal.expansion_edges,
2911 expansion_edges,
2912 stats: PartitionContractionStats {
2913 input_non_diagonal_edges: scan.input_non_diagonal_edges,
2914 compressed_diagonal_edges: diagonal.stats.compressed_edges,
2915 output_edges: edges.len() as u64,
2916 meta_vertices: meta_vertices.len() as u64,
2917 phantom_edges,
2918 isolated_cordless_cycles: diagonal.stats.isolated_cordless_cycles,
2919 },
2920 edges,
2921 meta_vertices,
2922 }
2923 }
2924
2925 #[allow(dead_code)]
2927 fn scan_partition_column_atomic<const K: usize>(
2928 matrix: &SerialEdgeMatrix<K>,
2929 partition: usize,
2930 threads: usize,
2931 atomic_table: &AtomicPartitionTable<K>,
2932 pool: &ThreadPool,
2933 ) -> PartitionColumnScanOutput<K> {
2934 atomic_table.clear(threads, pool);
2935 const EDGES_PER_TASK: usize = 16 * 1024;
2936 let tasks = (0..partition)
2937 .flat_map(|row| matrix.block(row, partition).chunks(EDGES_PER_TASK))
2938 .collect::<Vec<_>>();
2939 let outputs = pool.install(|| {
2940 tasks
2941 .into_par_iter()
2942 .map(|input| {
2943 let mut edges = Vec::with_capacity(input.len() / 2);
2944 let mut meta_vertices = Vec::new();
2945 let mut input_edges = 0u64;
2946 for edge in input {
2947 let Some((lower, current)) = endpoint_in_partition(matrix, edge, partition)
2948 else {
2949 continue;
2950 };
2951 input_edges += 1;
2952 atomic_table.absorb(
2953 current.vertex,
2954 PartitionOtherEnd {
2955 endpoint: lower,
2956 side_at_current: current.side,
2957 weight: edge.weight,
2958 in_same_part: false,
2959 processed: false,
2960 },
2961 partition,
2962 &mut edges,
2963 &mut meta_vertices,
2964 );
2965 }
2966 (edges, meta_vertices, input_edges)
2967 })
2968 .collect::<Vec<_>>()
2969 });
2970 let mut edges = Vec::new();
2971 let mut meta_vertices = Vec::new();
2972 let mut input_non_diagonal_edges = 0;
2973 for (worker_edges, worker_meta, worker_input) in outputs {
2974 edges.extend(worker_edges);
2975 meta_vertices.extend(worker_meta);
2976 input_non_diagonal_edges += worker_input;
2977 }
2978 PartitionColumnScanOutput {
2979 edges,
2980 meta_vertices,
2981 input_non_diagonal_edges,
2982 }
2983 }
2984
2985 #[allow(clippy::too_many_arguments)]
2986 fn contract_blocked_partition_atomic<const K: usize>(
2987 matrix: &mut BlockedEdgeMatrix<K>,
2988 appenders: &ConcurrentBlockedEdgeWriters,
2989 diagonal_matrix: &mut SerialEdgeMatrix<K>,
2990 partition: usize,
2991 threads: usize,
2992 table: &AtomicPartitionTable<K>,
2993 diagonal_ends: &mut FastHashMap<Kmer<K>, DiagonalOtherEnd<K>>,
2994 pool: &ThreadPool,
2995 timings: &mut BlockedContractTimings,
2996 ) -> Result<(PartitionContraction<K>, u64), SerialCollationError> {
2997 let phase = Instant::now();
2998 for row in 0..=partition {
3001 matrix.flush_block(row, partition)?;
3002 let added = appenders.merge_block_into(matrix, row, partition)?;
3003 if added != 0 {
3004 matrix.stats.edges += added as u64;
3005 matrix.stats.phi_edges += u64::from(row == 0) * added as u64;
3006 matrix.stats.diagonal_edges += u64::from(row == partition) * added as u64;
3007 }
3008 }
3009 timings.flush += phase.elapsed();
3010 let phase = Instant::now();
3011 table.clear(threads, pool);
3012 timings.clear += phase.elapsed();
3013
3014 let diagonal_read_started = Instant::now();
3015 let diagonal_edges = matrix.read_flushed_block(partition, partition)?;
3016 diagonal_matrix.blocks[partition][partition] = diagonal_edges;
3017 let diagonal_read_elapsed = diagonal_read_started.elapsed();
3018
3019 let setup_started = Instant::now();
3020 let record_len = discontinuity_edge_record_len::<K>();
3021 let block_ids = (0..partition)
3022 .filter_map(|row| {
3023 let block_id = matrix.block_index(row, partition);
3024 (matrix.blocks[block_id].edges != 0).then_some(block_id)
3025 })
3026 .collect::<Vec<_>>();
3027 let records_per_task = (1024 * 1024 / record_len).max(1);
3028 let bytes_per_task = records_per_task * record_len;
3029 let mut ordered = Vec::new();
3040 for &block_id in &block_ids {
3041 let block = &matrix.blocks[block_id];
3042 let flushed: usize = block.extents.iter().map(|e| e.len as usize).sum();
3043 let expected = block
3044 .edges
3045 .checked_mul(record_len)
3046 .and_then(|bytes| bytes.checked_sub(block.buffer.len()))
3047 .ok_or_else(|| {
3048 SerialCollationError::MalformedCoordBucket(matrix.containers.path_for(block_id))
3049 })?;
3050 if flushed != expected || flushed % record_len != 0 {
3051 return Err(SerialCollationError::MalformedCoordBucket(
3052 matrix.containers.path_for(block_id),
3053 ));
3054 }
3055 let container_index = matrix.containers.container_index(block_id);
3056 let container = matrix.containers.container_for(block_id);
3057 for extent in &block.extents {
3058 let mut done = 0usize;
3061 while done < extent.len as usize {
3062 let len = bytes_per_task.min(extent.len as usize - done);
3063 let offset = extent.offset + done as u64;
3064 ordered.push((
3065 container_index,
3066 offset,
3067 BlockedReadTask::File {
3068 file: &container.file,
3069 path: &container.path,
3070 offset,
3071 len,
3072 },
3073 ));
3074 done += len;
3075 }
3076 }
3077 ordered.extend(
3078 block
3079 .buffer
3080 .chunks(bytes_per_task)
3081 .map(|chunk| (usize::MAX, 0, BlockedReadTask::Memory(chunk))),
3082 );
3083 }
3084 ordered.sort_unstable_by_key(|(container, offset, _)| (*container, *offset));
3090 let tasks = ordered
3091 .into_iter()
3092 .map(|(_, _, task)| task)
3093 .collect::<Vec<_>>();
3094 timings.tasks += setup_started.elapsed();
3095 let join_started = Instant::now();
3096 let ((outputs, scan_elapsed), (diagonal, diagonal_elapsed)) = pool.install(|| {
3097 rayon::join(
3098 || {
3099 let scan_started = Instant::now();
3100 let outputs = tasks
3101 .into_par_iter()
3102 .map_init(
3103 || vec![0u8; bytes_per_task],
3104 |scratch, task| {
3105 let bytes = match task {
3106 BlockedReadTask::File {
3107 file,
3108 path,
3109 offset,
3110 len,
3111 } => {
3112 file.read_exact_at(&mut scratch[..len], offset).map_err(
3113 |source| SerialCollationError::Io {
3114 path: path.to_path_buf(),
3115 source,
3116 },
3117 )?;
3118 &scratch[..len]
3119 }
3120 BlockedReadTask::Memory(bytes) => bytes,
3121 };
3122 let mut prepared = Vec::with_capacity(bytes.len() / record_len / 2);
3123 let mut meta_vertices = Vec::new();
3124 let mut input_edges = 0u64;
3125 for encoded in bytes.chunks_exact(record_len) {
3126 let Some((vertex, incoming)) =
3127 decode_partition_incoming::<K>(encoded)
3128 else {
3129 continue;
3130 };
3131 input_edges += 1;
3132 table.absorb_prepared(
3133 vertex,
3134 incoming,
3135 partition,
3136 matrix.vertex_partitions,
3137 &mut prepared,
3138 &mut meta_vertices,
3139 );
3140 }
3141 let emitted = emit_prepared_edge_batch(&mut prepared, appenders)?;
3142 Ok::<_, SerialCollationError>((meta_vertices, input_edges, emitted))
3143 },
3144 )
3145 .collect::<Result<Vec<_>, SerialCollationError>>();
3146 (outputs, scan_started.elapsed())
3147 },
3148 || {
3149 let diagonal_started = Instant::now();
3150 let diagonal = Self::compress_diagonal_block_with_ends(
3151 diagonal_matrix,
3152 partition,
3153 diagonal_ends,
3154 );
3155 (diagonal, diagonal_started.elapsed())
3156 },
3157 )
3158 });
3159 let join_elapsed = join_started.elapsed();
3160 let outputs = outputs?;
3161 diagonal_matrix.blocks[partition][partition].clear();
3162 timings.join += join_elapsed;
3163 timings.scan += scan_elapsed;
3164 timings.diagonal += diagonal_read_elapsed + diagonal_elapsed;
3165 let gather_started = Instant::now();
3166 let mut meta_vertices =
3169 Vec::with_capacity(outputs.iter().map(|(meta, _, _)| meta.len()).sum());
3170 let mut input_non_diagonal_edges = 0;
3171 let mut pre_reinserted_edges = 0u64;
3172 for (local_meta, local_input, local_emitted) in outputs {
3173 meta_vertices.extend(local_meta);
3174 input_non_diagonal_edges += local_input;
3175 pre_reinserted_edges += local_emitted;
3176 }
3177 let scan = PartitionColumnScanOutput {
3178 edges: Vec::new(),
3179 meta_vertices,
3180 input_non_diagonal_edges,
3181 };
3182 timings.gather += gather_started.elapsed();
3183 let phase = Instant::now();
3184 let contracted =
3185 Self::finish_atomic_partition(partition, threads, table, pool, diagonal, scan);
3186 timings.finish += phase.elapsed();
3187 let mut contracted = contracted;
3188 contracted.stats.output_edges += pre_reinserted_edges;
3189 Ok((contracted, pre_reinserted_edges))
3190 }
3191
3192 #[allow(dead_code)]
3194 fn contract_partition_with_owned_tables<const K: usize>(
3195 matrix: &SerialEdgeMatrix<K>,
3196 partition: usize,
3197 threads: usize,
3198 tables: &mut OwnedPartitionTables<K>,
3199 ) -> PartitionContraction<K> {
3200 let (diagonal, scan) = if threads > 1 {
3201 std::thread::scope(|scope| {
3202 let diagonal = scope.spawn(|| Self::compress_diagonal_block(matrix, partition));
3203 let scan = Self::scan_partition_column_owned(matrix, partition, threads, tables);
3204 (
3205 diagonal
3206 .join()
3207 .expect("diagonal compression worker panicked"),
3208 scan,
3209 )
3210 })
3211 } else {
3212 let diagonal = Self::compress_diagonal_block(matrix, partition);
3213 let scan = Self::scan_partition_column_owned(matrix, partition, 1, tables);
3214 (diagonal, scan)
3215 };
3216
3217 let mut edges = scan.edges;
3218 let mut expansion_edges = Vec::new();
3219 let mut meta_vertices = scan.meta_vertices;
3220 meta_vertices.extend(diagonal.meta_vertices);
3221 let mut phantom_edges = 0u64;
3222
3223 for edge in &diagonal.edges {
3224 let (MatrixEndpoint::Vertex(x), MatrixEndpoint::Vertex(y)) = (edge.first, edge.second)
3225 else {
3226 continue;
3227 };
3228 for endpoint in [x, y] {
3229 if !tables.contains_key(&endpoint.vertex) {
3230 phantom_edges += 1;
3231 let phantom = DiscontinuityEndpoint {
3232 vertex: endpoint.vertex,
3233 side: endpoint.side.inverse(),
3234 };
3235 expansion_edges.push(join_other_ends_with_phantom(
3236 MatrixEndpoint::Phi,
3237 MatrixEndpoint::Vertex(phantom),
3238 1,
3239 Some(phantom),
3240 ));
3241 tables.insert(
3242 endpoint.vertex,
3243 PartitionOtherEnd {
3244 endpoint: MatrixEndpoint::Phi,
3245 side_at_current: endpoint.side.inverse(),
3246 weight: 1,
3247 in_same_part: false,
3248 processed: false,
3249 },
3250 );
3251 }
3252 }
3253
3254 let x_end = *tables.get(&x.vertex).expect("inserted missing x endpoint");
3255 let y_end = *tables.get(&y.vertex).expect("inserted missing y endpoint");
3256 if x_end.endpoint.is_phi() && y_end.endpoint.is_phi() {
3257 meta_vertices.push(two_weight_meta_vertex(
3258 x.vertex,
3259 partition,
3260 x.side.inverse(),
3261 x_end.weight,
3262 edge.weight + y_end.weight,
3263 false,
3264 ));
3265 } else {
3266 edges.push(join_other_ends(
3267 x_end.endpoint,
3268 y_end.endpoint,
3269 x_end.weight + edge.weight + y_end.weight,
3270 ));
3271 }
3272 tables
3273 .get_mut(&x.vertex)
3274 .expect("x endpoint exists")
3275 .processed = true;
3276 tables
3277 .get_mut(&y.vertex)
3278 .expect("y endpoint exists")
3279 .processed = true;
3280 }
3281
3282 for map in &tables.maps {
3283 for (&vertex, &end) in map {
3284 if end.processed {
3285 continue;
3286 }
3287 phantom_edges += 1;
3288 let phantom = DiscontinuityEndpoint {
3289 vertex,
3290 side: end.side_at_current.inverse(),
3291 };
3292 expansion_edges.push(join_other_ends_with_phantom(
3293 MatrixEndpoint::Phi,
3294 MatrixEndpoint::Vertex(phantom),
3295 1,
3296 Some(phantom),
3297 ));
3298 if end.endpoint.is_phi() {
3299 meta_vertices.push(two_weight_meta_vertex(
3300 vertex,
3301 partition,
3302 end.side_at_current,
3303 end.weight,
3304 1,
3305 false,
3306 ));
3307 } else {
3308 edges.push(join_other_ends(
3309 MatrixEndpoint::Phi,
3310 end.endpoint,
3311 1 + end.weight,
3312 ));
3313 }
3314 }
3315 }
3316
3317 PartitionContraction {
3318 partition,
3319 compressed_diagonal_edges: diagonal.expansion_edges,
3320 expansion_edges,
3321 stats: PartitionContractionStats {
3322 input_non_diagonal_edges: scan.input_non_diagonal_edges,
3323 compressed_diagonal_edges: diagonal.stats.compressed_edges,
3324 output_edges: edges.len() as u64,
3325 meta_vertices: meta_vertices.len() as u64,
3326 phantom_edges,
3327 isolated_cordless_cycles: diagonal.stats.isolated_cordless_cycles,
3328 },
3329 edges,
3330 meta_vertices,
3331 }
3332 }
3333
3334 #[allow(dead_code)]
3336 fn contract_partition_with_atomic_table<const K: usize>(
3337 matrix: &SerialEdgeMatrix<K>,
3338 partition: usize,
3339 threads: usize,
3340 table: &AtomicPartitionTable<K>,
3341 pool: &ThreadPool,
3342 ) -> PartitionContraction<K> {
3343 let (diagonal, scan) = std::thread::scope(|scope| {
3344 let diagonal = scope.spawn(|| Self::compress_diagonal_block(matrix, partition));
3345 let scan = Self::scan_partition_column_atomic(matrix, partition, threads, table, pool);
3346 (
3347 diagonal
3348 .join()
3349 .expect("diagonal compression worker panicked"),
3350 scan,
3351 )
3352 });
3353 Self::finish_atomic_partition(partition, threads, table, pool, diagonal, scan)
3354 }
3355
3356 fn finish_atomic_partition<const K: usize>(
3357 partition: usize,
3358 threads: usize,
3359 table: &AtomicPartitionTable<K>,
3360 pool: &ThreadPool,
3361 diagonal: DiagonalCompression<K>,
3362 scan: PartitionColumnScanOutput<K>,
3363 ) -> PartitionContraction<K> {
3364 let mut edges = scan.edges;
3365 let mut expansion_edges = Vec::new();
3366 let mut meta_vertices = scan.meta_vertices;
3367 meta_vertices.extend(diagonal.meta_vertices);
3368 let mut phantom_edges = 0u64;
3369
3370 let diagonal_chunk = diagonal.edges.len().div_ceil(threads.max(1)).max(1);
3371 let diagonal_outputs = pool.install(|| {
3372 diagonal
3373 .edges
3374 .par_chunks(diagonal_chunk)
3375 .map(|diagonal_edges| {
3376 let mut local_edges = Vec::new();
3377 let mut local_expansion = Vec::new();
3378 let mut local_meta = Vec::new();
3379 let mut local_phantoms = 0u64;
3380 for edge in diagonal_edges {
3381 let (MatrixEndpoint::Vertex(x), MatrixEndpoint::Vertex(y)) =
3382 (edge.first, edge.second)
3383 else {
3384 continue;
3385 };
3386 let x_stored = table.get(x.vertex);
3387 let y_stored = table.get(y.vertex);
3388 let endpoint = |vertex: DiscontinuityEndpoint<K>| PartitionOtherEnd {
3389 endpoint: MatrixEndpoint::Phi,
3390 side_at_current: vertex.side.inverse(),
3391 weight: 1,
3392 in_same_part: false,
3393 processed: true,
3394 };
3395 let x_end = x_stored.unwrap_or_else(|| endpoint(x));
3396 let y_end = y_stored.unwrap_or_else(|| endpoint(y));
3397 for (stored, endpoint) in [(x_stored, x), (y_stored, y)] {
3398 if stored.is_none() {
3399 local_phantoms += 1;
3400 let phantom = DiscontinuityEndpoint {
3401 vertex: endpoint.vertex,
3402 side: endpoint.side.inverse(),
3403 };
3404 local_expansion.push(join_other_ends_with_phantom(
3405 MatrixEndpoint::Phi,
3406 MatrixEndpoint::Vertex(phantom),
3407 1,
3408 Some(phantom),
3409 ));
3410 }
3411 }
3412 if x_end.endpoint.is_phi() && y_end.endpoint.is_phi() {
3413 local_meta.push(two_weight_meta_vertex(
3414 x.vertex,
3415 partition,
3416 x.side.inverse(),
3417 x_end.weight,
3418 edge.weight + y_end.weight,
3419 false,
3420 ));
3421 } else {
3422 local_edges.push(join_other_ends(
3423 x_end.endpoint,
3424 y_end.endpoint,
3425 x_end.weight + edge.weight + y_end.weight,
3426 ));
3427 }
3428 if x_stored.is_some() {
3429 table.mark_processed(x.vertex);
3430 }
3431 if y_stored.is_some() {
3432 table.mark_processed(y.vertex);
3433 }
3434 }
3435 (local_edges, local_expansion, local_meta, local_phantoms)
3436 })
3437 .collect::<Vec<_>>()
3438 });
3439 for (local_edges, local_expansion, local_meta, local_phantoms) in diagonal_outputs {
3440 edges.extend(local_edges);
3441 expansion_edges.extend(local_expansion);
3442 meta_vertices.extend(local_meta);
3443 phantom_edges += local_phantoms;
3444 }
3445
3446 let scan_chunk = table.slots.len().div_ceil(threads.max(1)).max(1);
3447 let phantom_outputs = pool.install(|| {
3448 table
3449 .slots
3450 .par_chunks(scan_chunk)
3451 .map(|slots| {
3452 let mut local_edges = Vec::new();
3453 let mut local_expansion = Vec::new();
3454 let mut local_meta = Vec::new();
3455 let mut local_phantoms = 0u64;
3456 for slot in slots {
3457 let key = slot.key.load(Ordering::Relaxed);
3458 if key == AtomicPartitionTable::<K>::EMPTY {
3459 continue;
3460 }
3461 let end = unsafe { (*slot.value.get()).assume_init() }.unpack::<K>();
3462 if end.processed {
3463 continue;
3464 }
3465 let vertex = Kmer::from_bits(key as u128);
3466 local_phantoms += 1;
3467 let phantom = DiscontinuityEndpoint {
3468 vertex,
3469 side: end.side_at_current.inverse(),
3470 };
3471 local_expansion.push(join_other_ends_with_phantom(
3472 MatrixEndpoint::Phi,
3473 MatrixEndpoint::Vertex(phantom),
3474 1,
3475 Some(phantom),
3476 ));
3477 if end.endpoint.is_phi() {
3478 local_meta.push(two_weight_meta_vertex(
3479 vertex,
3480 partition,
3481 end.side_at_current,
3482 end.weight,
3483 1,
3484 false,
3485 ));
3486 } else {
3487 local_edges.push(join_other_ends(
3488 MatrixEndpoint::Phi,
3489 end.endpoint,
3490 1 + end.weight,
3491 ));
3492 }
3493 }
3494 (local_edges, local_expansion, local_meta, local_phantoms)
3495 })
3496 .collect::<Vec<_>>()
3497 });
3498 for (local_edges, local_expansion, local_meta, local_phantoms) in phantom_outputs {
3499 edges.extend(local_edges);
3500 expansion_edges.extend(local_expansion);
3501 meta_vertices.extend(local_meta);
3502 phantom_edges += local_phantoms;
3503 }
3504 PartitionContraction {
3505 partition,
3506 compressed_diagonal_edges: diagonal.expansion_edges,
3507 expansion_edges,
3508 stats: PartitionContractionStats {
3509 input_non_diagonal_edges: scan.input_non_diagonal_edges,
3510 compressed_diagonal_edges: diagonal.stats.compressed_edges,
3511 output_edges: edges.len() as u64,
3512 meta_vertices: meta_vertices.len() as u64,
3513 phantom_edges,
3514 isolated_cordless_cycles: diagonal.stats.isolated_cordless_cycles,
3515 },
3516 edges,
3517 meta_vertices,
3518 }
3519 }
3520
3521 #[allow(dead_code)]
3523 fn contract_partition_with_flat_table<const K: usize>(
3524 matrix: &SerialEdgeMatrix<K>,
3525 partition: usize,
3526 table: &mut FlatPartitionTable<K>,
3527 ) -> PartitionContraction<K> {
3528 let (diagonal, scan) = std::thread::scope(|scope| {
3529 let diagonal = scope.spawn(|| Self::compress_diagonal_block(matrix, partition));
3530 let scan = Self::scan_partition_column_flat(matrix, partition, table);
3531 (
3532 diagonal
3533 .join()
3534 .expect("diagonal compression worker panicked"),
3535 scan,
3536 )
3537 });
3538 let mut edges = scan.edges;
3539 let mut expansion_edges = Vec::new();
3540 let mut meta_vertices = scan.meta_vertices;
3541 meta_vertices.extend(diagonal.meta_vertices);
3542 let mut phantom_edges = 0u64;
3543 for edge in &diagonal.edges {
3544 let (MatrixEndpoint::Vertex(x), MatrixEndpoint::Vertex(y)) = (edge.first, edge.second)
3545 else {
3546 continue;
3547 };
3548 for endpoint in [x, y] {
3549 if table.get(endpoint.vertex).is_none() {
3550 phantom_edges += 1;
3551 let phantom = DiscontinuityEndpoint {
3552 vertex: endpoint.vertex,
3553 side: endpoint.side.inverse(),
3554 };
3555 expansion_edges.push(join_other_ends_with_phantom(
3556 MatrixEndpoint::Phi,
3557 MatrixEndpoint::Vertex(phantom),
3558 1,
3559 Some(phantom),
3560 ));
3561 table.insert(
3562 endpoint.vertex,
3563 PartitionOtherEnd {
3564 endpoint: MatrixEndpoint::Phi,
3565 side_at_current: endpoint.side.inverse(),
3566 weight: 1,
3567 in_same_part: false,
3568 processed: false,
3569 },
3570 );
3571 }
3572 }
3573 let x_end = table.get(x.vertex).expect("inserted x endpoint");
3574 let y_end = table.get(y.vertex).expect("inserted y endpoint");
3575 if x_end.endpoint.is_phi() && y_end.endpoint.is_phi() {
3576 meta_vertices.push(two_weight_meta_vertex(
3577 x.vertex,
3578 partition,
3579 x.side.inverse(),
3580 x_end.weight,
3581 edge.weight + y_end.weight,
3582 false,
3583 ));
3584 } else {
3585 edges.push(join_other_ends(
3586 x_end.endpoint,
3587 y_end.endpoint,
3588 x_end.weight + edge.weight + y_end.weight,
3589 ));
3590 }
3591 table.mark_processed(x.vertex);
3592 table.mark_processed(y.vertex);
3593 }
3594 for &idx in &table.occupied {
3595 let end = unsafe { table.values[idx].assume_init() };
3596 if end.processed {
3597 continue;
3598 }
3599 let vertex = Kmer::from_bits(table.keys[idx] as u128);
3600 phantom_edges += 1;
3601 let phantom = DiscontinuityEndpoint {
3602 vertex,
3603 side: end.side_at_current.inverse(),
3604 };
3605 expansion_edges.push(join_other_ends_with_phantom(
3606 MatrixEndpoint::Phi,
3607 MatrixEndpoint::Vertex(phantom),
3608 1,
3609 Some(phantom),
3610 ));
3611 if end.endpoint.is_phi() {
3612 meta_vertices.push(two_weight_meta_vertex(
3613 vertex,
3614 partition,
3615 end.side_at_current,
3616 end.weight,
3617 1,
3618 false,
3619 ));
3620 } else {
3621 edges.push(join_other_ends(
3622 MatrixEndpoint::Phi,
3623 end.endpoint,
3624 1 + end.weight,
3625 ));
3626 }
3627 }
3628 PartitionContraction {
3629 partition,
3630 compressed_diagonal_edges: diagonal.expansion_edges,
3631 expansion_edges,
3632 stats: PartitionContractionStats {
3633 input_non_diagonal_edges: scan.input_non_diagonal_edges,
3634 compressed_diagonal_edges: diagonal.stats.compressed_edges,
3635 output_edges: edges.len() as u64,
3636 meta_vertices: meta_vertices.len() as u64,
3637 phantom_edges,
3638 isolated_cordless_cycles: diagonal.stats.isolated_cordless_cycles,
3639 },
3640 edges,
3641 meta_vertices,
3642 }
3643 }
3644
3645 fn scan_partition_column_flat<const K: usize>(
3646 matrix: &SerialEdgeMatrix<K>,
3647 partition: usize,
3648 table: &mut FlatPartitionTable<K>,
3649 ) -> PartitionColumnScanOutput<K> {
3650 table.clear();
3651 let column_edges = (0..partition)
3652 .map(|row| matrix.block(row, partition).len())
3653 .sum::<usize>();
3654 let mut edges = Vec::with_capacity(column_edges / 2);
3655 let mut meta_vertices = Vec::new();
3656 let mut input_non_diagonal_edges = 0;
3657 for row in 0..partition {
3658 for edge in matrix.block(row, partition) {
3659 let Some((lower, current)) = endpoint_in_partition(matrix, edge, partition) else {
3660 continue;
3661 };
3662 input_non_diagonal_edges += 1;
3663 table.absorb(
3664 current.vertex,
3665 PartitionOtherEnd {
3666 endpoint: lower,
3667 side_at_current: current.side,
3668 weight: edge.weight,
3669 in_same_part: false,
3670 processed: false,
3671 },
3672 partition,
3673 &mut edges,
3674 &mut meta_vertices,
3675 );
3676 }
3677 }
3678 PartitionColumnScanOutput {
3679 edges,
3680 meta_vertices,
3681 input_non_diagonal_edges,
3682 }
3683 }
3684
3685 fn scan_partition_column_owned<const K: usize>(
3686 matrix: &SerialEdgeMatrix<K>,
3687 partition: usize,
3688 threads: usize,
3689 tables: &mut OwnedPartitionTables<K>,
3690 ) -> PartitionColumnScanOutput<K> {
3691 tables.clear();
3692 let routing_workers = threads.max(1).min(partition);
3693 let rows_per_worker = partition.div_ceil(routing_workers);
3694 let owner_count = tables.maps.len();
3695 let owner_mask = tables.mask;
3696 let routed = std::thread::scope(|scope| {
3697 let mut handles = Vec::with_capacity(routing_workers);
3698 for row_start in (0..partition).step_by(rows_per_worker) {
3699 let row_end = (row_start + rows_per_worker).min(partition);
3700 handles.push(scope.spawn(move || {
3701 let mut owners = (0..owner_count)
3702 .map(|_| Vec::<PartitionColumnIncoming<K>>::new())
3703 .collect::<Vec<_>>();
3704 let mut input_edges = 0u64;
3705 for row in row_start..row_end {
3706 for edge in matrix.block(row, partition) {
3707 let Some((lower, current)) =
3708 endpoint_in_partition(matrix, edge, partition)
3709 else {
3710 continue;
3711 };
3712 input_edges += 1;
3713 let owner = partition_column_vertex_shard(current.vertex, owner_mask);
3714 owners[owner].push(PartitionColumnIncoming {
3715 vertex: current.vertex,
3716 end: PartitionOtherEnd {
3717 endpoint: lower,
3718 side_at_current: current.side,
3719 weight: edge.weight,
3720 in_same_part: false,
3721 processed: false,
3722 },
3723 });
3724 }
3725 }
3726 (owners, input_edges)
3727 }));
3728 }
3729 handles
3730 .into_iter()
3731 .map(|handle| handle.join().expect("partition routing worker panicked"))
3732 .collect::<Vec<_>>()
3733 });
3734
3735 let input_non_diagonal_edges = routed.iter().map(|(_, count)| *count).sum();
3736 let mut owner_inputs = (0..owner_count)
3737 .map(|_| Vec::<PartitionColumnIncoming<K>>::new())
3738 .collect::<Vec<_>>();
3739 for (mut worker_owners, _) in routed {
3740 for (owner, incoming) in worker_owners.iter_mut().enumerate() {
3741 owner_inputs[owner].append(incoming);
3742 }
3743 }
3744
3745 let outputs = std::thread::scope(|scope| {
3746 let mut handles = Vec::with_capacity(owner_count);
3747 for (map, incoming) in tables.maps.iter_mut().zip(owner_inputs) {
3748 handles.push(scope.spawn(move || {
3749 map.reserve(incoming.len().saturating_sub(map.capacity()));
3750 let mut edges = Vec::with_capacity(incoming.len() / 2);
3751 let mut meta_vertices = Vec::new();
3752 for incoming in incoming {
3753 absorb_partition_other_end(
3754 incoming.vertex,
3755 incoming.end,
3756 partition,
3757 map,
3758 &mut edges,
3759 &mut meta_vertices,
3760 );
3761 }
3762 (edges, meta_vertices)
3763 }));
3764 }
3765 handles
3766 .into_iter()
3767 .map(|handle| handle.join().expect("partition owner worker panicked"))
3768 .collect::<Vec<_>>()
3769 });
3770 let mut edges = Vec::new();
3771 let mut meta_vertices = Vec::new();
3772 for (worker_edges, worker_meta) in outputs {
3773 edges.extend(worker_edges);
3774 meta_vertices.extend(worker_meta);
3775 }
3776 PartitionColumnScanOutput {
3777 edges,
3778 meta_vertices,
3779 input_non_diagonal_edges,
3780 }
3781 }
3782
3783 pub fn contract_all_partitions<const K: usize>(
3784 matrix: &SerialEdgeMatrix<K>,
3785 ) -> FullSerialDiscontinuityContraction<K> {
3786 Self::contract_all_partitions_with_threads(matrix, 1)
3787 }
3788
3789 pub fn contract_all_partitions_with_threads<const K: usize>(
3790 matrix: &SerialEdgeMatrix<K>,
3791 threads: usize,
3792 ) -> FullSerialDiscontinuityContraction<K> {
3793 Self::contract_all_partitions_owned(matrix.clone(), threads)
3794 }
3795
3796 fn contract_all_partitions_owned<const K: usize>(
3797 working: SerialEdgeMatrix<K>,
3798 threads: usize,
3799 ) -> FullSerialDiscontinuityContraction<K> {
3800 Self::contract_all_partitions_owned_impl::<K, true>(working, threads)
3801 }
3802
3803 fn contract_blocked_external<const K: usize>(
3804 mut working: BlockedEdgeMatrix<K>,
3805 threads: usize,
3806 ) -> Result<ExternalBlockedContraction<K>, SerialCollationError> {
3807 let setup_started = Instant::now();
3808 let partition_count = working.partition_count();
3809 let total_partitions = partition_count - 1;
3810 let mut compressed_diagonal_edges =
3811 (0..partition_count).map(|_| Vec::new()).collect::<Vec<_>>();
3812 let meta_vertex_dir = working.dir.join("contracted-meta-path-info");
3813 if meta_vertex_dir.exists() {
3814 fs::remove_dir_all(&meta_vertex_dir).map_err(|source| SerialCollationError::Io {
3815 path: meta_vertex_dir.clone(),
3816 source,
3817 })?;
3818 }
3819 fs::create_dir_all(&meta_vertex_dir).map_err(|source| SerialCollationError::Io {
3820 path: meta_vertex_dir.clone(),
3821 source,
3822 })?;
3823 let mut meta_vertex_count = 0u64;
3824 let mut meta_vertices_per_partition = vec![0usize; partition_count];
3825 let mut reusable_table =
3826 FastHashMap::<Kmer<K>, PartitionOtherEnd<K>>::with_hasher(FastBuildHasher::default());
3827 let mut diagonal_ends =
3828 FastHashMap::<Kmer<K>, DiagonalOtherEnd<K>>::with_hasher(FastBuildHasher::default());
3829 let mut column = SerialEdgeMatrix::new(working.vertex_partitions)
3830 .map_err(|_| SerialCollationError::MalformedCoordBucket(working.dir.clone()))?;
3831 let max_partition_vertices = (1..partition_count)
3832 .map(|col| {
3833 let column_edges = (0..=col)
3834 .map(|row| working.blocks[working.block_index(row, col)].edges)
3835 .sum::<usize>();
3836 column_edges + working.blocks[working.block_index(col, col)].edges
3837 })
3838 .max()
3839 .unwrap_or(1);
3840 let atomic_table =
3841 (K <= 31).then(|| AtomicPartitionTable::<K>::with_max_entries(max_partition_vertices));
3842 let contraction_pool = ThreadPoolBuilder::new()
3843 .num_threads(threads.max(1))
3844 .build()
3845 .map_err(|_| SerialCollationError::WorkerPanic)?;
3846 let appenders = ConcurrentBlockedEdgeWriters::new(&working);
3847 let setup_elapsed = setup_started.elapsed();
3848 if let Some(table) = &atomic_table {
3849 eprintln!(
3850 "cuttlefish: contraction table setup {:.3}s; max entries {}, capacity {}, slot {} byte(s)",
3851 setup_elapsed.as_secs_f64(),
3852 max_partition_vertices,
3853 table.slots.len(),
3854 std::mem::size_of::<AtomicPartitionSlot>(),
3855 );
3856 }
3857 let mut stats = FullSerialContractionStats {
3858 input_edges: working.stats.edges,
3859 ..FullSerialContractionStats::default()
3860 };
3861 let started = Instant::now();
3862 let mut block_load_elapsed = Duration::default();
3863 let mut partition_contract_elapsed = Duration::default();
3864 let mut reinsert_elapsed = Duration::default();
3865 let mut direct_timings = BlockedContractTimings::default();
3866 let mut meta_write_elapsed = Duration::default();
3867 eprintln!(
3868 "cuttlefish: contracting {} blocked discontinuity partition(s) with {} worker(s)",
3869 total_partitions, threads
3870 );
3871
3872 for (completed, partition) in (1..partition_count).rev().enumerate() {
3873 let phase_started = Instant::now();
3874 let (mut contracted, pre_reinserted_edges) = if let Some(table) = &atomic_table {
3875 Self::contract_blocked_partition_atomic(
3876 &mut working,
3877 &appenders,
3878 &mut column,
3879 partition,
3880 threads,
3881 table,
3882 &mut diagonal_ends,
3883 &contraction_pool,
3884 &mut direct_timings,
3885 )?
3886 } else {
3887 let load_started = Instant::now();
3888 working.load_column_into(partition, threads, &mut column)?;
3889 block_load_elapsed += load_started.elapsed();
3890 (
3891 Self::contract_partition_with_reusable_table(
3892 &column,
3893 partition,
3894 threads,
3895 &mut reusable_table,
3896 ),
3897 0,
3898 )
3899 };
3900 partition_contract_elapsed += phase_started.elapsed();
3901 stats.partition_output_edges += contracted.stats.output_edges;
3902 stats.phantom_edges += contracted.stats.phantom_edges;
3903 stats.isolated_cordless_cycles += contracted.stats.isolated_cordless_cycles;
3904 stats.reinserted_edges += pre_reinserted_edges;
3905 compressed_diagonal_edges[partition] =
3906 std::mem::take(&mut contracted.compressed_diagonal_edges);
3907
3908 let phase_started = Instant::now();
3909 let vertex_partitions = working.vertex_partitions;
3910 let reinserted = contraction_pool.install(|| {
3911 let (expansion_result, reinsert_result) = rayon::join(
3912 || {
3913 emit_contracted_edge_chunks(
3914 contracted.expansion_edges,
3915 vertex_partitions,
3916 None,
3917 &appenders,
3918 )
3919 },
3920 || {
3921 emit_contracted_edge_chunks(
3922 contracted.edges,
3923 vertex_partitions,
3924 Some(partition),
3925 &appenders,
3926 )
3927 },
3928 );
3929 expansion_result?;
3930 reinsert_result
3931 })?;
3932 stats.reinserted_edges += reinserted;
3933 reinsert_elapsed += phase_started.elapsed();
3934 let meta_write_started = Instant::now();
3935 write_meta_vertex_bucket_parallel(
3936 &meta_vertex_dir,
3937 partition,
3938 &contracted.meta_vertices,
3939 &contraction_pool,
3940 )?;
3941 meta_write_elapsed += meta_write_started.elapsed();
3942 meta_vertex_count += contracted.meta_vertices.len() as u64;
3943 meta_vertices_per_partition[partition] += contracted.meta_vertices.len();
3944 for row in 0..=partition {
3945 column.blocks[row][partition].clear();
3946 }
3947 column.stats = SerialEdgeMatrixStats::default();
3948 report_discontinuity_contraction_progress(completed + 1, total_partitions, started);
3949 }
3950 let flush_started = Instant::now();
3951 for row in 0..partition_count {
3952 for col in row..partition_count {
3953 let added = appenders.merge_block_into(&mut working, row, col)?;
3954 if added != 0 {
3955 working.stats.edges += added as u64;
3956 working.stats.phi_edges += u64::from(row == 0) * added as u64;
3957 working.stats.diagonal_edges += u64::from(row == col) * added as u64;
3958 }
3959 }
3960 }
3961 working.flush_all_with_threads(threads)?;
3962 let block_flush_elapsed = flush_started.elapsed();
3963 eprintln!(
3964 "cuttlefish: blocked contraction detail: setup {:.3}s, block load/decode {:.3}s, partition contraction {:.3}s, reinsertion {:.3}s, meta write {:.3}s, block flush {:.3}s",
3965 setup_elapsed.as_secs_f64(),
3966 block_load_elapsed.as_secs_f64(),
3967 partition_contract_elapsed.as_secs_f64(),
3968 reinsert_elapsed.as_secs_f64(),
3969 meta_write_elapsed.as_secs_f64(),
3970 block_flush_elapsed.as_secs_f64(),
3971 );
3972 if atomic_table.is_some() {
3973 eprintln!(
3974 "cuttlefish: fused contraction phases: column flush {:.3}s, table clear {:.3}s, diagonal {:.3}s, raw read {:.3}s, table scan {:.3}s, scan/diagonal join wall {:.3}s, task setup {:.3}s, gather {:.3}s, finalize {:.3}s",
3975 direct_timings.flush.as_secs_f64(),
3976 direct_timings.clear.as_secs_f64(),
3977 direct_timings.diagonal.as_secs_f64(),
3978 direct_timings.read.as_secs_f64(),
3979 direct_timings.scan.as_secs_f64(),
3980 direct_timings.join.as_secs_f64(),
3981 direct_timings.tasks.as_secs_f64(),
3982 direct_timings.gather.as_secs_f64(),
3983 direct_timings.finish.as_secs_f64(),
3984 );
3985 }
3986 stats.partitions = total_partitions as u64;
3987 stats.meta_vertices = meta_vertex_count;
3988 Ok(ExternalBlockedContraction {
3989 vertex_partitions: working.vertex_partitions,
3990 expansion_matrix: working,
3991 compressed_diagonal_edges,
3992 meta_vertex_dir,
3993 meta_vertex_count,
3994 meta_vertices_per_partition,
3995 stats,
3996 })
3997 }
3998
3999 fn contract_all_partitions_owned_impl<const K: usize, const SORT_EXPANSION_EDGES: bool>(
4000 mut working: SerialEdgeMatrix<K>,
4001 threads: usize,
4002 ) -> FullSerialDiscontinuityContraction<K> {
4003 let mut partitions = Vec::new();
4004 let mut final_edges = Vec::new();
4005 let mut meta_vertices = Vec::new();
4006 let mut compressed_diagonal_edges = (0..working.partition_count())
4007 .map(|_| Vec::new())
4008 .collect::<Vec<_>>();
4009 let mut reusable_table =
4010 FastHashMap::<Kmer<K>, PartitionOtherEnd<K>>::with_hasher(FastBuildHasher::default());
4011 let mut stats = FullSerialContractionStats {
4012 input_edges: working.stats().edges,
4013 ..FullSerialContractionStats::default()
4014 };
4015 let total_partitions = working.partition_count().saturating_sub(1);
4016 let started = Instant::now();
4017 if total_partitions >= 16 {
4018 eprintln!(
4019 "cuttlefish: contracting {} discontinuity partition(s) with {} worker(s)",
4020 total_partitions, threads
4021 );
4022 }
4023
4024 for (completed, partition) in (1..working.partition_count()).rev().enumerate() {
4025 let mut contracted = if SORT_EXPANSION_EDGES {
4026 Self::contract_partition_with_threads(&working, partition, threads)
4027 } else {
4028 Self::contract_partition_with_reusable_table(
4029 &working,
4030 partition,
4031 threads,
4032 &mut reusable_table,
4033 )
4034 };
4035 stats.partition_output_edges += contracted.stats.output_edges;
4036 stats.phantom_edges += contracted.stats.phantom_edges;
4037 stats.isolated_cordless_cycles += contracted.stats.isolated_cordless_cycles;
4038 if !SORT_EXPANSION_EDGES && contracted.partition < compressed_diagonal_edges.len() {
4039 compressed_diagonal_edges[contracted.partition] =
4040 std::mem::take(&mut contracted.compressed_diagonal_edges);
4041 }
4042 for edge in &contracted.expansion_edges {
4043 working.add_edge_with_orientation_and_phantom(
4044 edge.first,
4045 edge.second,
4046 edge.weight,
4047 edge.unitig_index,
4048 edge.unitig_exit_side,
4049 edge.phantom_unitig,
4050 );
4051 }
4052
4053 for edge in &contracted.edges {
4054 if max_endpoint_partition(&working, edge.first, edge.second) < partition {
4055 working.add_edge_with_orientation_and_phantom(
4056 edge.first,
4057 edge.second,
4058 edge.weight,
4059 edge.unitig_index,
4060 edge.unitig_exit_side,
4061 edge.phantom_unitig,
4062 );
4063 stats.reinserted_edges += 1;
4064 } else if SORT_EXPANSION_EDGES {
4065 final_edges.push(edge.clone());
4066 }
4067 }
4068
4069 meta_vertices.extend(contracted.meta_vertices.iter().cloned());
4070 if SORT_EXPANSION_EDGES {
4071 if contracted.partition < compressed_diagonal_edges.len() {
4072 compressed_diagonal_edges[contracted.partition] =
4073 contracted.compressed_diagonal_edges.clone();
4074 }
4075 partitions.push(contracted);
4076 }
4077 report_discontinuity_contraction_progress(completed + 1, total_partitions, started);
4078 }
4079
4080 if SORT_EXPANSION_EDGES {
4081 final_edges.sort_by_key(|edge| {
4082 (
4083 endpoint_sort_key(edge.first),
4084 endpoint_sort_key(edge.second),
4085 edge.weight,
4086 )
4087 });
4088 meta_vertices.sort_by_key(|meta| {
4089 (
4090 meta.partition,
4091 meta.vertex.as_u128(),
4092 meta.entry_side as u8,
4093 meta.weight,
4094 meta.is_cycle,
4095 )
4096 });
4097 }
4098
4099 stats.partitions = total_partitions as u64;
4100 stats.final_edges = final_edges.len() as u64;
4101 stats.meta_vertices = meta_vertices.len() as u64;
4102
4103 let expansion_edges = if SORT_EXPANSION_EDGES {
4104 let mut expansion_edges = working.edges().cloned().collect::<Vec<_>>();
4105 expansion_edges.sort_by_key(|edge| {
4106 (
4107 endpoint_sort_key(edge.first),
4108 endpoint_sort_key(edge.second),
4109 edge.weight,
4110 edge.unitig_index,
4111 )
4112 });
4113 expansion_edges
4114 } else {
4115 Vec::new()
4116 };
4117
4118 FullSerialDiscontinuityContraction {
4119 vertex_partitions: working.vertex_partitions(),
4120 final_edges,
4121 expansion_edges,
4122 expansion_matrix: working,
4123 compressed_diagonal_edges,
4124 meta_vertices,
4125 partitions,
4126 stats,
4127 }
4128 }
4129
4130 pub fn contract<const K: usize>(matrix: &SerialEdgeMatrix<K>) -> SerialContraction {
4131 let mut endpoint_ids = FastHashMap::<EndpointKey<K>, usize>::with_capacity_and_hasher(
4132 matrix.stats.edges as usize * 2,
4133 FastBuildHasher::default(),
4134 );
4135 let mut adjacency = Vec::<Vec<(usize, u64, bool)>>::new();
4136
4137 for edge in matrix.edges() {
4138 let first = endpoint_id(&mut endpoint_ids, &mut adjacency, edge.first);
4139 let second = endpoint_id(&mut endpoint_ids, &mut adjacency, edge.second);
4140 let phi_edge = edge.first.is_phi() || edge.second.is_phi();
4141 adjacency[first].push((second, edge.weight, phi_edge));
4142 adjacency[second].push((first, edge.weight, phi_edge));
4143 }
4144
4145 let mut seen = vec![false; adjacency.len()];
4146 let mut components = Vec::new();
4147 for start in 0..adjacency.len() {
4148 if seen[start] {
4149 continue;
4150 }
4151
4152 let mut stack = vec![start];
4153 let mut vertices = FastHashSet::default();
4154 let mut edge_visits = 0u64;
4155 let mut weight_visits = 0u64;
4156 let mut phi_visits = 0u64;
4157 seen[start] = true;
4158
4159 while let Some(v) = stack.pop() {
4160 vertices.insert(v);
4161 for &(next, weight, phi_edge) in &adjacency[v] {
4162 edge_visits += 1;
4163 weight_visits += weight;
4164 phi_visits += u64::from(phi_edge);
4165 if !seen[next] {
4166 seen[next] = true;
4167 stack.push(next);
4168 }
4169 }
4170 }
4171
4172 let edges = edge_visits / 2;
4173 let phi_edges = phi_visits / 2;
4174 let cyclic = edges > 0 && vertices.iter().all(|&v| adjacency[v].len() == 2);
4175 components.push(SerialComponent {
4176 endpoints: vertices.len(),
4177 edges,
4178 weight: weight_visits / 2,
4179 phi_edges,
4180 cyclic,
4181 });
4182 }
4183
4184 components.sort_by_key(|component| {
4185 (
4186 std::cmp::Reverse(component.edges),
4187 std::cmp::Reverse(component.weight),
4188 component.endpoints,
4189 )
4190 });
4191
4192 let stats = SerialContractionStats {
4193 input_edges: matrix.stats.edges,
4194 components: components.len() as u64,
4195 phi_edges: matrix.stats.phi_edges,
4196 cyclic_components: components
4197 .iter()
4198 .filter(|component| component.cyclic)
4199 .count() as u64,
4200 };
4201
4202 SerialContraction { components, stats }
4203 }
4204}
4205
4206pub struct SerialDiscontinuityExpander;
4208
4209impl SerialDiscontinuityExpander {
4210 pub fn infer<const K: usize>(
4211 source: PathInfo<K>,
4212 source_side: Side,
4213 target_side: Side,
4214 weight: u64,
4215 ) -> PathInfo<K> {
4216 let rank = if source_side == source.exit_side {
4217 source.rank + weight
4218 } else {
4219 source.rank.saturating_sub(weight)
4220 };
4221 let exit_side = if source_side == source.exit_side {
4222 target_side.inverse()
4223 } else {
4224 target_side
4225 };
4226
4227 PathInfo {
4228 path_id: source.path_id,
4229 rank,
4230 exit_side,
4231 is_cycle: source.is_cycle,
4232 }
4233 }
4234
4235 #[inline(always)]
4236 fn infer_compact(
4237 source: CompactExpansionPathInfo,
4238 source_side: Side,
4239 target_side: Side,
4240 weight: u64,
4241 ) -> CompactExpansionPathInfo {
4242 let rank = if source_side == source.exit_side() {
4243 source.rank() + weight
4244 } else {
4245 source.rank().saturating_sub(weight)
4246 };
4247 let exit_side = if source_side == source.exit_side() {
4248 target_side.inverse()
4249 } else {
4250 target_side
4251 };
4252 CompactExpansionPathInfo::new(source.path_id, rank, exit_side, source.is_cycle())
4253 }
4254
4255 #[allow(clippy::too_many_arguments)]
4256 fn expand_non_diagonal_raw<const K: usize>(
4257 matrix: &BlockedEdgeMatrix<K>,
4258 partition: usize,
4259 map: &ExpansionPathInfoTable<K>,
4260 vertex_path_info_dir: &Path,
4261 edge_writers: &ConcurrentStitchedCoordWriters<'_>,
4262 ranges: &[ExternalLocalUnitigRange],
4263 range_index: &ExternalRangeIndex,
4264 ranges_per_bucket: usize,
4265 error_path: &Path,
4266 pool: &ThreadPool,
4267 ) -> Result<ExpandedPartition<K>, SerialCollationError> {
4268 let blocks = (partition + 1..matrix.partition_count())
4274 .filter_map(|col| {
4275 let block_id = matrix.block_index(partition, col);
4276 let block = &matrix.blocks[block_id];
4277 (block.edges != 0).then_some((col, block_id))
4278 })
4279 .collect::<Vec<_>>();
4280 let pass_inputs = blocks
4281 .iter()
4282 .map(|(_, block_id)| (*block_id, matrix.blocks[*block_id].extents.as_slice()))
4283 .collect::<Vec<_>>();
4284 let pass = pool.install(|| matrix.containers.read_pass(&pass_inputs))?;
4285 let mut tasks = Vec::new();
4286 for (slot, bytes) in pass.extents() {
4287 let (col, block_id) = blocks[slot];
4288 let record_len = matrix.blocks[block_id].record_len;
4289 if bytes.len() % record_len != 0 {
4290 return Err(SerialCollationError::MalformedCoordBucket(
4291 matrix.containers.path_for(block_id),
4292 ));
4293 }
4294 let bytes_per_task = (1024 * 1024 / record_len).max(1) * record_len;
4295 for chunk in bytes.chunks(bytes_per_task) {
4296 tasks.push((col, record_len, chunk));
4297 }
4298 }
4299 for (col, block_id) in blocks.iter().copied() {
4300 let block = &matrix.blocks[block_id];
4301 let bytes_per_task = (1024 * 1024 / block.record_len).max(1) * block.record_len;
4302 for chunk in block.buffer.chunks(bytes_per_task) {
4303 tasks.push((col, block.record_len, chunk));
4304 }
4305 }
4306 let outputs = pool.install(|| {
4307 tasks
4308 .into_par_iter()
4309 .map_init(
4310 || {
4311 (
4312 Vec::<u8>::new(),
4313 (0..edge_writers.writers.len())
4314 .map(|_| Vec::<StitchedCoordRecord>::new())
4315 .collect::<Vec<_>>(),
4316 Vec::<usize>::new(),
4317 )
4318 },
4319 |(vertex_records, path_records, used_path_buckets),
4320 (col, record_len, bytes)| {
4321 vertex_records.clear();
4322 debug_assert!(used_path_buckets.is_empty());
4323 let mut phantoms = Vec::new();
4324 let mut unresolved = 0u64;
4325 let mut inferred = 0u64;
4326 let mut emitted = 0u64;
4327 for encoded in bytes.chunks_exact(record_len) {
4328 let edge = decode_discontinuity_edge::<K>(encoded);
4329 let (MatrixEndpoint::Vertex(x), MatrixEndpoint::Vertex(y)) =
4330 (edge.first, edge.second)
4331 else {
4332 continue;
4333 };
4334 let Some(x_info) = map.get_compact(x.vertex) else {
4335 unresolved += 1;
4336 continue;
4337 };
4338 let y_info = Self::infer_compact(x_info, x.side, y.side, edge.weight);
4339 if y_info.rank() > 0 {
4340 append_encoded_compact_vertex_path_info::<K>(
4341 vertex_records,
4342 y.vertex,
4343 y_info,
4344 );
4345 inferred += 1;
4346 }
4347 if edge.weight == 1 {
4348 let record = stitched_record_from_compact_edge_path_info(
4349 &edge,
4350 compact_edge_path_info(&edge, x_info, y_info),
4351 error_path,
4352 )?;
4353 if let Some(phantom) = edge.phantom_unitig {
4354 phantoms.push((record, phantom));
4355 } else {
4356 let bucket_id = edge_path_info_bucket(
4357 &edge,
4358 ranges,
4359 range_index,
4360 ranges_per_bucket,
4361 )
4362 .ok_or_else(|| {
4363 SerialCollationError::MalformedCoordBucket(
4364 error_path.to_path_buf(),
4365 )
4366 })?;
4367 let bucket = &mut path_records[bucket_id];
4368 if bucket.is_empty() {
4369 used_path_buckets.push(bucket_id);
4370 }
4371 bucket.push(record);
4372 }
4373 emitted += 1;
4374 }
4375 }
4376 if !vertex_records.is_empty() {
4377 let path = vertex_path_info_bucket_path(vertex_path_info_dir, col);
4378 let mut file = OpenOptions::new()
4379 .create(true)
4380 .append(true)
4381 .open(&path)
4382 .map_err(|source| SerialCollationError::Io {
4383 path: path.clone(),
4384 source,
4385 })?;
4386 file.write_all(vertex_records)
4387 .map_err(|source| SerialCollationError::Io { path, source })?;
4388 }
4389 for bucket_id in used_path_buckets.drain(..) {
4390 edge_writers.write_path_records(bucket_id, &path_records[bucket_id])?;
4391 path_records[bucket_id].clear();
4392 }
4393 Ok::<_, SerialCollationError>((phantoms, unresolved, inferred, emitted))
4394 },
4395 )
4396 .collect::<Result<Vec<_>, SerialCollationError>>()
4397 })?;
4398 let mut phantoms = Vec::new();
4399 let mut unresolved = 0;
4400 let mut inferred = 0;
4401 let mut emitted = 0;
4402 for (local_phantoms, local_unresolved, local_inferred, local_emitted) in outputs {
4403 phantoms.extend(local_phantoms);
4404 unresolved += local_unresolved;
4405 inferred += local_inferred;
4406 emitted += local_emitted;
4407 }
4408 Ok((phantoms, unresolved, inferred, emitted))
4409 }
4410
4411 pub fn expand<const K: usize>(
4412 contraction: &FullSerialDiscontinuityContraction<K>,
4413 ) -> SerialExpansion<K> {
4414 Self::expand_impl(contraction, false)
4415 }
4416
4417 pub fn expand_with_original_edges<const K: usize>(
4418 contraction: &FullSerialDiscontinuityContraction<K>,
4419 ) -> SerialExpansion<K> {
4420 Self::expand_impl(contraction, true)
4421 }
4422
4423 pub fn expand_cpp_ordered<const K: usize>(
4424 contraction: &FullSerialDiscontinuityContraction<K>,
4425 ) -> SerialExpansion<K> {
4426 Self::expand_cpp_ordered_impl::<K, true, true>(contraction)
4427 }
4428
4429 #[allow(dead_code)]
4431 fn expand_cpp_ordered_external<const K: usize>(
4432 contraction: &FullSerialDiscontinuityContraction<K>,
4433 ) -> SerialExpansion<K> {
4434 Self::expand_cpp_ordered_impl::<K, false, false>(contraction)
4435 }
4436
4437 fn expand_cpp_ordered_external_to_range_buckets<const K: usize>(
4438 contraction: &mut ExternalBlockedContraction<K>,
4439 ranges: &[ExternalLocalUnitigRange],
4440 ranges_per_bucket: usize,
4441 bucket_count: usize,
4442 error_path: &Path,
4443 expansion_dir: &Path,
4444 threads: usize,
4445 ) -> Result<RangeBucketedExpansion<K>, SerialCollationError> {
4446 if expansion_dir.exists() {
4447 fs::remove_dir_all(expansion_dir).map_err(|source| SerialCollationError::Io {
4448 path: expansion_dir.to_path_buf(),
4449 source,
4450 })?;
4451 }
4452 fs::create_dir_all(expansion_dir).map_err(|source| SerialCollationError::Io {
4453 path: expansion_dir.to_path_buf(),
4454 source,
4455 })?;
4456 let expansion_pool = ThreadPoolBuilder::new()
4457 .num_threads(threads.max(1))
4458 .build()
4459 .map_err(|_| SerialCollationError::WorkerPanic)?;
4460
4461 let partition_count = contraction.vertex_partitions + 1;
4462 let vertex_path_info_dir = contraction.meta_vertex_dir.clone();
4463 let mut vertex_writers =
4464 VertexPathInfoBucketWriters::<K>::open_existing(&vertex_path_info_dir, partition_count);
4465 let meta_vertices_per_partition = &contraction.meta_vertices_per_partition;
4466 let seed_vertices = contraction.meta_vertex_count;
4467
4468 let mut diagonal_by_partition = std::mem::take(&mut contraction.compressed_diagonal_edges);
4471 diagonal_by_partition.resize_with(partition_count, Vec::new);
4472 let matrix = &contraction.expansion_matrix;
4473 let edge_path_info_dir = expansion_dir.join("P_e");
4474 fs::create_dir_all(&edge_path_info_dir).map_err(|source| SerialCollationError::Io {
4475 path: edge_path_info_dir.clone(),
4476 source,
4477 })?;
4478 let edge_writers =
4479 ConcurrentStitchedCoordWriters::new(&edge_path_info_dir, bucket_count, threads);
4480 let range_index = ExternalRangeIndex::new(ranges);
4481 let mut phantom_records = Vec::new();
4482 let mut unresolved_edges = 0u64;
4483 let mut inferred_vertices = 0u64;
4484 let mut edge_path_infos = 0u64;
4485 let max_partition_entries = (1..partition_count)
4486 .map(|partition| {
4487 let incoming = (1..partition)
4488 .map(|row| matrix.blocks[matrix.block_index(row, partition)].edges)
4489 .sum::<usize>();
4490 let diagonal = matrix.blocks[matrix.block_index(partition, partition)].edges;
4491 incoming + 2 * diagonal + meta_vertices_per_partition[partition]
4492 })
4493 .max()
4494 .unwrap_or(1);
4495 let map = ExpansionPathInfoTable::<K>::with_max_entries(max_partition_entries);
4496 eprintln!(
4497 "cuttlefish: expansion table max entries {}, capacity {}, slot {} byte(s)",
4498 max_partition_entries,
4499 map.capacity(),
4500 map.slot_size(),
4501 );
4502 let mut row_load_elapsed = Duration::default();
4503 let mut path_info_load_elapsed = Duration::default();
4504 let mut path_info_read_elapsed = Duration::default();
4505 let mut path_info_insert_elapsed = Duration::default();
4506 let mut path_info_clear_elapsed = Duration::default();
4507 let mut compressed_diagonal_elapsed = Duration::default();
4508 let mut non_diagonal_elapsed = Duration::default();
4509 let mut original_edge_elapsed = Duration::default();
4510 let expansion_work_started = Instant::now();
4511
4512 for partition in 1..partition_count {
4513 let row_load_started = Instant::now();
4514 let row_blocks = if K <= 31 {
4515 vec![matrix.read_flushed_block(partition, partition)?]
4516 } else {
4517 matrix.read_flushed_row(partition, threads)?
4518 };
4519 row_load_elapsed += row_load_started.elapsed();
4520 let diagonal_block = &row_blocks[0];
4521 let phase_started = Instant::now();
4522 map.clear();
4523 path_info_clear_elapsed += phase_started.elapsed();
4524 vertex_writers.flush_bucket(partition)?;
4525 read_vertex_path_info_bucket_into::<K>(
4526 &vertex_path_info_dir,
4527 partition,
4528 error_path,
4529 &map,
4530 &expansion_pool,
4531 &mut path_info_read_elapsed,
4532 &mut path_info_insert_elapsed,
4533 )?;
4534 path_info_load_elapsed += phase_started.elapsed();
4535
4536 let phase_started = Instant::now();
4537 for edge in diagonal_by_partition[partition].iter().rev() {
4538 let (MatrixEndpoint::Vertex(x), MatrixEndpoint::Vertex(y)) =
4539 (edge.first, edge.second)
4540 else {
4541 continue;
4542 };
4543
4544 if K <= 31 {
4545 if let Some(y_info) = map.get_compact(y.vertex) {
4546 let x_info = Self::infer_compact(y_info, y.side, x.side, edge.weight);
4547 if x_info.rank() > 0 && map.insert_compact(x.vertex, x_info) {
4548 inferred_vertices += 1;
4549 }
4550 } else if let Some(x_info) = map.get_compact(x.vertex) {
4551 let y_info = Self::infer_compact(x_info, x.side, y.side, edge.weight);
4552 if y_info.rank() > 0 && map.insert_compact(y.vertex, y_info) {
4553 inferred_vertices += 1;
4554 }
4555 } else {
4556 unresolved_edges += 1;
4557 }
4558 } else if let Some(y_info) = map.get(y.vertex) {
4559 let x_info = Self::infer(y_info, y.side, x.side, edge.weight);
4560 if x_info.rank > 0 && map.insert(x.vertex, x_info) {
4561 inferred_vertices += 1;
4562 }
4563 } else if let Some(x_info) = map.get(x.vertex) {
4564 let y_info = Self::infer(x_info, x.side, y.side, edge.weight);
4565 if y_info.rank > 0 && map.insert(y.vertex, y_info) {
4566 inferred_vertices += 1;
4567 }
4568 } else {
4569 unresolved_edges += 1;
4570 }
4571 }
4572 compressed_diagonal_elapsed += phase_started.elapsed();
4573
4574 let phase_started = Instant::now();
4575 let non_diagonal_cols = partition_count.saturating_sub(partition + 1);
4576 let non_diagonal_edges = row_blocks[1..].iter().map(Vec::len).sum::<usize>();
4577 let workers = threads.max(1).min(non_diagonal_cols.max(1));
4578 if K <= 31 {
4579 let (local_phantoms, local_unresolved, local_inferred, local_edge_infos) =
4580 Self::expand_non_diagonal_raw(
4581 matrix,
4582 partition,
4583 &map,
4584 &vertex_path_info_dir,
4585 &edge_writers,
4586 ranges,
4587 &range_index,
4588 ranges_per_bucket,
4589 error_path,
4590 &expansion_pool,
4591 )?;
4592 phantom_records.extend(local_phantoms);
4593 unresolved_edges += local_unresolved;
4594 inferred_vertices += local_inferred;
4595 edge_path_infos += local_edge_infos;
4596 } else if workers == 1 || non_diagonal_cols < 2 || non_diagonal_edges < 64 * 1024 {
4597 for col in partition + 1..partition_count {
4598 for edge in &row_blocks[col - partition] {
4599 let (MatrixEndpoint::Vertex(x), MatrixEndpoint::Vertex(y)) =
4600 (edge.first, edge.second)
4601 else {
4602 continue;
4603 };
4604 let Some(x_info) = map.get(x.vertex) else {
4605 unresolved_edges += 1;
4606 continue;
4607 };
4608 let y_info = Self::infer(x_info, x.side, y.side, edge.weight);
4609 if y_info.rank > 0 {
4610 vertex_writers.write_record(
4611 col,
4612 &VertexPathInfo {
4613 vertex: y.vertex,
4614 info: y_info,
4615 },
4616 )?;
4617 inferred_vertices += 1;
4618 }
4619 if edge.weight == 1 {
4620 push_edge_path_record_to_range_bucket_writer(
4621 edge,
4622 edge_path_info(edge, x_info, y_info),
4623 ranges,
4624 &range_index,
4625 ranges_per_bucket,
4626 &edge_writers,
4627 &mut phantom_records,
4628 error_path,
4629 )?;
4630 edge_path_infos += 1;
4631 }
4632 }
4633 }
4634 } else {
4635 for col in partition + 1..partition_count {
4636 vertex_writers.flush_bucket(col)?;
4637 }
4638 let chunk = non_diagonal_cols.div_ceil(workers);
4639 let worker_outputs = expansion_pool.install(|| {
4640 (0..workers)
4641 .into_par_iter()
4642 .map(|worker_id| {
4643 let col_start = partition + 1 + worker_id * chunk;
4644 if col_start >= partition_count {
4645 return Ok(None);
4646 }
4647 let col_end = (col_start + chunk).min(partition_count);
4648 let map = ↦
4649 let row_blocks = &row_blocks;
4650 let edge_writers = &edge_writers;
4651 let vertex_path_info_dir = &vertex_path_info_dir;
4652 let mut range_records = (0..bucket_count)
4653 .map(|_| Vec::<StitchedCoordRecord>::new())
4654 .collect::<Vec<_>>();
4655 let mut local_phantoms =
4656 Vec::<(StitchedCoordRecord, DiscontinuityEndpoint<K>)>::new();
4657 let mut local_unresolved = 0u64;
4658 let mut local_inferred = 0u64;
4659 let mut local_edge_infos = 0u64;
4660
4661 for col in col_start..col_end {
4662 let vertex_path =
4663 vertex_path_info_bucket_path(vertex_path_info_dir, col);
4664 let mut vertex_out: Option<BufWriter<File>> = None;
4665 for edge in &row_blocks[col - partition] {
4666 let (MatrixEndpoint::Vertex(x), MatrixEndpoint::Vertex(y)) =
4667 (edge.first, edge.second)
4668 else {
4669 continue;
4670 };
4671 let Some(x_info) = map.get(x.vertex) else {
4672 local_unresolved += 1;
4673 continue;
4674 };
4675 let y_info = Self::infer(x_info, x.side, y.side, edge.weight);
4676 if y_info.rank > 0 {
4677 if vertex_out.is_none() {
4678 let file = OpenOptions::new()
4679 .create(true)
4680 .append(true)
4681 .open(&vertex_path)
4682 .map_err(|source| SerialCollationError::Io {
4683 path: vertex_path.clone(),
4684 source,
4685 })?;
4686 vertex_out = Some(BufWriter::with_capacity(
4687 VERTEX_PATH_INFO_WRITE_BUFFER,
4688 file,
4689 ));
4690 }
4691 vertex_out
4692 .as_mut()
4693 .expect("vertex writer was just created")
4694 .write_all(
4695 &encoded_vertex_path_info_record(&VertexPathInfo {
4696 vertex: y.vertex,
4697 info: y_info,
4698 })[..vertex_path_info_record_len::<K>()],
4699 )
4700 .map_err(|source| SerialCollationError::Io {
4701 path: vertex_path.clone(),
4702 source,
4703 })?;
4704 local_inferred += 1;
4705 }
4706 if edge.weight == 1 {
4707 let info = edge_path_info(edge, x_info, y_info);
4708 let record = stitched_record_from_edge_path_info(
4709 edge, info, error_path,
4710 )?;
4711 if let Some(phantom) = edge.phantom_unitig {
4712 local_phantoms.push((record, phantom));
4713 } else {
4714 let bucket_id = edge_path_info_bucket(
4715 edge,
4716 ranges,
4717 &range_index,
4718 ranges_per_bucket,
4719 )
4720 .ok_or_else(|| {
4721 SerialCollationError::MalformedCoordBucket(
4722 error_path.to_path_buf(),
4723 )
4724 })?;
4725 range_records
4726 .get_mut(bucket_id)
4727 .ok_or_else(|| {
4728 SerialCollationError::MalformedCoordBucket(
4729 error_path.to_path_buf(),
4730 )
4731 })?
4732 .push(record);
4733 }
4734 local_edge_infos += 1;
4735 }
4736 }
4737 if let Some(mut out) = vertex_out {
4738 out.flush().map_err(|source| SerialCollationError::Io {
4739 path: vertex_path,
4740 source,
4741 })?;
4742 }
4743 }
4744
4745 for (bucket_id, records) in range_records.iter().enumerate() {
4746 edge_writers.write_path_records(bucket_id, records)?;
4747 }
4748
4749 Ok::<_, SerialCollationError>(Some((
4750 local_phantoms,
4751 local_unresolved,
4752 local_inferred,
4753 local_edge_infos,
4754 )))
4755 })
4756 .collect::<Result<Vec<_>, SerialCollationError>>()
4757 })?;
4758
4759 for (local_phantoms, local_unresolved, local_inferred, local_edge_infos) in
4760 worker_outputs.into_iter().flatten()
4761 {
4762 unresolved_edges += local_unresolved;
4763 inferred_vertices += local_inferred;
4764 edge_path_infos += local_edge_infos;
4765 phantom_records.extend(local_phantoms);
4766 }
4767 }
4768 non_diagonal_elapsed += phase_started.elapsed();
4769
4770 let phase_started = Instant::now();
4771 let diagonal_records_per_chunk =
4772 (1024 * 1024 / std::mem::size_of::<DiscontinuityEdge<K>>()).max(1);
4773 let diagonal_outputs = expansion_pool.install(|| {
4774 diagonal_block
4775 .par_chunks(diagonal_records_per_chunk)
4776 .map(|edges| {
4777 let mut records = (0..bucket_count)
4778 .map(|_| None::<Vec<StitchedCoordRecord>>)
4779 .collect::<Vec<_>>();
4780 let mut used_buckets = Vec::new();
4781 let mut phantoms = Vec::new();
4782 let mut local_unresolved = 0u64;
4783 let mut local_emitted = 0u64;
4784 for edge in edges {
4785 if edge.weight != 1 {
4786 continue;
4787 }
4788 let (MatrixEndpoint::Vertex(x), MatrixEndpoint::Vertex(y)) =
4789 (edge.first, edge.second)
4790 else {
4791 continue;
4792 };
4793 let record = if K <= 31 {
4794 let Some(x_info) = map.get_compact(x.vertex) else {
4795 local_unresolved += 1;
4796 continue;
4797 };
4798 let y_info =
4799 Self::infer_compact(x_info, x.side, y.side, edge.weight);
4800 stitched_record_from_compact_edge_path_info(
4801 edge,
4802 compact_edge_path_info(edge, x_info, y_info),
4803 error_path,
4804 )?
4805 } else {
4806 let Some(x_info) = map.get(x.vertex) else {
4807 local_unresolved += 1;
4808 continue;
4809 };
4810 let y_info = Self::infer(x_info, x.side, y.side, edge.weight);
4811 stitched_record_from_edge_path_info(
4812 edge,
4813 edge_path_info(edge, x_info, y_info),
4814 error_path,
4815 )?
4816 };
4817 if let Some(phantom) = edge.phantom_unitig {
4818 phantoms.push((record, phantom));
4819 } else {
4820 let bucket_id = edge_path_info_bucket(
4821 edge,
4822 ranges,
4823 &range_index,
4824 ranges_per_bucket,
4825 )
4826 .ok_or_else(|| {
4827 SerialCollationError::MalformedCoordBucket(
4828 error_path.to_path_buf(),
4829 )
4830 })?;
4831 let bucket = &mut records[bucket_id];
4832 if bucket.is_none() {
4833 *bucket = Some(Vec::new());
4834 used_buckets.push(bucket_id);
4835 }
4836 bucket
4837 .as_mut()
4838 .expect("edge bucket was just initialized")
4839 .push(record);
4840 }
4841 local_emitted += 1;
4842 }
4843 for bucket_id in used_buckets {
4844 edge_writers.write_path_records(
4845 bucket_id,
4846 records[bucket_id]
4847 .as_deref()
4848 .expect("used edge bucket is initialized"),
4849 )?;
4850 }
4851 Ok::<_, SerialCollationError>((phantoms, local_unresolved, local_emitted))
4852 })
4853 .collect::<Result<Vec<_>, SerialCollationError>>()
4854 })?;
4855 for (phantoms, local_unresolved, local_emitted) in diagonal_outputs {
4856 phantom_records.extend(phantoms);
4857 unresolved_edges += local_unresolved;
4858 edge_path_infos += local_emitted;
4859 }
4860
4861 let phi_index = matrix.block_index(0, partition);
4862 let phi_block = &matrix.blocks[phi_index];
4863 if phi_block.edges != 0 {
4864 let bytes = matrix
4865 .containers
4866 .read_block(phi_index, &phi_block.extents)?;
4867 if bytes.len() + phi_block.buffer.len() != phi_block.edges * phi_block.record_len {
4868 return Err(SerialCollationError::MalformedCoordBucket(
4869 matrix.containers.path_for(phi_index),
4870 ));
4871 }
4872 let records_per_chunk = (1024 * 1024 / phi_block.record_len).max(1);
4873 let bytes_per_chunk = records_per_chunk * phi_block.record_len;
4874 let phi_outputs = expansion_pool.install(|| {
4875 bytes
4876 .par_chunks(bytes_per_chunk)
4877 .chain(phi_block.buffer.par_chunks(bytes_per_chunk))
4878 .map(|chunk| {
4879 let mut records = (0..bucket_count)
4880 .map(|_| None::<Vec<StitchedCoordRecord>>)
4881 .collect::<Vec<_>>();
4882 let mut used_buckets = Vec::new();
4883 let mut phantoms = Vec::new();
4884 let mut local_unresolved = 0u64;
4885 let mut local_emitted = 0u64;
4886 for encoded in chunk.chunks_exact(phi_block.record_len) {
4887 let edge = decode_discontinuity_edge::<K>(encoded);
4888 if edge.weight != 1 {
4889 continue;
4890 }
4891 let (MatrixEndpoint::Phi, MatrixEndpoint::Vertex(v)) =
4892 (edge.first, edge.second)
4893 else {
4894 continue;
4895 };
4896 let record = if K <= 31 {
4897 let Some(v_info) = map.get_compact(v.vertex) else {
4898 local_unresolved += 1;
4899 continue;
4900 };
4901 stitched_record_from_compact_edge_path_info(
4902 &edge,
4903 compact_phi_edge_path_info(&edge, v_info),
4904 error_path,
4905 )?
4906 } else {
4907 let Some(v_info) = map.get(v.vertex) else {
4908 local_unresolved += 1;
4909 continue;
4910 };
4911 stitched_record_from_edge_path_info(
4912 &edge,
4913 phi_edge_path_info(&edge, v_info),
4914 error_path,
4915 )?
4916 };
4917 if let Some(phantom) = edge.phantom_unitig {
4918 phantoms.push((record, phantom));
4919 } else {
4920 let bucket_id = edge_path_info_bucket(
4921 &edge,
4922 ranges,
4923 &range_index,
4924 ranges_per_bucket,
4925 )
4926 .ok_or_else(|| {
4927 SerialCollationError::MalformedCoordBucket(
4928 error_path.to_path_buf(),
4929 )
4930 })?;
4931 let bucket = &mut records[bucket_id];
4932 if bucket.is_none() {
4933 *bucket = Some(Vec::new());
4934 used_buckets.push(bucket_id);
4935 }
4936 bucket
4937 .as_mut()
4938 .expect("edge bucket was just initialized")
4939 .push(record);
4940 }
4941 local_emitted += 1;
4942 }
4943 for bucket_id in used_buckets {
4944 edge_writers.write_path_records(
4945 bucket_id,
4946 records[bucket_id]
4947 .as_deref()
4948 .expect("used edge bucket is initialized"),
4949 )?;
4950 }
4951 Ok::<_, SerialCollationError>((
4952 phantoms,
4953 local_unresolved,
4954 local_emitted,
4955 ))
4956 })
4957 .collect::<Result<Vec<_>, SerialCollationError>>()
4958 })?;
4959 for (phantoms, local_unresolved, local_emitted) in phi_outputs {
4960 phantom_records.extend(phantoms);
4961 unresolved_edges += local_unresolved;
4962 edge_path_infos += local_emitted;
4963 }
4964 }
4965 original_edge_elapsed += phase_started.elapsed();
4966 }
4967
4968 drop(vertex_writers);
4969 eprintln!(
4970 "cuttlefish: blocked expansion detail: row load/decode {:.3}s, propagation/emission {:.3}s",
4971 row_load_elapsed.as_secs_f64(),
4972 expansion_work_started
4973 .elapsed()
4974 .saturating_sub(row_load_elapsed)
4975 .as_secs_f64()
4976 );
4977 eprintln!(
4978 "cuttlefish: path-info load phases: clear {:.3}s, read {:.3}s, decode/insert {:.3}s",
4979 path_info_clear_elapsed.as_secs_f64(),
4980 path_info_read_elapsed.as_secs_f64(),
4981 path_info_insert_elapsed.as_secs_f64(),
4982 );
4983 eprintln!(
4984 "cuttlefish: blocked expansion phases: path-info load {:.3}s, compressed diagonal {:.3}s, non-diagonal {:.3}s, original diagonal/phi {:.3}s",
4985 path_info_load_elapsed.as_secs_f64(),
4986 compressed_diagonal_elapsed.as_secs_f64(),
4987 non_diagonal_elapsed.as_secs_f64(),
4988 original_edge_elapsed.as_secs_f64(),
4989 );
4990 let mut record_manifest = edge_writers.finish(&expansion_pool)?;
4991 record_manifest.sort_by(|left, right| {
4992 left.bucket_id
4993 .cmp(&right.bucket_id)
4994 .then_with(|| left.path.cmp(&right.path))
4995 });
4996
4997 Ok(RangeBucketedExpansion {
4998 stats: SerialExpansionStats {
4999 seed_vertices,
5000 inferred_vertices,
5001 edge_path_infos,
5002 unresolved_edges,
5003 },
5004 records: Vec::new(),
5005 record_manifest,
5006 phantom_records,
5007 })
5008 }
5009
5010 fn expand_cpp_ordered_impl<
5011 const K: usize,
5012 const COLLECT_VERTICES: bool,
5013 const SORT_DEDUP_EDGES: bool,
5014 >(
5015 contraction: &FullSerialDiscontinuityContraction<K>,
5016 ) -> SerialExpansion<K> {
5017 let partition_count = contraction.vertex_partitions + 1;
5018 let mut vertex_buckets = vec![Vec::<VertexPathInfo<K>>::new(); partition_count];
5019 for meta in &contraction.meta_vertices {
5020 if meta.partition < partition_count {
5021 vertex_buckets[meta.partition].push(VertexPathInfo {
5022 vertex: meta.vertex,
5023 info: PathInfo {
5024 path_id: meta.vertex,
5025 rank: meta.weight,
5026 exit_side: meta.entry_side,
5027 is_cycle: meta.is_cycle,
5028 },
5029 });
5030 }
5031 }
5032 let seed_vertices = vertex_buckets.iter().map(Vec::len).sum::<usize>() as u64;
5033
5034 let mut diagonal_by_partition = contraction.compressed_diagonal_edges.clone();
5035 diagonal_by_partition.resize_with(partition_count, Vec::new);
5036 if diagonal_by_partition.iter().all(Vec::is_empty) {
5037 for partition in &contraction.partitions {
5038 if partition.partition < partition_count {
5039 diagonal_by_partition[partition.partition]
5040 .extend(partition.compressed_diagonal_edges.iter().cloned());
5041 }
5042 }
5043 }
5044
5045 let mut non_diagonal_by_row =
5046 vec![Vec::<(usize, DiscontinuityEdge<K>)>::new(); partition_count];
5047 let mut diagonal_by_row = vec![Vec::<DiscontinuityEdge<K>>::new(); partition_count];
5048 let mut phi_by_partition = vec![Vec::<DiscontinuityEdge<K>>::new(); partition_count];
5049 for edge in &contraction.expansion_edges {
5050 let (row, col) =
5051 edge_matrix_row_col(contraction.vertex_partitions, edge.first, edge.second);
5052 if row == 0 {
5053 if col < partition_count {
5054 phi_by_partition[col].push(edge.clone());
5055 }
5056 } else if row == col {
5057 if row < partition_count {
5058 diagonal_by_row[row].push(edge.clone());
5059 }
5060 } else if row < partition_count {
5061 non_diagonal_by_row[row].push((col, edge.clone()));
5062 }
5063 }
5064 for row in 1..partition_count {
5065 non_diagonal_by_row[row].sort_by_key(|(col, edge)| {
5066 (
5067 *col,
5068 endpoint_sort_key(edge.first),
5069 endpoint_sort_key(edge.second),
5070 edge.weight,
5071 edge.unitig_index,
5072 )
5073 });
5074 diagonal_by_row[row].sort_by_key(|edge| {
5075 (
5076 endpoint_sort_key(edge.first),
5077 endpoint_sort_key(edge.second),
5078 edge.weight,
5079 edge.unitig_index,
5080 )
5081 });
5082 phi_by_partition[row].sort_by_key(|edge| {
5083 (
5084 endpoint_sort_key(edge.first),
5085 endpoint_sort_key(edge.second),
5086 edge.weight,
5087 edge.unitig_index,
5088 )
5089 });
5090 }
5091 let mut edges = Vec::new();
5092 let mut unresolved_edges = 0u64;
5093 let mut inferred_vertices = 0u64;
5094 let mut map = FastHashMap::<Kmer<K>, PathInfo<K>>::default();
5095
5096 for partition in 1..partition_count {
5097 map.clear();
5098 for vertex_info in &vertex_buckets[partition] {
5099 map.insert(vertex_info.vertex, vertex_info.info);
5100 }
5101
5102 for edge in diagonal_by_partition[partition].iter().rev() {
5103 let (MatrixEndpoint::Vertex(x), MatrixEndpoint::Vertex(y)) =
5104 (edge.first, edge.second)
5105 else {
5106 continue;
5107 };
5108
5109 if let Some(&y_info) = map.get(&y.vertex) {
5110 let x_info = Self::infer(y_info, y.side, x.side, edge.weight);
5111 if x_info.rank > 0 && !map.contains_key(&x.vertex) {
5112 map.insert(x.vertex, x_info);
5113 inferred_vertices += 1;
5114 }
5115 } else if let Some(&x_info) = map.get(&x.vertex) {
5116 let y_info = Self::infer(x_info, x.side, y.side, edge.weight);
5117 if y_info.rank > 0 && !map.contains_key(&y.vertex) {
5118 map.insert(y.vertex, y_info);
5119 inferred_vertices += 1;
5120 }
5121 } else {
5122 unresolved_edges += 1;
5123 }
5124 }
5125
5126 for (col, edge) in &non_diagonal_by_row[partition] {
5127 let (MatrixEndpoint::Vertex(x), MatrixEndpoint::Vertex(y)) =
5128 (edge.first, edge.second)
5129 else {
5130 continue;
5131 };
5132 let Some(&x_info) = map.get(&x.vertex) else {
5133 unresolved_edges += 1;
5134 continue;
5135 };
5136 let y_info = Self::infer(x_info, x.side, y.side, edge.weight);
5137 if y_info.rank > 0 {
5138 vertex_buckets[*col].push(VertexPathInfo {
5139 vertex: y.vertex,
5140 info: y_info,
5141 });
5142 inferred_vertices += 1;
5143 }
5144 if edge.weight == 1 {
5145 edges.push(EdgePathInfo {
5146 unitig_index: edge.unitig_index,
5147 phantom_unitig: edge.phantom_unitig,
5148 info: edge_path_info(edge, x_info, y_info),
5149 });
5150 }
5151 }
5152
5153 for edge in &diagonal_by_row[partition] {
5154 if edge.weight != 1 {
5155 continue;
5156 }
5157 let (MatrixEndpoint::Vertex(x), MatrixEndpoint::Vertex(y)) =
5158 (edge.first, edge.second)
5159 else {
5160 continue;
5161 };
5162 let Some(&x_info) = map.get(&x.vertex) else {
5163 unresolved_edges += 1;
5164 continue;
5165 };
5166 let y_info = Self::infer(x_info, x.side, y.side, edge.weight);
5167 edges.push(EdgePathInfo {
5168 unitig_index: edge.unitig_index,
5169 phantom_unitig: edge.phantom_unitig,
5170 info: edge_path_info(edge, x_info, y_info),
5171 });
5172 }
5173
5174 for edge in &phi_by_partition[partition] {
5175 if edge.weight != 1 {
5176 continue;
5177 }
5178 let (MatrixEndpoint::Phi, MatrixEndpoint::Vertex(v)) = (edge.first, edge.second)
5179 else {
5180 continue;
5181 };
5182 if let Some(&v_info) = map.get(&v.vertex) {
5183 edges.push(EdgePathInfo {
5184 unitig_index: edge.unitig_index,
5185 phantom_unitig: edge.phantom_unitig,
5186 info: phi_edge_path_info(edge, v_info),
5187 });
5188 } else {
5189 unresolved_edges += 1;
5190 }
5191 }
5192 }
5193
5194 if SORT_DEDUP_EDGES {
5195 edges.sort_by_key(|edge| {
5196 (
5197 edge.unitig_index,
5198 edge.info.path_id.as_u128(),
5199 edge.info.rank,
5200 edge.info.exit_side as u8,
5201 )
5202 });
5203 edges.dedup();
5204 }
5205
5206 let vertices = if COLLECT_VERTICES {
5207 let mut vertices = vertex_buckets
5208 .into_iter()
5209 .flatten()
5210 .collect::<Vec<VertexPathInfo<K>>>();
5211 vertices.sort_by_key(|entry| {
5212 (
5213 entry.vertex.as_u128(),
5214 entry.info.path_id.as_u128(),
5215 entry.info.rank,
5216 entry.info.exit_side as u8,
5217 )
5218 });
5219 vertices.dedup();
5220 vertices
5221 } else {
5222 Vec::new()
5223 };
5224
5225 SerialExpansion {
5226 stats: SerialExpansionStats {
5227 seed_vertices,
5228 inferred_vertices,
5229 edge_path_infos: edges.len() as u64,
5230 unresolved_edges,
5231 },
5232 vertices,
5233 edges,
5234 }
5235 }
5236
5237 fn expand_impl<const K: usize>(
5238 contraction: &FullSerialDiscontinuityContraction<K>,
5239 include_original_edges: bool,
5240 ) -> SerialExpansion<K> {
5241 let mut vertex_info = FastHashMap::<Kmer<K>, PathInfo<K>>::with_capacity_and_hasher(
5242 contraction.meta_vertices.len(),
5243 FastBuildHasher::default(),
5244 );
5245 let mut edges = Vec::new();
5246 let mut unresolved_edges = 0u64;
5247
5248 for meta in &contraction.meta_vertices {
5249 vertex_info.insert(
5250 meta.vertex,
5251 PathInfo {
5252 path_id: meta.vertex,
5253 rank: meta.weight,
5254 exit_side: meta.entry_side,
5255 is_cycle: meta.is_cycle,
5256 },
5257 );
5258 }
5259 let seed_vertices = vertex_info.len() as u64;
5260
5261 let mut expansion_edges =
5262 if include_original_edges && !contraction.expansion_edges.is_empty() {
5263 contraction.expansion_edges.clone()
5264 } else {
5265 contraction.final_edges.clone()
5266 };
5267 expansion_edges.sort_by_key(|edge| {
5268 (
5269 endpoint_sort_key(edge.first),
5270 endpoint_sort_key(edge.second),
5271 edge.weight,
5272 edge.unitig_index,
5273 )
5274 });
5275
5276 let mut changed = true;
5277 while changed {
5278 changed = false;
5279 for edge in &expansion_edges {
5280 match (edge.first, edge.second) {
5281 (MatrixEndpoint::Phi, MatrixEndpoint::Vertex(v))
5282 | (MatrixEndpoint::Vertex(v), MatrixEndpoint::Phi) => {
5283 if edge.weight == 1 {
5284 if let Some(&info) = vertex_info.get(&v.vertex) {
5285 edges.push(EdgePathInfo {
5286 unitig_index: edge.unitig_index,
5287 phantom_unitig: edge.phantom_unitig,
5288 info: phi_edge_path_info(edge, info),
5289 });
5290 }
5291 }
5292 }
5293 (MatrixEndpoint::Vertex(x), MatrixEndpoint::Vertex(y)) => {
5294 let x_info = vertex_info.get(&x.vertex).copied();
5295 let y_info = vertex_info.get(&y.vertex).copied();
5296 match (x_info, y_info) {
5297 (Some(x_info), None) => {
5298 let inferred = Self::infer(x_info, x.side, y.side, edge.weight);
5299 if inferred.rank > 0 {
5300 vertex_info.insert(y.vertex, inferred);
5301 changed = true;
5302 }
5303 }
5304 (None, Some(y_info)) => {
5305 let inferred = Self::infer(y_info, y.side, x.side, edge.weight);
5306 if inferred.rank > 0 {
5307 vertex_info.insert(x.vertex, inferred);
5308 changed = true;
5309 }
5310 }
5311 (Some(x_info), Some(y_info))
5312 if edge.weight == 1 && x_info.path_id == y_info.path_id =>
5313 {
5314 edges.push(EdgePathInfo {
5315 unitig_index: edge.unitig_index,
5316 phantom_unitig: edge.phantom_unitig,
5317 info: edge_path_info(edge, x_info, y_info),
5318 });
5319 }
5320 _ => {}
5321 }
5322 }
5323 _ => {}
5324 }
5325 }
5326 }
5327
5328 for edge in &expansion_edges {
5329 match (edge.first, edge.second) {
5330 (MatrixEndpoint::Phi, MatrixEndpoint::Vertex(v))
5331 | (MatrixEndpoint::Vertex(v), MatrixEndpoint::Phi)
5332 if !vertex_info.contains_key(&v.vertex) =>
5333 {
5334 unresolved_edges += 1;
5335 }
5336 (MatrixEndpoint::Vertex(x), MatrixEndpoint::Vertex(y))
5337 if !vertex_info.contains_key(&x.vertex)
5338 || !vertex_info.contains_key(&y.vertex) =>
5339 {
5340 unresolved_edges += 1;
5341 }
5342 _ => {}
5343 }
5344 }
5345
5346 edges.sort_by_key(|edge| {
5347 (
5348 edge.unitig_index,
5349 edge.info.path_id.as_u128(),
5350 edge.info.rank,
5351 edge.info.exit_side as u8,
5352 )
5353 });
5354 edges.dedup();
5355
5356 let mut vertices = vertex_info
5357 .into_iter()
5358 .map(|(vertex, info)| VertexPathInfo { vertex, info })
5359 .collect::<Vec<_>>();
5360 vertices.sort_by_key(|entry| {
5361 (
5362 entry.vertex.as_u128(),
5363 entry.info.path_id.as_u128(),
5364 entry.info.rank,
5365 entry.info.exit_side as u8,
5366 )
5367 });
5368
5369 SerialExpansion {
5370 stats: SerialExpansionStats {
5371 seed_vertices,
5372 inferred_vertices: (vertices.len() as u64).saturating_sub(seed_vertices),
5373 edge_path_infos: edges.len() as u64,
5374 unresolved_edges,
5375 },
5376 vertices,
5377 edges,
5378 }
5379 }
5380}
5381
5382pub struct SerialUncoloredCollator;
5387
5388impl SerialUncoloredCollator {
5389 pub fn collate<const K: usize>(
5390 inputs: &DiscontinuityInputs<K>,
5391 expansion: &SerialExpansion<K>,
5392 ) -> SerialCollation {
5393 Self::collate_with_threads(inputs, expansion, 1)
5394 }
5395
5396 pub fn collate_with_threads<const K: usize>(
5397 inputs: &DiscontinuityInputs<K>,
5398 expansion: &SerialExpansion<K>,
5399 threads: usize,
5400 ) -> SerialCollation {
5401 eprintln!(
5402 "cuttlefish: collating {} path edge info record(s) against {} local unitig(s)",
5403 expansion.edges.len(),
5404 inputs.unitigs.len()
5405 );
5406 let mut records = Vec::<CollationRecord<K>>::new();
5407 let mut missing_unitig_labels = 0u64;
5408 let mut path_unitigs = vec![false; inputs.unitigs.len()];
5409
5410 let collect_started = Instant::now();
5411 for edge in &expansion.edges {
5412 let Some(label) = inputs.try_label(edge.unitig_index) else {
5413 missing_unitig_labels += 1;
5414 continue;
5415 };
5416 path_unitigs[edge.unitig_index] = true;
5417 records.push(CollationRecord {
5418 info: edge.info,
5419 label: label.to_vec(),
5420 });
5421 }
5422 eprintln!(
5423 "cuttlefish: collation collected path records in {:.3}s",
5424 collect_started.elapsed().as_secs_f64()
5425 );
5426
5427 let record_sort_started = Instant::now();
5428 records.sort_by_key(|record| {
5429 (
5430 record.info.path_id.as_u128(),
5431 record.info.rank,
5432 record.info.exit_side as u8,
5433 )
5434 });
5435 eprintln!(
5436 "cuttlefish: collation sorted path records in {:.3}s",
5437 record_sort_started.elapsed().as_secs_f64()
5438 );
5439
5440 let expanded_started = Instant::now();
5441 let mut expanded_unitigs = Vec::new();
5442 let mut start = 0;
5443 while start < records.len() {
5444 let path_id = records[start].info.path_id;
5445 let is_cycle = records[start].info.is_cycle;
5446 let mut end = start + 1;
5447 while end < records.len() && records[end].info.path_id == path_id {
5448 end += 1;
5449 }
5450
5451 let mut label = Vec::new();
5452 if end - start == 2
5453 && !is_cycle
5454 && records[start].info.rank == 0
5455 && records[start + 1].info.rank == 0
5456 {
5457 append_or_init::<K>(
5458 &mut label,
5459 oriented_label(
5460 &records[start].label,
5461 records[start].info.exit_side == Side::Front,
5462 ),
5463 );
5464 append_or_init::<K>(
5465 &mut label,
5466 oriented_label(
5467 &records[start + 1].label,
5468 records[start + 1].info.exit_side != Side::Front,
5469 ),
5470 );
5471 } else {
5472 for record in &records[start..end] {
5473 append_or_init::<K>(
5474 &mut label,
5475 oriented_label(&record.label, record.info.exit_side == Side::Front),
5476 );
5477 }
5478 }
5479
5480 if is_cycle && label.len() >= K {
5481 label.truncate(label.len() - (K - 1));
5482 }
5483
5484 expanded_unitigs.push(canonical_label(label));
5485 start = end;
5486 }
5487 expanded_unitigs.sort_unstable();
5488 expanded_unitigs.dedup();
5489 eprintln!(
5490 "cuttlefish: collation built expanded path unitigs in {:.3}s",
5491 expanded_started.elapsed().as_secs_f64()
5492 );
5493 eprintln!(
5494 "cuttlefish: collated {} expanded path unitig(s)",
5495 expanded_unitigs.len()
5496 );
5497
5498 let stitch_started = Instant::now();
5499 let mut stitched_unitigs = stitch_discontinuity_paths::<K>(inputs, &[], threads);
5500 stitched_unitigs.sort_unstable();
5501 stitched_unitigs.dedup();
5502 let stitched_discontinuity_unitigs = stitched_unitigs.len() as u64;
5503 eprintln!(
5504 "cuttlefish: collation stitched discontinuity paths in {:.3}s",
5505 stitch_started.elapsed().as_secs_f64()
5506 );
5507 eprintln!(
5508 "cuttlefish: added {} stitched discontinuity unitig(s)",
5509 stitched_discontinuity_unitigs
5510 );
5511
5512 let suppress_started = Instant::now();
5513 let suppressed_expanded =
5514 suppress_labels_contained_in_sources(&mut expanded_unitigs, &stitched_unitigs, threads);
5515 eprintln!(
5516 "cuttlefish: collation suppressed contained expanded labels in {:.3}s",
5517 suppress_started.elapsed().as_secs_f64()
5518 );
5519 if suppressed_expanded > 0 {
5520 eprintln!(
5521 "cuttlefish: suppressed {} expanded path unitig(s) contained in stitched paths",
5522 suppressed_expanded
5523 );
5524 }
5525
5526 let mut unitigs = Vec::with_capacity(
5527 expanded_unitigs.len() + stitched_unitigs.len() + inputs.unitigs.len(),
5528 );
5529 let mut unitig_origins = Vec::new();
5530 let merge_started = Instant::now();
5531 for unitig in expanded_unitigs {
5532 unitigs.push(unitig);
5533 unitig_origins.push("expanded");
5534 }
5535 for unitig in stitched_unitigs {
5536 unitigs.push(unitig);
5537 unitig_origins.push("stitched");
5538 }
5539
5540 let mut direct_local_unitigs = 0u64;
5541 for (index, unitig) in inputs.unitigs.iter().enumerate() {
5542 if !path_unitigs[index] && unitig.left_exit().is_none() && unitig.right_exit().is_none()
5543 {
5544 unitigs.push(canonical_label(unitig.label(inputs).to_vec()));
5545 unitig_origins.push("direct");
5546 direct_local_unitigs += 1;
5547 }
5548 }
5549 eprintln!(
5550 "cuttlefish: collation merged direct local unitigs in {:.3}s",
5551 merge_started.elapsed().as_secs_f64()
5552 );
5553 eprintln!(
5554 "cuttlefish: added {} direct local unitig(s)",
5555 direct_local_unitigs
5556 );
5557
5558 let final_sort_started = Instant::now();
5559 unitigs = bucketed_maximal_unitig_reduce(unitigs, threads);
5560 eprintln!(
5561 "cuttlefish: collation bucket reduce completed in {:.3}s",
5562 final_sort_started.elapsed().as_secs_f64()
5563 );
5564
5565 SerialCollation {
5566 stats: SerialCollationStats {
5567 input_path_infos: expansion.edges.len() as u64,
5568 emitted_unitigs: unitigs.len() as u64,
5569 emitted_bases: unitigs.iter().map(|unitig| unitig.len() as u64).sum(),
5570 missing_unitig_labels,
5571 direct_local_unitigs,
5572 stitched_discontinuity_unitigs,
5573 },
5574 unitigs,
5575 }
5576 }
5577
5578 pub fn collate_path_info_only_with_threads<const K: usize>(
5579 inputs: &DiscontinuityInputs<K>,
5580 expansion: &SerialExpansion<K>,
5581 threads: usize,
5582 ) -> SerialCollation {
5583 eprintln!(
5584 "cuttlefish: collating {} path edge info record(s) against {} local unitig(s)",
5585 expansion.edges.len(),
5586 inputs.unitigs.len()
5587 );
5588 let mut records = Vec::<CollationRecord<K>>::with_capacity(expansion.edges.len());
5589 let mut missing_unitig_labels = 0u64;
5590 let mut path_unitigs = vec![false; inputs.unitigs.len()];
5591
5592 let collect_started = Instant::now();
5593 for edge in &expansion.edges {
5594 let Some(label) = inputs.try_label(edge.unitig_index) else {
5595 missing_unitig_labels += 1;
5596 continue;
5597 };
5598 path_unitigs[edge.unitig_index] = true;
5599 records.push(CollationRecord {
5600 info: edge.info,
5601 label: label.to_vec(),
5602 });
5603 }
5604 eprintln!(
5605 "cuttlefish: path-info collation collected records in {:.3}s",
5606 collect_started.elapsed().as_secs_f64()
5607 );
5608
5609 let record_sort_started = Instant::now();
5610 records.sort_by_key(|record| {
5611 (
5612 record.info.path_id.as_u128(),
5613 record.info.rank,
5614 record.info.exit_side as u8,
5615 )
5616 });
5617 eprintln!(
5618 "cuttlefish: path-info collation sorted records in {:.3}s",
5619 record_sort_started.elapsed().as_secs_f64()
5620 );
5621
5622 let build_started = Instant::now();
5623 let mut unitigs = Vec::new();
5624 let mut start = 0;
5625 while start < records.len() {
5626 let path_id = records[start].info.path_id;
5627 let is_cycle = records[start].info.is_cycle;
5628 let mut end = start + 1;
5629 while end < records.len() && records[end].info.path_id == path_id {
5630 end += 1;
5631 }
5632
5633 let mut label = Vec::new();
5634 if end - start == 2
5635 && !is_cycle
5636 && records[start].info.rank == 0
5637 && records[start + 1].info.rank == 0
5638 {
5639 append_or_init::<K>(
5640 &mut label,
5641 oriented_label(
5642 &records[start].label,
5643 records[start].info.exit_side == Side::Front,
5644 ),
5645 );
5646 append_or_init::<K>(
5647 &mut label,
5648 oriented_label(
5649 &records[start + 1].label,
5650 records[start + 1].info.exit_side != Side::Front,
5651 ),
5652 );
5653 } else {
5654 for record in &records[start..end] {
5655 append_or_init::<K>(
5656 &mut label,
5657 oriented_label(&record.label, record.info.exit_side == Side::Front),
5658 );
5659 }
5660 }
5661
5662 if label.len() >= K {
5663 if is_cycle {
5664 label.truncate(label.len().saturating_sub(K - 1));
5665 unitigs.push(normalize_stitched_cycle::<K>(&label));
5666 } else {
5667 unitigs.push(canonical_label(label));
5668 }
5669 }
5670 start = end;
5671 }
5672 let stitched_discontinuity_unitigs = unitigs.len() as u64;
5673 eprintln!(
5674 "cuttlefish: path-info collation built path unitigs in {:.3}s",
5675 build_started.elapsed().as_secs_f64()
5676 );
5677 eprintln!(
5678 "cuttlefish: added {} path-info discontinuity unitig(s)",
5679 stitched_discontinuity_unitigs
5680 );
5681
5682 let merge_started = Instant::now();
5683 let mut direct_local_unitigs = 0u64;
5684 for (index, unitig) in inputs.unitigs.iter().enumerate() {
5685 if !path_unitigs[index] && unitig.left_exit().is_none() && unitig.right_exit().is_none()
5686 {
5687 unitigs.push(canonical_label(unitig.label(inputs).to_vec()));
5688 direct_local_unitigs += 1;
5689 }
5690 }
5691 eprintln!(
5692 "cuttlefish: path-info collation merged direct local unitigs in {:.3}s",
5693 merge_started.elapsed().as_secs_f64()
5694 );
5695 eprintln!(
5696 "cuttlefish: added {} direct local unitig(s)",
5697 direct_local_unitigs
5698 );
5699
5700 let final_sort_started = Instant::now();
5701 unitigs = bucketed_maximal_unitig_reduce(unitigs, threads);
5702 eprintln!(
5703 "cuttlefish: path-info collation bucket reduce completed in {:.3}s",
5704 final_sort_started.elapsed().as_secs_f64()
5705 );
5706
5707 SerialCollation {
5708 stats: SerialCollationStats {
5709 input_path_infos: expansion.edges.len() as u64,
5710 emitted_unitigs: unitigs.len() as u64,
5711 emitted_bases: unitigs.iter().map(|unitig| unitig.len() as u64).sum(),
5712 missing_unitig_labels,
5713 direct_local_unitigs,
5714 stitched_discontinuity_unitigs,
5715 },
5716 unitigs,
5717 }
5718 }
5719
5720 pub fn collate_stitched_with_threads<const K: usize>(
5721 inputs: &DiscontinuityInputs<K>,
5722 threads: usize,
5723 ) -> SerialCollation {
5724 Self::collate_stitched_with_threads_impl(inputs, threads, None)
5725 .expect("in-memory stitched collation cannot fail")
5726 }
5727
5728 pub fn collate_stitched_with_threads_in_dir<const K: usize>(
5729 inputs: &DiscontinuityInputs<K>,
5730 threads: usize,
5731 coord_dir: &Path,
5732 ) -> Result<SerialCollation, SerialCollationError> {
5733 Self::collate_stitched_with_threads_impl(inputs, threads, Some(coord_dir))
5734 }
5735
5736 pub fn collate_stitched_to_fasta_with_threads_in_dir<const K: usize>(
5737 inputs: &DiscontinuityInputs<K>,
5738 threads: usize,
5739 coord_dir: &Path,
5740 final_dir: &Path,
5741 output_path: &Path,
5742 ) -> Result<SerialCollationStats, SerialCollationError> {
5743 eprintln!(
5744 "cuttlefish: collating {} local unitig(s) by discontinuity-end stitching",
5745 inputs.unitigs.len()
5746 );
5747
5748 let spill_started = Instant::now();
5749 let mut final_buckets = FinalUnitigBucketWriters::create(final_dir, threads)?;
5750 let stitched_discontinuity_unitigs = stitch_discontinuity_paths_to_final_buckets::<K>(
5751 inputs,
5752 &[],
5753 threads,
5754 coord_dir,
5755 &mut final_buckets,
5756 )?;
5757
5758 let mut direct_local_unitigs = 0u64;
5759 for unitig in &inputs.unitigs {
5760 if unitig.left_exit().is_none() && unitig.right_exit().is_none() {
5761 let label = canonical_label(unitig.label(inputs).to_vec());
5762 final_buckets.write_label(&label)?;
5763 direct_local_unitigs += 1;
5764 }
5765 }
5766 let manifest = final_buckets.finish()?;
5767 eprintln!(
5768 "cuttlefish: collation stitched and spilled final-label candidates in {:.3}s",
5769 spill_started.elapsed().as_secs_f64()
5770 );
5771 eprintln!(
5772 "cuttlefish: added {} stitched discontinuity unitig(s)",
5773 stitched_discontinuity_unitigs
5774 );
5775 eprintln!(
5776 "cuttlefish: added {} direct local unitig(s)",
5777 direct_local_unitigs
5778 );
5779
5780 let reduce_started = Instant::now();
5781 let (emitted_unitigs, emitted_bases) =
5782 reduce_final_unitig_buckets_to_fasta(&manifest, output_path)?;
5783 eprintln!(
5784 "cuttlefish: collation external bucket reduce and FASTA write completed in {:.3}s",
5785 reduce_started.elapsed().as_secs_f64()
5786 );
5787
5788 Ok(SerialCollationStats {
5789 input_path_infos: 0,
5790 emitted_unitigs,
5791 emitted_bases,
5792 missing_unitig_labels: 0,
5793 direct_local_unitigs,
5794 stitched_discontinuity_unitigs,
5795 })
5796 }
5797
5798 pub fn collate_external_stitched_to_fasta_with_threads_in_dir<const K: usize>(
5799 inputs: &mut ExternalDiscontinuityInputs<K>,
5800 threads: usize,
5801 coord_dir: &Path,
5802 final_dir: &Path,
5803 output_path: &Path,
5804 ) -> Result<SerialCollationStats, SerialCollationError> {
5805 eprintln!(
5806 "cuttlefish: collating {} local unitig(s) by external discontinuity-end stitching",
5807 inputs.unitig_count()
5808 );
5809
5810 let spill_started = Instant::now();
5811 let colored = inputs.color_runs.is_some()
5812 || inputs
5813 .local_unitig_buckets
5814 .as_ref()
5815 .is_some_and(|buckets| buckets.iter().any(|bucket| bucket.colored));
5816 let mut direct_local_unitigs = 0u64;
5817 let trivial_started = Instant::now();
5818 let trivial = inputs.trivial_fasta.take();
5819 let mut final_buckets = if inputs.trivial_is_output {
5820 let (records, bases) = trivial.map_or((0, 0), |(_, records, bases)| (records, bases));
5822 direct_local_unitigs = records;
5823 FinalUnitigBucketWriters::adopt_direct(output_path, colored, records, bases)?
5824 } else {
5825 let mut writers = FinalUnitigBucketWriters::create_direct(output_path, colored)?;
5826 if let Some((path, records, bases)) = trivial {
5827 writers.append_direct_fasta_file(&path, records, bases)?;
5828 remove_serial_file(&path)?;
5829 direct_local_unitigs = records;
5830 }
5831 writers
5832 };
5833 let trivial_elapsed = trivial_started.elapsed();
5834 let use_cpp_path_info = std::env::var_os("CF3_RS_ENDPOINT_STITCH").is_none();
5835 let mut direct_local_complete = false;
5836 let stitched_discontinuity_unitigs = if use_cpp_path_info {
5837 let result = collate_external_cpp_path_info_to_final_buckets::<K>(
5838 inputs,
5839 threads,
5840 coord_dir,
5841 &mut final_buckets,
5842 )?;
5843 direct_local_unitigs += result.direct_local_unitigs;
5844 direct_local_complete = result.direct_local_unitigs_complete;
5845 result.stitched_unitigs
5846 } else {
5847 stitch_external_discontinuity_paths_to_final_buckets::<K>(
5848 inputs,
5849 threads,
5850 coord_dir,
5851 &mut final_buckets,
5852 )?
5853 };
5854
5855 let fallback_started = Instant::now();
5856 if !direct_local_complete {
5857 let reader = ExternalDiscontinuityReader::open(inputs)?;
5858 let mut color_reader = None;
5859 let mut range_index = 0usize;
5860 let mut scratch = Vec::new();
5861 for (unitig_index, unitig) in reader.iter()?.enumerate() {
5862 let unitig = unitig?;
5863 if inputs
5864 .ranges
5865 .get(range_index)
5866 .is_some_and(|range| range.start_unitig == unitig_index)
5867 {
5868 color_reader = inputs
5869 .color_runs
5870 .as_ref()
5871 .map(|sidecar| sidecar.reader_at(inputs.ranges[range_index].color_start))
5872 .transpose()?;
5873 range_index += 1;
5874 }
5875 let colors = color_reader
5876 .as_mut()
5877 .map(|reader| reader.read_next())
5878 .transpose()?;
5879 if unitig.left_exit().is_none() && unitig.right_exit().is_none() {
5880 reader.read_label(&unitig, &mut scratch)?;
5881 let reverse = reverse_complement_is_less(&scratch);
5882 let label = canonical_label(scratch.clone());
5883 if let Some(mut colors) = colors {
5884 if reverse {
5885 colors = reverse_color_runs(
5886 &colors,
5887 (unitig.label_len as usize - K + 1) as u32,
5888 );
5889 }
5890 final_buckets.write_colored_label(&label, &colors)?;
5891 } else {
5892 final_buckets.write_label(&label)?;
5893 }
5894 direct_local_unitigs += 1;
5895 }
5896 }
5897 }
5898 let fallback_elapsed = fallback_started.elapsed();
5899 let cleanup_started = Instant::now();
5900 if !keep_intermediates() {
5901 remove_serial_file(&inputs.unitig_path)?;
5902 remove_serial_file(&inputs.label_path)?;
5903 if let Some(sidecar) = inputs.color_runs.as_ref() {
5904 remove_serial_file(&sidecar.run_path)?;
5905 }
5906 }
5907 let cleanup_elapsed = cleanup_started.elapsed();
5908 let finish_started = Instant::now();
5909 let manifest = final_buckets.finish()?;
5910 eprintln!(
5911 "cuttlefish: collation serial detail: trivial-fasta adopt/copy {:.3}s (in-place {}), direct-local fallback {:.3}s (ran {}), cleanup {:.3}s, finish {:.3}s",
5912 trivial_elapsed.as_secs_f64(),
5913 inputs.trivial_is_output,
5914 fallback_elapsed.as_secs_f64(),
5915 !direct_local_complete,
5916 cleanup_elapsed.as_secs_f64(),
5917 finish_started.elapsed().as_secs_f64(),
5918 );
5919 eprintln!(
5920 "cuttlefish: external collation stitched and spilled final-label candidates in {:.3}s",
5921 spill_started.elapsed().as_secs_f64()
5922 );
5923 eprintln!(
5924 "cuttlefish: added {} stitched discontinuity unitig(s)",
5925 stitched_discontinuity_unitigs
5926 );
5927 eprintln!(
5928 "cuttlefish: added {} direct local unitig(s)",
5929 direct_local_unitigs
5930 );
5931
5932 let reduce_started = Instant::now();
5933 let (emitted_unitigs, emitted_bases) =
5934 reduce_final_unitig_buckets_to_fasta(&manifest, output_path)?;
5935 if !keep_intermediates() {
5936 let _ = fs::remove_dir(coord_dir);
5937 let _ = fs::remove_dir(final_dir);
5938 }
5939 eprintln!(
5940 "cuttlefish: collation external bucket reduce and FASTA write completed in {:.3}s",
5941 reduce_started.elapsed().as_secs_f64()
5942 );
5943
5944 Ok(SerialCollationStats {
5945 input_path_infos: 0,
5946 emitted_unitigs,
5947 emitted_bases,
5948 missing_unitig_labels: 0,
5949 direct_local_unitigs,
5950 stitched_discontinuity_unitigs,
5951 })
5952 }
5953
5954 fn collate_stitched_with_threads_impl<const K: usize>(
5955 inputs: &DiscontinuityInputs<K>,
5956 threads: usize,
5957 coord_dir: Option<&Path>,
5958 ) -> Result<SerialCollation, SerialCollationError> {
5959 eprintln!(
5960 "cuttlefish: collating {} local unitig(s) by discontinuity-end stitching",
5961 inputs.unitigs.len()
5962 );
5963
5964 let stitch_started = Instant::now();
5965 let mut unitigs = if let Some(coord_dir) = coord_dir {
5966 stitch_discontinuity_paths_with_coord_dir::<K>(inputs, &[], threads, coord_dir)?
5967 } else {
5968 stitch_discontinuity_paths::<K>(inputs, &[], threads)
5969 };
5970 let stitched_discontinuity_unitigs = unitigs.len() as u64;
5971 eprintln!(
5972 "cuttlefish: collation stitched discontinuity paths in {:.3}s",
5973 stitch_started.elapsed().as_secs_f64()
5974 );
5975 eprintln!(
5976 "cuttlefish: added {} stitched discontinuity unitig(s)",
5977 stitched_discontinuity_unitigs
5978 );
5979
5980 let merge_started = Instant::now();
5981 let mut direct_local_unitigs = 0u64;
5982 for unitig in &inputs.unitigs {
5983 if unitig.left_exit().is_none() && unitig.right_exit().is_none() {
5984 unitigs.push(canonical_label(unitig.label(inputs).to_vec()));
5985 direct_local_unitigs += 1;
5986 }
5987 }
5988 eprintln!(
5989 "cuttlefish: collation merged direct local unitigs in {:.3}s",
5990 merge_started.elapsed().as_secs_f64()
5991 );
5992 eprintln!(
5993 "cuttlefish: added {} direct local unitig(s)",
5994 direct_local_unitigs
5995 );
5996
5997 let final_sort_started = Instant::now();
5998 unitigs = bucketed_maximal_unitig_reduce(unitigs, threads);
5999 eprintln!(
6000 "cuttlefish: collation bucket reduce completed in {:.3}s",
6001 final_sort_started.elapsed().as_secs_f64()
6002 );
6003
6004 Ok(SerialCollation {
6005 stats: SerialCollationStats {
6006 input_path_infos: 0,
6007 emitted_unitigs: unitigs.len() as u64,
6008 emitted_bases: unitigs.iter().map(|unitig| unitig.len() as u64).sum(),
6009 missing_unitig_labels: 0,
6010 direct_local_unitigs,
6011 stitched_discontinuity_unitigs,
6012 },
6013 unitigs,
6014 })
6015 }
6016}
6017
6018#[derive(Debug, Clone, PartialEq, Eq)]
6019struct CollationRecord<const K: usize> {
6020 info: PathInfo<K>,
6021 label: Vec<u8>,
6022}
6023
6024fn append_or_init<const K: usize>(label: &mut Vec<u8>, next: Vec<u8>) {
6025 if label.is_empty() {
6026 label.extend_from_slice(&next);
6027 } else {
6028 label.extend_from_slice(&next[K..]);
6029 }
6030}
6031
6032fn oriented_label(label: &[u8], reverse_complement: bool) -> Vec<u8> {
6033 if reverse_complement {
6034 reverse_complement_label(label)
6035 } else {
6036 label.to_vec()
6037 }
6038}
6039
6040fn reverse_complement_label(label: &[u8]) -> Vec<u8> {
6041 label
6042 .iter()
6043 .rev()
6044 .map(|&base| complement_ascii(base))
6045 .collect()
6046}
6047
6048fn canonical_label(label: Vec<u8>) -> Vec<u8> {
6049 if reverse_complement_is_less(&label) {
6050 reverse_complement_label(&label)
6051 } else {
6052 label
6053 }
6054}
6055
6056fn bucketed_maximal_unitig_reduce(mut unitigs: Vec<Vec<u8>>, threads: usize) -> Vec<Vec<u8>> {
6057 if unitigs.len() <= 2048 {
6058 unitigs.sort_unstable();
6059 unitigs.dedup();
6060 return unitigs;
6061 }
6062
6063 let bucket_count = (threads.max(1) * 1024)
6064 .next_power_of_two()
6065 .clamp(1024, 16_384);
6066 let bucket_mask = bucket_count - 1;
6067 let mut buckets = (0..bucket_count).map(|_| Vec::new()).collect::<Vec<_>>();
6068 for unitig in unitigs {
6069 let bucket = (hash_bytes(&unitig, 0) as usize) & bucket_mask;
6070 buckets[bucket].push(unitig);
6071 }
6072
6073 let workers = threads.max(1).min(bucket_count);
6074 if workers == 1 {
6075 for bucket in &mut buckets {
6076 bucket.sort_unstable();
6077 bucket.dedup();
6078 }
6079 } else {
6080 let chunk_size = bucket_count.div_ceil(workers);
6081 std::thread::scope(|scope| {
6082 for chunk in buckets.chunks_mut(chunk_size) {
6083 scope.spawn(move || {
6084 for bucket in chunk {
6085 bucket.sort_unstable();
6086 bucket.dedup();
6087 }
6088 });
6089 }
6090 });
6091 }
6092
6093 let total = buckets.iter().map(Vec::len).sum();
6094 let mut reduced = Vec::with_capacity(total);
6095 for mut bucket in buckets {
6096 reduced.append(&mut bucket);
6097 }
6098 reduced
6099}
6100
6101#[derive(Debug, Clone, PartialEq, Eq)]
6102struct FinalUnitigBucketEntry {
6103 bucket_id: usize,
6104 path: PathBuf,
6105 records: u64,
6106 bases: u64,
6107 colored: bool,
6108 direct_output: bool,
6109}
6110
6111struct FinalUnitigBucketWriters {
6112 dir: PathBuf,
6113 bucket_mask: usize,
6114 next_bucket: usize,
6115 writers: Vec<Option<FinalUnitigBucketWriter>>,
6116 records: Vec<u64>,
6117 bases: Vec<u64>,
6118 open_writers: usize,
6119 colored: bool,
6120 direct_output: Option<BufWriter<File>>,
6121 direct_output_path: Option<PathBuf>,
6122 direct_header_scratch: Vec<u8>,
6124 direct_records: u64,
6125 direct_record_id_highwater: u64,
6126 direct_bases: u64,
6127}
6128
6129impl FinalUnitigBucketWriters {
6130 fn create(dir: &Path, threads: usize) -> Result<Self, SerialCollationError> {
6131 Self::create_with_color(dir, threads, false)
6132 }
6133
6134 fn create_with_color(
6135 dir: &Path,
6136 threads: usize,
6137 colored: bool,
6138 ) -> Result<Self, SerialCollationError> {
6139 if dir.exists() {
6140 fs::remove_dir_all(dir).map_err(|source| SerialCollationError::Io {
6141 path: dir.to_path_buf(),
6142 source,
6143 })?;
6144 }
6145 fs::create_dir_all(dir).map_err(|source| SerialCollationError::Io {
6146 path: dir.to_path_buf(),
6147 source,
6148 })?;
6149 let bucket_count = final_unitig_bucket_count(threads);
6150 let mut writers = Vec::with_capacity(bucket_count);
6151 writers.resize_with(bucket_count, || None);
6152 Ok(Self {
6153 dir: dir.to_path_buf(),
6154 bucket_mask: bucket_count - 1,
6155 next_bucket: 0,
6156 writers,
6157 records: vec![0; bucket_count],
6158 bases: vec![0; bucket_count],
6159 open_writers: 0,
6160 colored,
6161 direct_output: None,
6162 direct_output_path: None,
6163 direct_header_scratch: Vec::new(),
6164 direct_records: 0,
6165 direct_record_id_highwater: 0,
6166 direct_bases: 0,
6167 })
6168 }
6169
6170 fn create_direct(output_path: &Path, colored: bool) -> Result<Self, SerialCollationError> {
6171 let file = File::create(output_path).map_err(|source| SerialCollationError::Io {
6172 path: output_path.to_path_buf(),
6173 source,
6174 })?;
6175 Ok(Self {
6176 dir: output_path.parent().unwrap_or(Path::new(".")).to_path_buf(),
6177 bucket_mask: 0,
6178 next_bucket: 0,
6179 writers: Vec::new(),
6180 records: Vec::new(),
6181 bases: Vec::new(),
6182 open_writers: 0,
6183 colored,
6184 direct_output: Some(BufWriter::with_capacity(8 * 1024 * 1024, file)),
6185 direct_output_path: Some(output_path.to_path_buf()),
6186 direct_header_scratch: Vec::new(),
6187 direct_records: 0,
6188 direct_record_id_highwater: 0,
6189 direct_bases: 0,
6190 })
6191 }
6192
6193 fn adopt_direct(
6196 output_path: &Path,
6197 colored: bool,
6198 records: u64,
6199 bases: u64,
6200 ) -> Result<Self, SerialCollationError> {
6201 let mut file = OpenOptions::new()
6202 .write(true)
6203 .open(output_path)
6204 .map_err(|source| SerialCollationError::Io {
6205 path: output_path.to_path_buf(),
6206 source,
6207 })?;
6208 file.seek(SeekFrom::End(0))
6211 .map_err(|source| SerialCollationError::Io {
6212 path: output_path.to_path_buf(),
6213 source,
6214 })?;
6215 Ok(Self {
6216 dir: output_path.parent().unwrap_or(Path::new(".")).to_path_buf(),
6217 bucket_mask: 0,
6218 next_bucket: 0,
6219 writers: Vec::new(),
6220 records: Vec::new(),
6221 bases: Vec::new(),
6222 open_writers: 0,
6223 colored,
6224 direct_output: Some(BufWriter::with_capacity(8 * 1024 * 1024, file)),
6225 direct_output_path: Some(output_path.to_path_buf()),
6226 direct_header_scratch: Vec::new(),
6227 direct_records: records,
6228 direct_record_id_highwater: records,
6229 direct_bases: bases,
6230 })
6231 }
6232
6233 fn write_label(&mut self, label: &[u8]) -> Result<(), SerialCollationError> {
6234 if self.direct_output.is_some() {
6235 return self.write_direct_record(label, &[]);
6236 }
6237 let bucket_id = self.next_bucket;
6238 self.next_bucket = (self.next_bucket + 1) & self.bucket_mask;
6239 if self.writers[bucket_id].is_none() {
6240 self.evict_writer_if_needed(bucket_id)?;
6241 self.writers[bucket_id] = Some(FinalUnitigBucketWriter::open(
6242 &self.dir,
6243 bucket_id,
6244 self.colored,
6245 )?);
6246 self.open_writers += 1;
6247 }
6248 self.writers[bucket_id]
6249 .as_mut()
6250 .expect("final unitig bucket writer was just created")
6251 .write_record(label, &[])?;
6252 self.records[bucket_id] += 1;
6253 self.bases[bucket_id] += label.len() as u64;
6254 Ok(())
6255 }
6256
6257 fn write_colored_label(
6258 &mut self,
6259 label: &[u8],
6260 colors: &[UnitigColor],
6261 ) -> Result<(), SerialCollationError> {
6262 if !self.colored {
6263 return Err(SerialCollationError::MalformedCoordBucket(self.dir.clone()));
6264 }
6265 if self.direct_output.is_some() {
6266 return self.write_direct_record(label, colors);
6267 }
6268 let bucket_id = self.next_bucket;
6269 self.next_bucket = (self.next_bucket + 1) & self.bucket_mask;
6270 if self.writers[bucket_id].is_none() {
6271 self.evict_writer_if_needed(bucket_id)?;
6272 self.writers[bucket_id] =
6273 Some(FinalUnitigBucketWriter::open(&self.dir, bucket_id, true)?);
6274 self.open_writers += 1;
6275 }
6276 self.writers[bucket_id]
6277 .as_mut()
6278 .expect("final unitig bucket writer was just created")
6279 .write_record(label, colors)?;
6280 self.records[bucket_id] += 1;
6281 self.bases[bucket_id] += label.len() as u64;
6282 Ok(())
6283 }
6284
6285 fn write_direct_record(
6286 &mut self,
6287 label: &[u8],
6288 colors: &[UnitigColor],
6289 ) -> Result<(), SerialCollationError> {
6290 self.direct_records += 1;
6291 self.direct_record_id_highwater += 1;
6292 self.direct_bases += label.len() as u64;
6293 let header = &mut self.direct_header_scratch;
6296 header.clear();
6297 header.reserve(4 + colors.len() * 12);
6298 header.extend_from_slice(b">0");
6299 for color in colors {
6300 header.push(b' ');
6301 append_decimal_u64(header, color.raw());
6302 }
6303 header.push(b'\n');
6304 let out = self.direct_output.as_mut().expect("direct output exists");
6305 out.write_all(header)
6306 .and_then(|_| out.write_all(label))
6307 .and_then(|_| out.write_all(b"\n"))
6308 .map_err(|source| SerialCollationError::Io {
6309 path: self
6310 .direct_output_path
6311 .clone()
6312 .expect("direct output path exists"),
6313 source,
6314 })
6315 }
6316
6317 fn write_direct_batch(
6318 &mut self,
6319 bytes: &[u8],
6320 records: u64,
6321 bases: u64,
6322 ) -> Result<(), SerialCollationError> {
6323 self.direct_output
6324 .as_mut()
6325 .expect("direct output exists")
6326 .write_all(bytes)
6327 .map_err(|source| SerialCollationError::Io {
6328 path: self
6329 .direct_output_path
6330 .clone()
6331 .expect("direct output path exists"),
6332 source,
6333 })?;
6334 self.direct_records += records;
6335 self.direct_bases += bases;
6336 Ok(())
6337 }
6338
6339 fn append_direct_fasta_file(
6340 &mut self,
6341 path: &Path,
6342 records: u64,
6343 bases: u64,
6344 ) -> Result<(), SerialCollationError> {
6345 let file = File::open(path).map_err(|source| SerialCollationError::Io {
6346 path: path.to_path_buf(),
6347 source,
6348 })?;
6349 std::io::copy(
6350 &mut BufReader::with_capacity(8 * 1024 * 1024, file),
6351 self.direct_output.as_mut().expect("direct output exists"),
6352 )
6353 .map_err(|source| SerialCollationError::Io {
6354 path: path.to_path_buf(),
6355 source,
6356 })?;
6357 self.direct_records += records;
6358 self.direct_record_id_highwater += records;
6359 self.direct_bases += bases;
6360 Ok(())
6361 }
6362
6363 fn prepare_parallel_direct_output(
6364 &mut self,
6365 ) -> Result<(File, PathBuf, u64), SerialCollationError> {
6366 let path = self
6367 .direct_output_path
6368 .clone()
6369 .expect("direct output path exists");
6370 let output = self.direct_output.as_mut().expect("direct output exists");
6371 output.flush().map_err(|source| SerialCollationError::Io {
6372 path: path.clone(),
6373 source,
6374 })?;
6375 let file = output
6376 .get_ref()
6377 .try_clone()
6378 .map_err(|source| SerialCollationError::Io {
6379 path: path.clone(),
6380 source,
6381 })?;
6382 let offset = file
6383 .metadata()
6384 .map_err(|source| SerialCollationError::Io {
6385 path: path.clone(),
6386 source,
6387 })?
6388 .len();
6389 Ok((file, path, offset))
6390 }
6391
6392 fn evict_writer_if_needed(
6393 &mut self,
6394 requested_bucket_id: usize,
6395 ) -> Result<(), SerialCollationError> {
6396 if self.open_writers < MAX_OPEN_STITCH_ENDPOINT_WRITERS {
6397 return Ok(());
6398 }
6399 let evict_bucket_id = self
6400 .writers
6401 .iter()
6402 .enumerate()
6403 .find_map(|(bucket_id, writer)| {
6404 (bucket_id != requested_bucket_id && writer.is_some()).then_some(bucket_id)
6405 })
6406 .unwrap_or(requested_bucket_id);
6407 if let Some(mut writer) = self.writers[evict_bucket_id].take() {
6408 writer.flush()?;
6409 self.open_writers -= 1;
6410 }
6411 Ok(())
6412 }
6413
6414 fn finish(mut self) -> Result<Vec<FinalUnitigBucketEntry>, SerialCollationError> {
6415 if let Some(mut output) = self.direct_output.take() {
6416 output.flush().map_err(|source| SerialCollationError::Io {
6417 path: self
6418 .direct_output_path
6419 .clone()
6420 .expect("direct output path exists"),
6421 source,
6422 })?;
6423 return Ok(vec![FinalUnitigBucketEntry {
6424 bucket_id: 0,
6425 path: self.direct_output_path.expect("direct output path exists"),
6426 records: self.direct_records,
6427 bases: self.direct_bases,
6428 colored: self.colored,
6429 direct_output: true,
6430 }]);
6431 }
6432 let mut manifest = Vec::new();
6433 for writer in self.writers.iter_mut().flatten() {
6434 writer.flush()?;
6435 }
6436 for (bucket_id, &records) in self.records.iter().enumerate() {
6437 if records != 0 {
6438 manifest.push(FinalUnitigBucketEntry {
6439 bucket_id,
6440 path: self.dir.join(format!("{bucket_id:05}.fub")),
6441 records,
6442 bases: self.bases[bucket_id],
6443 colored: self.colored,
6444 direct_output: false,
6445 });
6446 }
6447 }
6448 manifest.sort_by_key(|entry| entry.bucket_id);
6449 Ok(manifest)
6450 }
6451}
6452
6453struct FinalUnitigBucketWriter {
6454 path: PathBuf,
6455 out: BufWriter<File>,
6456 colored: bool,
6457}
6458
6459impl FinalUnitigBucketWriter {
6460 fn open(dir: &Path, bucket_id: usize, colored: bool) -> Result<Self, SerialCollationError> {
6461 let path = dir.join(format!("{bucket_id:05}.fub"));
6462 let file = OpenOptions::new()
6463 .create(true)
6464 .append(true)
6465 .open(&path)
6466 .map_err(|source| SerialCollationError::Io {
6467 path: path.clone(),
6468 source,
6469 })?;
6470 Ok(Self {
6471 path,
6472 out: BufWriter::with_capacity(FINAL_UNITIG_BUCKET_WRITE_BUFFER, file),
6473 colored,
6474 })
6475 }
6476
6477 fn write_record(
6478 &mut self,
6479 label: &[u8],
6480 colors: &[UnitigColor],
6481 ) -> Result<(), SerialCollationError> {
6482 let len = u32::try_from(label.len())
6483 .map_err(|_| SerialCollationError::MalformedCoordBucket(self.path.clone()))?;
6484 self.out
6485 .write_all(&len.to_le_bytes())
6486 .and_then(|_| {
6487 if self.colored {
6488 let count = u32::try_from(colors.len())
6489 .map_err(|_| std::io::Error::from(std::io::ErrorKind::InvalidData))?;
6490 self.out.write_all(&count.to_le_bytes())?;
6491 }
6492 self.out.write_all(label)
6493 })
6494 .and_then(|_| {
6495 if self.colored {
6496 for color in colors {
6497 self.out.write_all(&color.raw().to_le_bytes())?;
6498 }
6499 }
6500 Ok(())
6501 })
6502 .map_err(|source| SerialCollationError::Io {
6503 path: self.path.clone(),
6504 source,
6505 })?;
6506 Ok(())
6507 }
6508
6509 fn flush(&mut self) -> Result<(), SerialCollationError> {
6510 self.out.flush().map_err(|source| SerialCollationError::Io {
6511 path: self.path.clone(),
6512 source,
6513 })
6514 }
6515}
6516
6517fn reduce_final_unitig_buckets_to_fasta(
6518 manifest: &[FinalUnitigBucketEntry],
6519 output_path: &Path,
6520) -> Result<(u64, u64), SerialCollationError> {
6521 if let [entry] = manifest
6522 && entry.direct_output
6523 {
6524 if entry.path != output_path {
6525 return Err(SerialCollationError::MalformedCoordBucket(
6526 entry.path.clone(),
6527 ));
6528 }
6529 return Ok((entry.records, entry.bases));
6530 }
6531 let file = File::create(output_path).map_err(|source| SerialCollationError::Io {
6532 path: output_path.to_path_buf(),
6533 source,
6534 })?;
6535 let mut out = file;
6536 let mut fasta_buffer = Vec::with_capacity(8 * 1024 * 1024);
6537 let mut emitted = 0u64;
6538 let mut bases = 0u64;
6539
6540 for entry in manifest {
6541 let bucket = read_final_unitig_bucket(entry)?;
6542 remove_serial_file(&entry.path)?;
6543 let bytes = &bucket.bytes;
6544 for (label_start, label_len, color_start, color_count) in bucket.labels {
6545 let label = &bytes[label_start..label_start + label_len];
6546 emitted += 1;
6547 bases += label.len() as u64;
6548 fasta_buffer.extend_from_slice(b">0");
6549 for color in &bucket.colors[color_start..color_start + color_count] {
6550 fasta_buffer.push(b' ');
6551 append_decimal_u64(&mut fasta_buffer, color.raw());
6552 }
6553 fasta_buffer.push(b'\n');
6554 fasta_buffer.extend_from_slice(label);
6555 fasta_buffer.push(b'\n');
6556 if fasta_buffer.len() >= 8 * 1024 * 1024 {
6557 out.write_all(&fasta_buffer)
6558 .map_err(|source| SerialCollationError::Io {
6559 path: output_path.to_path_buf(),
6560 source,
6561 })?;
6562 fasta_buffer.clear();
6563 }
6564 }
6565 }
6566
6567 out.write_all(&fasta_buffer)
6568 .map_err(|source| SerialCollationError::Io {
6569 path: output_path.to_path_buf(),
6570 source,
6571 })?;
6572
6573 out.flush().map_err(|source| SerialCollationError::Io {
6574 path: output_path.to_path_buf(),
6575 source,
6576 })?;
6577 Ok((emitted, bases))
6578}
6579
6580#[inline]
6581fn append_decimal_u64(output: &mut Vec<u8>, mut value: u64) {
6582 let mut digits = [0u8; 20];
6583 let mut start = digits.len();
6584 loop {
6585 start -= 1;
6586 digits[start] = b'0' + (value % 10) as u8;
6587 value /= 10;
6588 if value == 0 {
6589 break;
6590 }
6591 }
6592 output.extend_from_slice(&digits[start..]);
6593}
6594
6595struct FinalUnitigBucketData {
6596 bytes: Vec<u8>,
6597 labels: Vec<(usize, usize, usize, usize)>,
6598 colors: Vec<UnitigColor>,
6599}
6600
6601fn read_final_unitig_bucket(
6602 entry: &FinalUnitigBucketEntry,
6603) -> Result<FinalUnitigBucketData, SerialCollationError> {
6604 let bytes = fs::read(&entry.path).map_err(|source| SerialCollationError::Io {
6605 path: entry.path.clone(),
6606 source,
6607 })?;
6608 let mut labels = Vec::with_capacity(entry.records as usize);
6609 let mut colors = Vec::new();
6610 let mut cursor = 0usize;
6611 for _ in 0..entry.records {
6612 let len_end = cursor
6613 .checked_add(4)
6614 .ok_or_else(|| SerialCollationError::MalformedCoordBucket(entry.path.clone()))?;
6615 let len_bytes = bytes
6616 .get(cursor..len_end)
6617 .ok_or_else(|| SerialCollationError::MalformedCoordBucket(entry.path.clone()))?;
6618 let len =
6619 u32::from_le_bytes(len_bytes.try_into().expect("four-byte label length")) as usize;
6620 let color_count = if entry.colored {
6621 let count_end = len_end + 4;
6622 let count_bytes = bytes
6623 .get(len_end..count_end)
6624 .ok_or_else(|| SerialCollationError::MalformedCoordBucket(entry.path.clone()))?;
6625 cursor = count_end;
6626 u32::from_le_bytes(count_bytes.try_into().expect("four-byte color count")) as usize
6627 } else {
6628 cursor = len_end;
6629 0
6630 };
6631 let label_end = cursor
6632 .checked_add(len)
6633 .ok_or_else(|| SerialCollationError::MalformedCoordBucket(entry.path.clone()))?;
6634 bytes
6635 .get(cursor..label_end)
6636 .ok_or_else(|| SerialCollationError::MalformedCoordBucket(entry.path.clone()))?;
6637 let label_start = cursor;
6638 cursor = label_end;
6639 let color_start = colors.len();
6640 for _ in 0..color_count {
6641 let color_end = cursor + 8;
6642 let raw = bytes
6643 .get(cursor..color_end)
6644 .ok_or_else(|| SerialCollationError::MalformedCoordBucket(entry.path.clone()))?;
6645 let raw = u64::from_le_bytes(raw.try_into().expect("eight-byte unitig color"));
6646 colors.push(UnitigColor::new(
6647 (raw & 0xff_ffff) as u32,
6648 crate::state::ColorCoordinate::from_u40(raw >> 24),
6649 ));
6650 cursor = color_end;
6651 }
6652 labels.push((label_start, len, color_start, color_count));
6653 }
6654 if cursor != bytes.len() {
6655 return Err(SerialCollationError::MalformedCoordBucket(
6656 entry.path.clone(),
6657 ));
6658 }
6659 Ok(FinalUnitigBucketData {
6660 bytes,
6661 labels,
6662 colors,
6663 })
6664}
6665
6666fn reverse_complement_is_less(label: &[u8]) -> bool {
6667 for (idx, &base) in label.iter().enumerate() {
6668 let rc_base = complement_ascii(label[label.len() - 1 - idx]);
6669 if rc_base != base {
6670 return rc_base < base;
6671 }
6672 }
6673 false
6674}
6675
6676fn suppress_labels_contained_in_sources(
6677 labels: &mut Vec<Vec<u8>>,
6678 sources: &[Vec<u8>],
6679 threads: usize,
6680) -> usize {
6681 if labels.is_empty() || sources.is_empty() {
6682 return 0;
6683 }
6684
6685 let mut lengths = labels.iter().map(Vec::len).collect::<Vec<_>>();
6686 lengths.sort_unstable();
6687 lengths.dedup();
6688
6689 let mut patterns_by_len = FastHashMap::<usize, FastHashMap<u128, Vec<usize>>>::default();
6690 for &len in &lengths {
6691 let mut patterns = FastHashMap::<u128, Vec<usize>>::default();
6692 for (index, label) in labels
6693 .iter()
6694 .enumerate()
6695 .filter(|(_, label)| label.len() == len)
6696 {
6697 patterns.entry(hash_label(label)).or_default().push(index);
6698 }
6699 patterns_by_len.insert(len, patterns);
6700 }
6701
6702 let powers = lengths
6703 .iter()
6704 .map(|&len| {
6705 (
6706 len,
6707 (hash_pow(HASH_BASE_1, len), hash_pow(HASH_BASE_2, len)),
6708 )
6709 })
6710 .collect::<FastHashMap<_, _>>();
6711
6712 let workers = threads.max(1).min(sources.len().max(1));
6713 let mut keep = vec![true; labels.len()];
6714 if workers == 1 {
6715 for source in sources {
6716 mark_contained_candidates_for_haystack(
6717 source,
6718 &lengths,
6719 &patterns_by_len,
6720 &powers,
6721 labels,
6722 &mut keep,
6723 );
6724 let reverse = reverse_complement_label(source);
6725 mark_contained_candidates_for_haystack(
6726 &reverse,
6727 &lengths,
6728 &patterns_by_len,
6729 &powers,
6730 labels,
6731 &mut keep,
6732 );
6733 }
6734 } else {
6735 let chunk_len = sources.len().div_ceil(workers);
6736 let local_keeps = std::thread::scope(|scope| {
6737 let mut handles = Vec::new();
6738 for chunk in sources.chunks(chunk_len) {
6739 let lengths = &lengths;
6740 let patterns_by_len = &patterns_by_len;
6741 let powers = &powers;
6742 let labels = &*labels;
6743 handles.push(scope.spawn(move || {
6744 let mut local_keep = vec![true; labels.len()];
6745 for source in chunk {
6746 mark_contained_candidates_for_haystack(
6747 source,
6748 lengths,
6749 patterns_by_len,
6750 powers,
6751 labels,
6752 &mut local_keep,
6753 );
6754 let reverse = reverse_complement_label(source);
6755 mark_contained_candidates_for_haystack(
6756 &reverse,
6757 lengths,
6758 patterns_by_len,
6759 powers,
6760 labels,
6761 &mut local_keep,
6762 );
6763 }
6764 local_keep
6765 }));
6766 }
6767
6768 handles
6769 .into_iter()
6770 .map(|handle| handle.join().expect("stitched-containment worker panicked"))
6771 .collect::<Vec<_>>()
6772 });
6773 for local_keep in local_keeps {
6774 for (dst, src) in keep.iter_mut().zip(local_keep) {
6775 *dst &= src;
6776 }
6777 }
6778 }
6779
6780 let before = labels.len();
6781 let mut idx = 0usize;
6782 labels.retain(|_| {
6783 let retain = keep[idx];
6784 idx += 1;
6785 retain
6786 });
6787 before - labels.len()
6788}
6789
6790const HASH_BASE_1: u64 = 1_099_511_628_211;
6791const HASH_BASE_2: u64 = 1_000_000_007;
6792
6793fn mark_contained_candidates_for_haystack(
6794 haystack: &[u8],
6795 lengths: &[usize],
6796 patterns_by_len: &FastHashMap<usize, FastHashMap<u128, Vec<usize>>>,
6797 powers: &FastHashMap<usize, (u64, u64)>,
6798 unitigs: &[Vec<u8>],
6799 keep: &mut [bool],
6800) {
6801 if haystack.is_empty() {
6802 return;
6803 }
6804 let mut prefix1 = Vec::with_capacity(haystack.len() + 1);
6805 let mut prefix2 = Vec::with_capacity(haystack.len() + 1);
6806 prefix1.push(0u64);
6807 prefix2.push(0u64);
6808 for &base in haystack {
6809 let code = hash_base_code(base);
6810 prefix1.push(
6811 prefix1
6812 .last()
6813 .copied()
6814 .unwrap()
6815 .wrapping_mul(HASH_BASE_1)
6816 .wrapping_add(code),
6817 );
6818 prefix2.push(
6819 prefix2
6820 .last()
6821 .copied()
6822 .unwrap()
6823 .wrapping_mul(HASH_BASE_2)
6824 .wrapping_add(code),
6825 );
6826 }
6827
6828 for &len in lengths {
6829 if len >= haystack.len() {
6830 continue;
6831 }
6832 let Some(patterns) = patterns_by_len.get(&len) else {
6833 continue;
6834 };
6835 let Some(&(pow1, pow2)) = powers.get(&len) else {
6836 continue;
6837 };
6838 for start in 0..=haystack.len() - len {
6839 let end = start + len;
6840 let h1 = prefix1[end].wrapping_sub(prefix1[start].wrapping_mul(pow1));
6841 let h2 = prefix2[end].wrapping_sub(prefix2[start].wrapping_mul(pow2));
6842 let hash = ((h1 as u128) << 64) | h2 as u128;
6843 let Some(indices) = patterns.get(&hash) else {
6844 continue;
6845 };
6846 let window = &haystack[start..end];
6847 for &index in indices {
6848 if keep[index] && unitigs[index].as_slice() == window {
6849 keep[index] = false;
6850 }
6851 }
6852 }
6853 }
6854}
6855
6856fn hash_label(label: &[u8]) -> u128 {
6857 let mut h1 = 0u64;
6858 let mut h2 = 0u64;
6859 for &base in label {
6860 let code = hash_base_code(base);
6861 h1 = h1.wrapping_mul(HASH_BASE_1).wrapping_add(code);
6862 h2 = h2.wrapping_mul(HASH_BASE_2).wrapping_add(code);
6863 }
6864 ((h1 as u128) << 64) | h2 as u128
6865}
6866
6867fn hash_pow(base: u64, len: usize) -> u64 {
6868 let mut pow = 1u64;
6869 for _ in 0..len {
6870 pow = pow.wrapping_mul(base);
6871 }
6872 pow
6873}
6874
6875fn hash_base_code(base: u8) -> u64 {
6876 match base {
6877 b'A' => 1,
6878 b'C' => 2,
6879 b'G' => 3,
6880 b'T' => 4,
6881 _ => 5,
6882 }
6883}
6884
6885fn phi_edge_path_info<const K: usize>(
6886 edge: &DiscontinuityEdge<K>,
6887 vertex_info: PathInfo<K>,
6888) -> PathInfo<K> {
6889 let rank = if vertex_info.rank == 1 {
6890 0
6891 } else {
6892 vertex_info.rank
6893 };
6894 let exit_side = if rank == 0 {
6895 edge.unitig_exit_side
6896 } else {
6897 edge.unitig_exit_side.inverse()
6898 };
6899
6900 PathInfo {
6901 path_id: vertex_info.path_id,
6902 rank,
6903 exit_side,
6904 is_cycle: vertex_info.is_cycle,
6905 }
6906}
6907
6908#[inline(always)]
6909fn compact_phi_edge_path_info<const K: usize>(
6910 edge: &DiscontinuityEdge<K>,
6911 vertex_info: CompactExpansionPathInfo,
6912) -> CompactExpansionPathInfo {
6913 compact_phi_edge_path_info_with_exit(edge.unitig_exit_side, vertex_info)
6914}
6915
6916#[inline(always)]
6917fn compact_phi_edge_path_info_with_exit(
6918 unitig_exit_side: Side,
6919 vertex_info: CompactExpansionPathInfo,
6920) -> CompactExpansionPathInfo {
6921 let rank = if vertex_info.rank() == 1 {
6922 0
6923 } else {
6924 vertex_info.rank()
6925 };
6926 let exit_side = if rank == 0 {
6927 unitig_exit_side
6928 } else {
6929 unitig_exit_side.inverse()
6930 };
6931 CompactExpansionPathInfo::new(vertex_info.path_id, rank, exit_side, vertex_info.is_cycle())
6932}
6933
6934fn edge_path_info<const K: usize>(
6935 edge: &DiscontinuityEdge<K>,
6936 first_info: PathInfo<K>,
6937 second_info: PathInfo<K>,
6938) -> PathInfo<K> {
6939 let rank = first_info.rank.min(second_info.rank);
6940 let exit_side = if rank == first_info.rank {
6941 edge.unitig_exit_side
6942 } else {
6943 edge.unitig_exit_side.inverse()
6944 };
6945
6946 PathInfo {
6947 path_id: first_info.path_id,
6948 rank,
6949 exit_side,
6950 is_cycle: first_info.is_cycle,
6951 }
6952}
6953
6954#[inline(always)]
6955fn compact_edge_path_info<const K: usize>(
6956 edge: &DiscontinuityEdge<K>,
6957 first_info: CompactExpansionPathInfo,
6958 second_info: CompactExpansionPathInfo,
6959) -> CompactExpansionPathInfo {
6960 compact_edge_path_info_with_exit(edge.unitig_exit_side, first_info, second_info)
6961}
6962
6963#[inline(always)]
6964fn compact_edge_path_info_with_exit(
6965 unitig_exit_side: Side,
6966 first_info: CompactExpansionPathInfo,
6967 second_info: CompactExpansionPathInfo,
6968) -> CompactExpansionPathInfo {
6969 let rank = first_info.rank().min(second_info.rank());
6970 let exit_side = if rank == first_info.rank() {
6971 unitig_exit_side
6972 } else {
6973 unitig_exit_side.inverse()
6974 };
6975 CompactExpansionPathInfo::new(first_info.path_id, rank, exit_side, first_info.is_cycle())
6976}
6977
6978#[derive(Debug, Clone, Copy, PartialEq, Eq)]
6979enum LabelEnd {
6980 Left,
6981 Right,
6982}
6983
6984#[derive(Debug, Clone, Copy, PartialEq, Eq)]
6985struct HalfEnd {
6986 unitig_index: u32,
6987}
6988
6989impl HalfEnd {
6990 #[inline]
6991 fn unitig_index(self) -> usize {
6992 self.unitig_index as usize
6993 }
6994}
6995
6996#[inline]
6997fn label_end_for_node(node: usize) -> LabelEnd {
6998 if node & 1 == 0 {
6999 LabelEnd::Left
7000 } else {
7001 LabelEnd::Right
7002 }
7003}
7004
7005#[inline]
7006fn reverse_for_stitch_node(node: usize) -> bool {
7007 node & 1 == 1
7008}
7009
7010const STITCH_NO_NODE: u32 = u32::MAX;
7011
7012#[inline]
7013fn stitch_node(node: usize) -> u32 {
7014 u32::try_from(node).expect("stitch node index exceeds u32")
7015}
7016
7017#[inline]
7018fn stitch_node_index(node: u32) -> usize {
7019 node as usize
7020}
7021
7022#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
7023struct StitchAdjacency {
7024 to: usize,
7025 unitig_index: Option<usize>,
7026}
7027
7028#[derive(Debug, Clone, Copy, Default)]
7029struct StitchAdjacencyList {
7030 edges: [StitchAdjacency; 3],
7031 len: u8,
7032}
7033
7034impl StitchAdjacencyList {
7035 fn push(&mut self, edge: StitchAdjacency) {
7036 let len = self.len as usize;
7037 debug_assert!(len < self.edges.len());
7038 if len < self.edges.len() {
7039 self.edges[len] = edge;
7040 self.len += 1;
7041 }
7042 }
7043
7044 fn len(&self) -> usize {
7045 self.len as usize
7046 }
7047
7048 fn iter(&self) -> impl Iterator<Item = &StitchAdjacency> {
7049 self.edges[..self.len()].iter()
7050 }
7051}
7052
7053#[derive(Debug, Clone, Copy, Default)]
7054struct StitchVertexEnds {
7055 fronts: [u32; 2],
7056 backs: [u32; 2],
7057 nodes: [u32; 2],
7058 front_len: u8,
7059 back_len: u8,
7060 total_len: u8,
7061}
7062
7063impl StitchVertexEnds {
7064 fn push(&mut self, side: Side, node: u32) {
7065 if self.total_len < 2 {
7066 self.nodes[self.total_len as usize] = node;
7067 }
7068 self.total_len = self.total_len.saturating_add(1);
7069
7070 match side {
7071 Side::Front => {
7072 if self.front_len < 2 {
7073 self.fronts[self.front_len as usize] = node;
7074 }
7075 self.front_len = self.front_len.saturating_add(1);
7076 }
7077 Side::Back => {
7078 if self.back_len < 2 {
7079 self.backs[self.back_len as usize] = node;
7080 }
7081 self.back_len = self.back_len.saturating_add(1);
7082 }
7083 }
7084 }
7085}
7086
7087#[derive(Debug, Clone, Copy, PartialEq, Eq)]
7088struct StitchEndpointRecord<const K: usize> {
7089 vertex: Kmer<K>,
7090 side: Side,
7091 node: u32,
7092}
7093
7094#[derive(Debug, Clone, PartialEq, Eq)]
7095struct StitchEndpointBucketEntry {
7096 path: PathBuf,
7097 records: u64,
7098}
7099
7100const MAX_OPEN_STITCH_ENDPOINT_WRITERS: usize = 512;
7101
7102struct StitchEndpointBucketWriters {
7103 dir: PathBuf,
7104 bucket_mask: usize,
7105 writers: Vec<Option<StitchEndpointBucketWriter>>,
7106 records: Vec<u64>,
7107 open_writers: usize,
7108}
7109
7110impl StitchEndpointBucketWriters {
7111 fn create(dir: &Path, threads: usize) -> Result<Self, SerialCollationError> {
7112 if dir.exists() {
7113 fs::remove_dir_all(dir).map_err(|source| SerialCollationError::Io {
7114 path: dir.to_path_buf(),
7115 source,
7116 })?;
7117 }
7118 fs::create_dir_all(dir).map_err(|source| SerialCollationError::Io {
7119 path: dir.to_path_buf(),
7120 source,
7121 })?;
7122 let bucket_count = stitch_endpoint_bucket_count(threads);
7123 let mut writers = Vec::with_capacity(bucket_count);
7124 writers.resize_with(bucket_count, || None);
7125 Ok(Self {
7126 dir: dir.to_path_buf(),
7127 bucket_mask: bucket_count - 1,
7128 writers,
7129 records: vec![0; bucket_count],
7130 open_writers: 0,
7131 })
7132 }
7133
7134 fn write_record<const K: usize>(
7135 &mut self,
7136 record: &StitchEndpointRecord<K>,
7137 ) -> Result<(), SerialCollationError> {
7138 let bits = record.vertex.as_u128();
7139 let bucket_id = (hash_bytes(&bits.to_le_bytes(), 0) as usize) & self.bucket_mask;
7140 if self.writers[bucket_id].is_none() {
7141 self.evict_writer_if_needed(bucket_id)?;
7142 self.writers[bucket_id] = Some(StitchEndpointBucketWriter::open(
7143 &self.dir,
7144 bucket_id,
7145 self.records[bucket_id],
7146 )?);
7147 self.open_writers += 1;
7148 }
7149 let writer = self.writers[bucket_id]
7150 .as_mut()
7151 .expect("stitch endpoint bucket writer was just created");
7152 writer.write_record(bits, record.side, record.node)?;
7153 self.records[bucket_id] += 1;
7154 Ok(())
7155 }
7156
7157 fn evict_writer_if_needed(
7158 &mut self,
7159 requested_bucket_id: usize,
7160 ) -> Result<(), SerialCollationError> {
7161 if self.open_writers < MAX_OPEN_STITCH_ENDPOINT_WRITERS {
7162 return Ok(());
7163 }
7164 let evict_bucket_id = self
7165 .writers
7166 .iter()
7167 .enumerate()
7168 .find_map(|(bucket_id, writer)| {
7169 (bucket_id != requested_bucket_id && writer.is_some()).then_some(bucket_id)
7170 })
7171 .unwrap_or(requested_bucket_id);
7172 if let Some(mut writer) = self.writers[evict_bucket_id].take() {
7173 writer.flush()?;
7174 self.open_writers -= 1;
7175 }
7176 Ok(())
7177 }
7178
7179 fn finish(mut self) -> Result<Vec<StitchEndpointBucketEntry>, SerialCollationError> {
7180 let mut manifest = Vec::new();
7181 for writer in self.writers.iter_mut().flatten() {
7182 writer.flush()?;
7183 }
7184 for (bucket_id, &records) in self.records.iter().enumerate() {
7185 if records != 0 {
7186 manifest.push(StitchEndpointBucketEntry {
7187 path: self.dir.join(format!("{bucket_id:05}.seb")),
7188 records,
7189 });
7190 }
7191 }
7192 Ok(manifest)
7193 }
7194}
7195
7196struct StitchEndpointBucketWriter {
7197 path: PathBuf,
7198 out: BufWriter<File>,
7199 records: u64,
7200}
7201
7202impl StitchEndpointBucketWriter {
7203 fn open(dir: &Path, bucket_id: usize, records: u64) -> Result<Self, SerialCollationError> {
7204 let path = dir.join(format!("{bucket_id:05}.seb"));
7205 let file = OpenOptions::new()
7206 .create(true)
7207 .append(true)
7208 .open(&path)
7209 .map_err(|source| SerialCollationError::Io {
7210 path: path.clone(),
7211 source,
7212 })?;
7213 Ok(Self {
7214 path,
7215 out: BufWriter::with_capacity(STITCH_COORD_SHARD_WRITE_BUFFER, file),
7216 records,
7217 })
7218 }
7219
7220 fn write_record(
7221 &mut self,
7222 vertex_bits: u128,
7223 side: Side,
7224 node: u32,
7225 ) -> Result<(), SerialCollationError> {
7226 let side = match side {
7227 Side::Front => 0u8,
7228 Side::Back => 1u8,
7229 };
7230 let mut bytes = [0u8; 21];
7231 bytes[..16].copy_from_slice(&vertex_bits.to_le_bytes());
7232 bytes[16] = side;
7233 bytes[17..21].copy_from_slice(&node.to_le_bytes());
7234 self.out
7235 .write_all(&bytes)
7236 .map_err(|source| SerialCollationError::Io {
7237 path: self.path.clone(),
7238 source,
7239 })?;
7240 self.records += 1;
7241 Ok(())
7242 }
7243
7244 fn flush(&mut self) -> Result<(), SerialCollationError> {
7245 self.out.flush().map_err(|source| SerialCollationError::Io {
7246 path: self.path.clone(),
7247 source,
7248 })
7249 }
7250}
7251
7252#[derive(Debug, Clone, PartialEq, Eq)]
7253struct StitchedPath {
7254 label: Vec<u8>,
7255 is_cycle: bool,
7256}
7257
7258#[derive(Debug, Clone, Copy, PartialEq, Eq)]
7259struct StitchedCoordRecord {
7260 path_id: u64,
7261 rank: u64,
7262 unitig_index: u32,
7263 reverse: bool,
7264 is_cycle: bool,
7265}
7266
7267#[derive(Clone, Copy)]
7268struct DenseLocalPathInfo {
7269 path_id: u64,
7270 rank_and_flags: u64,
7271}
7272
7273impl DenseLocalPathInfo {
7274 const EMPTY: Self = Self {
7275 path_id: 0,
7276 rank_and_flags: u64::MAX,
7277 };
7278
7279 #[allow(dead_code)]
7281 fn from_record(record: StitchedCoordRecord) -> Self {
7282 Self {
7283 path_id: record.path_id,
7284 rank_and_flags: (record.rank << 2)
7285 | u64::from(record.reverse)
7286 | (u64::from(record.is_cycle) << 1),
7287 }
7288 }
7289
7290 fn to_record(self, unitig_index: usize) -> Option<StitchedCoordRecord> {
7291 (self.rank_and_flags != u64::MAX).then_some(StitchedCoordRecord {
7292 path_id: self.path_id,
7293 rank: self.rank_and_flags >> 2,
7294 unitig_index: unitig_index as u32,
7295 reverse: self.rank_and_flags & 1 != 0,
7296 is_cycle: self.rank_and_flags & 2 != 0,
7297 })
7298 }
7299}
7300
7301#[derive(Debug, Clone, PartialEq, Eq)]
7302struct StitchedCoordBucketEntry {
7303 bucket_id: usize,
7304 records: u64,
7305 path: PathBuf,
7306}
7307
7308#[derive(Debug, Clone, Copy, PartialEq, Eq)]
7309struct MaterializedStitchedCoordRecord {
7310 path_id: u64,
7311 rank: u64,
7312 label_offset: u64,
7313 label_len: u32,
7314 reverse: bool,
7315 is_cycle: bool,
7316 color_index: u32,
7317 color_count: u32,
7318}
7319
7320#[repr(C)]
7321#[derive(Debug, Clone, Copy, PartialEq, Eq)]
7322struct LoadedMaterializedStitchedCoordRecord {
7323 path_id: u64,
7324 label_offset: u32,
7325 color_start: u32,
7326 rank: u16,
7327 label_len: u16,
7328 color_count: u16,
7329 flags: u16,
7330}
7331
7332const _: () = assert!(std::mem::size_of::<LoadedMaterializedStitchedCoordRecord>() == 24);
7333
7334impl LoadedMaterializedStitchedCoordRecord {
7335 const REVERSE_FLAG: u16 = 1;
7336 const CYCLE_FLAG: u16 = 1 << 1;
7337
7338 #[inline(always)]
7339 fn reverse(self) -> bool {
7340 self.flags & Self::REVERSE_FLAG != 0
7341 }
7342
7343 #[inline(always)]
7344 fn is_cycle(self) -> bool {
7345 self.flags & Self::CYCLE_FLAG != 0
7346 }
7347
7348 #[inline(always)]
7349 fn color_count(self) -> u32 {
7350 u32::from(self.color_count)
7351 }
7352
7353 #[allow(clippy::too_many_arguments)]
7354 fn new(
7355 path_id: u64,
7356 rank: u64,
7357 label_offset: u32,
7358 label_len: u32,
7359 reverse: bool,
7360 is_cycle: bool,
7361 color_start: u32,
7362 color_count: u32,
7363 ) -> Self {
7364 Self {
7365 path_id,
7366 label_offset,
7367 color_start,
7368 rank: u16::try_from(rank).expect("materialized rank fits C++ weight_t"),
7369 label_len: u16::try_from(label_len)
7370 .expect("materialized label length fits C++ uni_len_t"),
7371 color_count: u16::try_from(color_count)
7372 .expect("materialized color count fits C++ uni_len_t"),
7373 flags: (u16::from(reverse) * Self::REVERSE_FLAG)
7374 | (u16::from(is_cycle) * Self::CYCLE_FLAG),
7375 }
7376 }
7377}
7378
7379#[derive(Debug, Clone, PartialEq, Eq)]
7380struct MaterializedStitchedCoordBucketEntry {
7381 bucket_id: usize,
7382 records: u64,
7383 label_bytes: u64,
7384 coord_path: PathBuf,
7385 label_path: PathBuf,
7386 color_path: Option<PathBuf>,
7387 color_runs: u64,
7388}
7389
7390const STITCH_COORD_MAGIC: &[u8; 8] = b"CF3SCB2\0";
7391const MATERIALIZED_STITCH_COORD_MAGIC: &[u8; 8] = b"CF3MCB2\0";
7392const STITCH_COORD_HEADER_LEN: u64 = 32;
7393const STITCH_PATH_INFO_RECORD_LEN: u64 = 24;
7394const STITCH_COORD_RECORD_LEN: u64 = 24;
7395const MATERIALIZED_STITCH_COORD_SHARD_WRITE_BUFFER: usize = 16 * 1024;
7396const FINAL_UNITIG_BUCKET_WRITE_BUFFER: usize = 128 * 1024;
7397const STITCH_COORD_SHARD_WRITE_BUFFER: usize = 1024 * 1024;
7398const STITCH_COORD_RECORD_WRITE_BUFFER: usize = 1024 * 1024;
7399const EDGE_PATH_INFO_WORKER_BUFFER: usize = 128 * 1024;
7400const STITCH_COORD_REVERSE_FLAG: u8 = 1;
7401const STITCH_COORD_CYCLE_FLAG: u8 = 2;
7402const MAX_OPEN_STITCH_PATH_INFO_WRITERS: usize = 768;
7403const MAX_OPEN_MATERIALIZED_STITCH_WRITERS: usize = 768;
7404const DEFAULT_MAX_UNITIG_COORD_BUCKETS: usize = 1024;
7406const DEFAULT_EXPECTED_COLOR_CEILING: u64 = 48 * 1024 * 1024;
7408const HIGH_THREAD_COORD_BUCKET_THRESHOLD: usize = 128;
7410const HIGH_THREAD_MAX_UNITIG_COORD_BUCKETS: usize = 256;
7412const MAX_NARROW_COORD_BUCKET_BASES: u64 = 64 * 1024 * 1024;
7419const MAX_OPEN_MATERIALIZED_STITCH_WRITERS_PER_SHARD: usize = 8;
7420const STITCH_PATH_INFO_BUCKET_TARGET: usize = 32;
7421const DEFAULT_INMEM_STITCH_PATH_INFO_LIMIT: usize = 768 * 1024 * 1024;
7422fn stitch_discontinuity_paths<const K: usize>(
7423 inputs: &DiscontinuityInputs<K>,
7424 skip_unitigs: &[bool],
7425 threads: usize,
7426) -> Vec<Vec<u8>> {
7427 stitch_discontinuity_paths_impl::<K>(inputs, skip_unitigs, threads, None)
7428 .expect("in-memory discontinuity stitching cannot fail")
7429}
7430
7431fn stitch_discontinuity_paths_with_coord_dir<const K: usize>(
7432 inputs: &DiscontinuityInputs<K>,
7433 skip_unitigs: &[bool],
7434 threads: usize,
7435 coord_dir: &Path,
7436) -> Result<Vec<Vec<u8>>, SerialCollationError> {
7437 stitch_discontinuity_paths_impl::<K>(inputs, skip_unitigs, threads, Some(coord_dir))
7438}
7439
7440fn stitch_discontinuity_paths_to_final_buckets<const K: usize>(
7441 inputs: &DiscontinuityInputs<K>,
7442 skip_unitigs: &[bool],
7443 threads: usize,
7444 coord_dir: &Path,
7445 final_buckets: &mut FinalUnitigBucketWriters,
7446) -> Result<u64, SerialCollationError> {
7447 let debug_stitch = false;
7448 let endpoint_dir = coord_dir.join("endpoints");
7449 let mut endpoint_writers = StitchEndpointBucketWriters::create(&endpoint_dir, threads)?;
7450 let build_started = Instant::now();
7451 let mut half_ends = Vec::<HalfEnd>::new();
7452
7453 for (unitig_index, unitig) in inputs.unitigs.iter().enumerate() {
7454 if skip_unitigs.get(unitig_index).copied().unwrap_or(false) {
7455 continue;
7456 }
7457 if unitig.left_exit().is_none() && unitig.right_exit().is_none() {
7458 continue;
7459 }
7460
7461 let (left_endpoint, right_endpoint) = endpoints_by_label_end(unitig);
7462 let left_node = half_ends.len();
7463 half_ends.push(HalfEnd {
7464 unitig_index: u32::try_from(unitig_index).expect("unitig index exceeds u32"),
7465 });
7466 if let Some(endpoint) = left_endpoint {
7467 endpoint_writers.write_record(&StitchEndpointRecord {
7468 vertex: endpoint.vertex,
7469 side: endpoint.side,
7470 node: stitch_node(left_node),
7471 })?;
7472 }
7473
7474 let right_node = half_ends.len();
7475 half_ends.push(HalfEnd {
7476 unitig_index: u32::try_from(unitig_index).expect("unitig index exceeds u32"),
7477 });
7478 if let Some(endpoint) = right_endpoint {
7479 endpoint_writers.write_record(&StitchEndpointRecord {
7480 vertex: endpoint.vertex,
7481 side: endpoint.side,
7482 node: stitch_node(right_node),
7483 })?;
7484 }
7485 }
7486 let half_end_elapsed = build_started.elapsed();
7487
7488 let mut join_neighbor = vec![STITCH_NO_NODE; half_ends.len()];
7489 let endpoint_sort_started = Instant::now();
7490 let endpoint_manifest = endpoint_writers.finish()?;
7491 let endpoint_sort_elapsed = endpoint_sort_started.elapsed();
7492 let endpoint_join_started = Instant::now();
7493 join_neighbors_from_endpoint_buckets::<K>(&endpoint_manifest, &mut join_neighbor, threads)?;
7494 let endpoint_join_elapsed = endpoint_join_started.elapsed();
7495
7496 if debug_stitch {
7497 eprintln!("stitch inputs: {} unitigs", inputs.unitigs.len());
7498 for (node, half_end) in half_ends.iter().enumerate() {
7499 let (left_endpoint, right_endpoint) =
7500 endpoints_by_label_end(&inputs.unitigs[half_end.unitig_index()]);
7501 let label_end = label_end_for_node(node);
7502 let endpoint = match label_end {
7503 LabelEnd::Left => left_endpoint,
7504 LabelEnd::Right => right_endpoint,
7505 };
7506 eprintln!(
7507 "node {node}: unitig={} end={:?} endpoint={:?} join={}",
7508 half_end.unitig_index(),
7509 label_end,
7510 endpoint.map(debug_endpoint),
7511 join_neighbor[node],
7512 );
7513 }
7514 }
7515
7516 let component_started = Instant::now();
7517 let component_starts = stitch_component_starts_from_neighbors(&half_ends, &join_neighbor);
7518 let component_elapsed = component_started.elapsed();
7519 let walk_started = Instant::now();
7520 let manifest = write_materialized_stitched_coord_buckets(
7521 inputs,
7522 coord_dir,
7523 threads,
7524 &half_ends,
7525 &join_neighbor,
7526 &component_starts,
7527 )?;
7528 let emitted = reduce_materialized_stitched_coord_bucket_files_to_final::<K>(
7529 &manifest,
7530 &[],
7531 threads,
7532 final_buckets,
7533 )?;
7534 let walk_elapsed = walk_started.elapsed();
7535
7536 eprintln!(
7537 "cuttlefish: stitch detail: half-ends {:.3}s, endpoint sort {:.3}s, endpoint join {:.3}s, components {:.3}s, materialized walk+reduce {:.3}s",
7538 half_end_elapsed.as_secs_f64(),
7539 endpoint_sort_elapsed.as_secs_f64(),
7540 endpoint_join_elapsed.as_secs_f64(),
7541 component_elapsed.as_secs_f64(),
7542 walk_elapsed.as_secs_f64(),
7543 );
7544 Ok(emitted)
7545}
7546
7547fn stitch_external_discontinuity_paths_to_final_buckets<const K: usize>(
7548 inputs: &ExternalDiscontinuityInputs<K>,
7549 threads: usize,
7550 coord_dir: &Path,
7551 final_buckets: &mut FinalUnitigBucketWriters,
7552) -> Result<u64, SerialCollationError> {
7553 report_process_memory("external stitch start");
7554 let endpoint_dir = coord_dir.join("endpoints");
7555 let mut endpoint_writers = StitchEndpointBucketWriters::create(&endpoint_dir, threads)?;
7556 let reader = ExternalDiscontinuityReader::open(inputs)?;
7557 let build_started = Instant::now();
7558 let mut half_ends = Vec::<HalfEnd>::new();
7559 let use_direct_label_reads = std::env::var_os("CF3_RS_DIRECT_STITCH_LABEL_READS").is_some();
7560 let use_ordered_labels = std::env::var_os("CF3_RS_ORDERED_STITCH_LABELS").is_some();
7561 let collect_label_refs = use_direct_label_reads || use_ordered_labels;
7562 let mut label_refs = if collect_label_refs {
7563 Vec::<ExternalLabelRef>::with_capacity(inputs.unitig_count())
7564 } else {
7565 Vec::<ExternalLabelRef>::new()
7566 };
7567
7568 for (unitig_index, unitig) in reader.iter()?.enumerate() {
7569 let unitig = unitig?;
7570 if collect_label_refs {
7571 label_refs.push(ExternalLabelRef {
7572 label_start: unitig.label_start,
7573 label_len: unitig.label_len,
7574 });
7575 }
7576 if unitig.left_exit().is_none() && unitig.right_exit().is_none() {
7577 continue;
7578 }
7579
7580 let (left_endpoint, right_endpoint) = endpoints_by_label_end(&unitig);
7581 let left_node = half_ends.len();
7582 half_ends.push(HalfEnd {
7583 unitig_index: u32::try_from(unitig_index).expect("unitig index exceeds u32"),
7584 });
7585 if let Some(endpoint) = left_endpoint {
7586 endpoint_writers.write_record(&StitchEndpointRecord {
7587 vertex: endpoint.vertex,
7588 side: endpoint.side,
7589 node: stitch_node(left_node),
7590 })?;
7591 }
7592
7593 let right_node = half_ends.len();
7594 half_ends.push(HalfEnd {
7595 unitig_index: u32::try_from(unitig_index).expect("unitig index exceeds u32"),
7596 });
7597 if let Some(endpoint) = right_endpoint {
7598 endpoint_writers.write_record(&StitchEndpointRecord {
7599 vertex: endpoint.vertex,
7600 side: endpoint.side,
7601 node: stitch_node(right_node),
7602 })?;
7603 }
7604 }
7605 let half_end_elapsed = build_started.elapsed();
7606 report_process_memory("external stitch after half-ends");
7607
7608 let mut join_neighbor = vec![STITCH_NO_NODE; half_ends.len()];
7609 let endpoint_sort_started = Instant::now();
7610 let endpoint_manifest = endpoint_writers.finish()?;
7611 let endpoint_sort_elapsed = endpoint_sort_started.elapsed();
7612 let endpoint_join_started = Instant::now();
7613 join_neighbors_from_endpoint_buckets::<K>(&endpoint_manifest, &mut join_neighbor, threads)?;
7614 let endpoint_join_elapsed = endpoint_join_started.elapsed();
7615 report_process_memory("external stitch after endpoint join");
7616
7617 let walk_started = Instant::now();
7618 let component_elapsed;
7619 let emitted = if use_direct_label_reads {
7620 let component_started = Instant::now();
7621 let component_starts = stitch_component_starts_from_neighbors(&half_ends, &join_neighbor);
7622 component_elapsed = component_started.elapsed();
7623 emit_external_stitched_labels_to_final_buckets::<K>(
7624 inputs,
7625 threads,
7626 &label_refs,
7627 &half_ends,
7628 &join_neighbor,
7629 &component_starts,
7630 final_buckets,
7631 )?
7632 } else if use_ordered_labels {
7633 let component_started = Instant::now();
7634 let component_starts = stitch_component_starts_from_neighbors(&half_ends, &join_neighbor);
7635 component_elapsed = component_started.elapsed();
7636 emit_external_ordered_stitched_labels_to_final_buckets::<K>(
7637 inputs,
7638 coord_dir,
7639 &label_refs,
7640 &half_ends,
7641 &join_neighbor,
7642 &component_starts,
7643 final_buckets,
7644 )?
7645 } else {
7646 let manifest =
7647 if std::env::var_os("CF3_RS_NEIGHBOR_STITCH_PATH_INFO").is_some() {
7648 let component_started = Instant::now();
7649 let manifest = write_external_materialized_stitched_coord_buckets_from_neighbors::<
7650 K,
7651 >(
7652 inputs, coord_dir, threads, &half_ends, &join_neighbor
7653 )?;
7654 component_elapsed = component_started.elapsed();
7655 report_process_memory("external stitch after materialized coord write");
7656 manifest
7657 } else {
7658 let component_started = Instant::now();
7659 let component_starts =
7660 stitch_component_starts_from_neighbors(&half_ends, &join_neighbor);
7661 component_elapsed = component_started.elapsed();
7662 write_external_materialized_stitched_coord_buckets::<K>(
7663 inputs,
7664 coord_dir,
7665 threads,
7666 &half_ends,
7667 &join_neighbor,
7668 &component_starts,
7669 )?
7670 };
7671 report_process_memory("external stitch before materialized coord reduce");
7672 let reduce_started = Instant::now();
7673 let emitted = reduce_materialized_stitched_coord_bucket_files_to_final::<K>(
7674 &manifest,
7675 &[],
7676 threads,
7677 final_buckets,
7678 )?;
7679 report_process_memory("external stitch after materialized coord reduce");
7680 eprintln!(
7681 "cuttlefish: external materialized stitch reduce {:.3}s",
7682 reduce_started.elapsed().as_secs_f64()
7683 );
7684 emitted
7685 };
7686 let walk_elapsed = walk_started.elapsed();
7687
7688 eprintln!(
7689 "cuttlefish: external stitch detail: half-ends {:.3}s, endpoint sort {:.3}s, endpoint join {:.3}s, components {:.3}s, stitch label materialization/reduce {:.3}s",
7690 half_end_elapsed.as_secs_f64(),
7691 endpoint_sort_elapsed.as_secs_f64(),
7692 endpoint_join_elapsed.as_secs_f64(),
7693 component_elapsed.as_secs_f64(),
7694 walk_elapsed.as_secs_f64(),
7695 );
7696 Ok(emitted)
7697}
7698
7699struct ExternalCppPathInfoCollation {
7700 stitched_unitigs: u64,
7701 direct_local_unitigs: u64,
7702 direct_local_unitigs_complete: bool,
7703}
7704
7705const BACKGROUND_UNLINK_WORKERS: usize = 8;
7712
7713pub(crate) fn spawn_background_dir_removal(dir: PathBuf) -> std::thread::JoinHandle<()> {
7721 std::thread::spawn(move || {
7722 let Ok(entries) = fs::read_dir(&dir) else {
7723 let _ = fs::remove_dir_all(&dir);
7724 return;
7725 };
7726 let paths: Vec<PathBuf> = entries.flatten().map(|entry| entry.path()).collect();
7727 let next = AtomicUsize::new(0);
7728 let workers = paths.len().clamp(1, BACKGROUND_UNLINK_WORKERS);
7729 std::thread::scope(|scope| {
7730 for _ in 0..workers {
7731 scope.spawn(|| {
7732 loop {
7733 let index = next.fetch_add(1, Ordering::Relaxed);
7734 let Some(path) = paths.get(index) else {
7735 break;
7736 };
7737 if fs::remove_file(path).is_err() {
7738 let _ = fs::remove_dir_all(path);
7739 }
7740 }
7741 });
7742 }
7743 });
7744 let _ = fs::remove_dir_all(&dir);
7745 })
7746}
7747
7748fn collate_external_cpp_path_info_to_final_buckets<const K: usize>(
7749 inputs: &mut ExternalDiscontinuityInputs<K>,
7750 threads: usize,
7751 coord_dir: &Path,
7752 final_buckets: &mut FinalUnitigBucketWriters,
7753) -> Result<ExternalCppPathInfoCollation, SerialCollationError> {
7754 eprintln!("cuttlefish: deriving discontinuity path-info by C++-style contraction/expansion");
7755
7756 let matrix_started = Instant::now();
7757 let matrix = inputs
7758 .edge_matrix
7759 .take()
7760 .ok_or_else(|| SerialCollationError::MalformedCoordBucket(inputs.unitig_path.clone()))?;
7761 let matrix_dir = matrix.dir.clone();
7762 eprintln!(
7763 "cuttlefish: discontinuity edge matrix handed off in {:.3}s; {} edge(s), {} phi edge(s), {} diagonal edge(s)",
7764 matrix_started.elapsed().as_secs_f64(),
7765 matrix.stats.edges,
7766 matrix.stats.phi_edges,
7767 matrix.stats.diagonal_edges
7768 );
7769
7770 let contract_started = Instant::now();
7771 report_process_memory("before C++-style discontinuity contraction");
7772 let mut contraction =
7773 SerialDiscontinuityContractor::contract_blocked_external(matrix, threads)?;
7774 report_process_memory("after C++-style discontinuity contraction");
7775 eprintln!(
7776 "cuttlefish: discontinuity graph contracted in {:.3}s; {} meta-vertices, {} final edge(s), {} reinserted edge(s)",
7777 contract_started.elapsed().as_secs_f64(),
7778 contraction.stats.meta_vertices,
7779 contraction.stats.final_edges,
7780 contraction.stats.reinserted_edges
7781 );
7782 let (ranges_per_bucket, path_info_bucket_count) = if let Some(buckets) =
7783 inputs.local_unitig_buckets.as_ref()
7784 {
7785 (1, buckets.len().max(1))
7786 } else {
7787 let ranges_per_bucket = ranges_per_path_info_bucket(inputs.ranges.len(), threads.max(1));
7788 (
7789 ranges_per_bucket,
7790 inputs.ranges.len().div_ceil(ranges_per_bucket).max(1),
7791 )
7792 };
7793 let expansion_dir = coord_dir.with_file_name(format!(
7794 "{}.cpp-expansion",
7795 coord_dir
7796 .file_name()
7797 .and_then(|name| name.to_str())
7798 .unwrap_or("coord")
7799 ));
7800 let expand_started = Instant::now();
7801 report_process_memory("before C++-style discontinuity expansion");
7802 let expansion = SerialDiscontinuityExpander::expand_cpp_ordered_external_to_range_buckets(
7803 &mut contraction,
7804 &inputs.ranges,
7805 ranges_per_bucket,
7806 path_info_bucket_count,
7807 &inputs.unitig_path,
7808 &expansion_dir,
7809 threads,
7810 )?;
7811 report_process_memory("after C++-style discontinuity expansion");
7812 eprintln!(
7813 "cuttlefish: contracted graph expanded in {:.3}s; {} path-info edge record(s), {} unresolved edge(s)",
7814 expand_started.elapsed().as_secs_f64(),
7815 expansion.stats.edge_path_infos,
7816 expansion.stats.unresolved_edges
7817 );
7818 drop(contraction);
7819 let matrix_reclaim = spawn_background_dir_removal(matrix_dir);
7820 trim_process_allocations();
7821 report_process_memory("after discontinuity contraction release");
7822
7823 let map_started = Instant::now();
7824 report_process_memory("before C++-style path-info map");
7825 let materialized = map_external_cpp_path_info_buckets_to_max_unitig_buckets::<K>(
7826 inputs,
7827 coord_dir,
7828 threads,
7829 ranges_per_bucket,
7830 expansion,
7831 final_buckets,
7832 )?;
7833 report_process_memory("after C++-style path-info map");
7834 let expansion_reclaim = spawn_background_dir_removal(expansion_dir);
7835 let local_unitig_reclaim = inputs
7837 .local_unitig_bucket_dir
7838 .take()
7839 .filter(|_| !keep_intermediates())
7840 .map(spawn_background_dir_removal);
7841 eprintln!(
7842 "cuttlefish: C++-style path-info map completed in {:.3}s",
7843 map_started.elapsed().as_secs_f64()
7844 );
7845 let reduce_started = Instant::now();
7846 report_process_memory("before C++-style path-info reduce");
7847 let emitted = reduce_materialized_stitched_coord_bucket_files_to_final::<K>(
7848 &materialized.manifest,
7849 &materialized.retained,
7850 threads,
7851 final_buckets,
7852 )?;
7853 report_process_memory("after C++-style path-info reduce");
7854 eprintln!(
7855 "cuttlefish: C++-style path-info reduce {:.3}s",
7856 reduce_started.elapsed().as_secs_f64()
7857 );
7858 let reclaim_started = Instant::now();
7859 let _ = matrix_reclaim.join();
7860 let _ = expansion_reclaim.join();
7861 if let Some(handle) = local_unitig_reclaim {
7862 let _ = handle.join();
7863 }
7864 eprintln!(
7865 "cuttlefish: waited {:.3}s for background intermediate removal",
7866 reclaim_started.elapsed().as_secs_f64()
7867 );
7868 Ok(ExternalCppPathInfoCollation {
7869 stitched_unitigs: emitted,
7870 direct_local_unitigs: materialized.direct_local_unitigs,
7871 direct_local_unitigs_complete: materialized.direct_local_unitigs_complete,
7872 })
7873}
7874
7875fn prepare_unitig_blocked_edge<const K: usize>(
7876 unitig: &DiscontinuityUnitig<K>,
7877 unitig_index: usize,
7878 vertex_partitions: usize,
7879) -> Option<PreparedBlockedEdge> {
7880 let mut edge = match (unitig.left_exit(), unitig.right_exit()) {
7881 (Some(left), Some(right)) => DiscontinuityEdge {
7882 first: MatrixEndpoint::Vertex(left),
7883 second: MatrixEndpoint::Vertex(right),
7884 weight: 1,
7885 unitig_bucket: 0,
7886 unitig_index,
7887 unitig_exit_side: Side::Back,
7888 phantom_unitig: None,
7889 swapped: false,
7890 },
7891 (Some(endpoint), None) => DiscontinuityEdge {
7892 first: MatrixEndpoint::Phi,
7893 second: MatrixEndpoint::Vertex(endpoint),
7894 weight: 1,
7895 unitig_bucket: 0,
7896 unitig_index,
7897 unitig_exit_side: Side::Front,
7898 phantom_unitig: None,
7899 swapped: false,
7900 },
7901 (None, Some(endpoint)) => DiscontinuityEdge {
7902 first: MatrixEndpoint::Phi,
7903 second: MatrixEndpoint::Vertex(endpoint),
7904 weight: 1,
7905 unitig_bucket: 0,
7906 unitig_index,
7907 unitig_exit_side: Side::Back,
7908 phantom_unitig: None,
7909 swapped: false,
7910 },
7911 (None, None) => return None,
7912 };
7913 let first_partition = edge_matrix_partition(vertex_partitions, edge.first);
7914 let second_partition = edge_matrix_partition(vertex_partitions, edge.second);
7915 if first_partition > second_partition {
7916 std::mem::swap(&mut edge.first, &mut edge.second);
7917 edge.unitig_exit_side = edge.unitig_exit_side.inverse();
7918 edge.swapped = true;
7919 }
7920 let row = first_partition.min(second_partition);
7921 let col = first_partition.max(second_partition);
7922 Some(PreparedBlockedEdge {
7923 block: row * (vertex_partitions + 1) + col,
7924 bytes: encode_discontinuity_edge(&edge),
7925 phi: row == 0,
7926 diagonal: row == col,
7927 })
7928}
7929
7930fn prepare_existing_blocked_edge<const K: usize>(
7931 mut edge: DiscontinuityEdge<K>,
7932 vertex_partitions: usize,
7933) -> PreparedBlockedEdge {
7934 let first_partition = edge_matrix_partition(vertex_partitions, edge.first);
7935 let second_partition = edge_matrix_partition(vertex_partitions, edge.second);
7936 if first_partition > second_partition {
7937 std::mem::swap(&mut edge.first, &mut edge.second);
7938 edge.unitig_exit_side = edge.unitig_exit_side.inverse();
7939 edge.swapped = !edge.swapped;
7940 }
7941 let row = first_partition.min(second_partition);
7942 let col = first_partition.max(second_partition);
7943 PreparedBlockedEdge {
7944 block: row * (vertex_partitions + 1) + col,
7945 bytes: encode_discontinuity_edge(&edge),
7946 phi: row == 0,
7947 diagonal: row == col,
7948 }
7949}
7950
7951fn emit_contracted_edge_chunks<const K: usize>(
7952 edges: Vec<DiscontinuityEdge<K>>,
7953 vertex_partitions: usize,
7954 max_partition_exclusive: Option<usize>,
7955 appenders: &ConcurrentBlockedEdgeWriters,
7956) -> Result<u64, SerialCollationError> {
7957 const EDGES_PER_BATCH: usize = 8 * 1024;
7958 edges
7959 .par_chunks(EDGES_PER_BATCH)
7960 .map(|chunk| {
7961 let mut prepared = Vec::with_capacity(chunk.len());
7962 for edge in chunk {
7963 if let Some(limit) = max_partition_exclusive {
7964 let first = edge_matrix_partition(vertex_partitions, edge.first);
7965 let second = edge_matrix_partition(vertex_partitions, edge.second);
7966 if first.max(second) >= limit {
7967 continue;
7968 }
7969 }
7970 prepared.push(prepare_existing_blocked_edge(
7971 edge.clone(),
7972 vertex_partitions,
7973 ));
7974 }
7975 prepared.sort_unstable_by_key(|edge| edge.block);
7976 let mut start = 0;
7977 while start < prepared.len() {
7978 let block = prepared[start].block;
7979 let mut end = start + 1;
7980 while end < prepared.len() && prepared[end].block == block {
7981 end += 1;
7982 }
7983 appenders.add_batch(&prepared[start..end])?;
7984 start = end;
7985 }
7986 Ok::<_, SerialCollationError>(prepared.len() as u64)
7987 })
7988 .try_reduce(|| 0, |left, right| Ok(left + right))
7989}
7990
7991#[allow(dead_code)]
7993fn emit_contracted_edge_batch<const K: usize>(
7994 edges: Vec<DiscontinuityEdge<K>>,
7995 vertex_partitions: usize,
7996 appenders: &ConcurrentBlockedEdgeWriters,
7997) -> Result<u64, SerialCollationError> {
7998 let mut prepared = edges
7999 .into_iter()
8000 .map(|edge| prepare_existing_blocked_edge(edge, vertex_partitions))
8001 .collect::<Vec<_>>();
8002 prepared.sort_unstable_by_key(|edge| edge.block);
8003 let mut start = 0;
8004 while start < prepared.len() {
8005 let block = prepared[start].block;
8006 let mut end = start + 1;
8007 while end < prepared.len() && prepared[end].block == block {
8008 end += 1;
8009 }
8010 appenders.add_batch(&prepared[start..end])?;
8011 start = end;
8012 }
8013 Ok(prepared.len() as u64)
8014}
8015
8016fn emit_prepared_edge_batch(
8017 prepared: &mut [PreparedBlockedEdge],
8018 appenders: &ConcurrentBlockedEdgeWriters,
8019) -> Result<u64, SerialCollationError> {
8020 prepared.sort_unstable_by_key(|edge| edge.block);
8021 let mut start = 0;
8022 while start < prepared.len() {
8023 let block = prepared[start].block;
8024 let mut end = start + 1;
8025 while end < prepared.len() && prepared[end].block == block {
8026 end += 1;
8027 }
8028 appenders.add_batch(&prepared[start..end])?;
8029 start = end;
8030 }
8031 Ok(prepared.len() as u64)
8032}
8033
8034fn serial_collation_to_input_error(error: SerialCollationError) -> DiscontinuityInputError {
8035 match error {
8036 SerialCollationError::Io { path, source } => DiscontinuityInputError::Io { path, source },
8037 SerialCollationError::MalformedCoordBucket(path) => DiscontinuityInputError::Io {
8038 path,
8039 source: std::io::Error::new(
8040 std::io::ErrorKind::InvalidData,
8041 "malformed blocked edge matrix",
8042 ),
8043 },
8044 SerialCollationError::WorkerPanic => DiscontinuityInputError::WorkerPanic,
8045 SerialCollationError::Color(err) => DiscontinuityInputError::Color(err),
8046 }
8047}
8048
8049struct ExternalCppPathInfoMaterialized {
8050 manifest: Vec<MaterializedStitchedCoordBucketEntry>,
8051 retained: Vec<Vec<PendingMaterializedBucket>>,
8052 direct_local_unitigs: u64,
8053 direct_local_unitigs_complete: bool,
8054}
8055
8056struct RangeBucketedExpansion<const K: usize> {
8057 stats: SerialExpansionStats,
8058 records: Vec<Vec<StitchedCoordRecord>>,
8059 record_manifest: Vec<StitchedCoordBucketEntry>,
8060 phantom_records: Vec<(StitchedCoordRecord, DiscontinuityEndpoint<K>)>,
8061}
8062
8063const VERTEX_PATH_INFO_WRITE_BUFFER: usize = 1024 * 1024;
8064
8065#[inline]
8066const fn vertex_path_info_record_len<const K: usize>() -> usize {
8067 if K <= 31 {
8068 std::mem::size_of::<CompactVertexPathInfoRecord>()
8069 } else {
8070 2 * discontinuity_edge_kmer_bytes::<K>() + 9
8071 }
8072}
8073
8074#[repr(C)]
8075#[derive(Clone, Copy)]
8076struct CompactVertexPathInfoRecord {
8077 vertex: u64,
8078 path_id: u64,
8079 rank_and_flags: u64,
8080}
8081
8082const _: () = assert!(std::mem::size_of::<CompactVertexPathInfoRecord>() == 24);
8083
8084struct VertexPathInfoBucketWriters<const K: usize> {
8085 dir: PathBuf,
8086 writers: Vec<Option<BufWriter<File>>>,
8087 buffers: Vec<Vec<u8>>,
8088 phantom: PhantomData<[(); K]>,
8089}
8090
8091impl<const K: usize> VertexPathInfoBucketWriters<K> {
8092 #[allow(dead_code)]
8094 fn create(dir: &Path, bucket_count: usize) -> Result<Self, SerialCollationError> {
8095 if dir.exists() {
8096 fs::remove_dir_all(dir).map_err(|source| SerialCollationError::Io {
8097 path: dir.to_path_buf(),
8098 source,
8099 })?;
8100 }
8101 fs::create_dir_all(dir).map_err(|source| SerialCollationError::Io {
8102 path: dir.to_path_buf(),
8103 source,
8104 })?;
8105 let mut writers = Vec::with_capacity(bucket_count);
8106 writers.resize_with(bucket_count, || None);
8107 let mut buffers = Vec::with_capacity(bucket_count);
8108 buffers.resize_with(bucket_count, Vec::new);
8109 Ok(Self {
8110 dir: dir.to_path_buf(),
8111 writers,
8112 buffers,
8113 phantom: PhantomData,
8114 })
8115 }
8116
8117 fn open_existing(dir: &Path, bucket_count: usize) -> Self {
8118 let mut writers = Vec::with_capacity(bucket_count);
8119 writers.resize_with(bucket_count, || None);
8120 let mut buffers = Vec::with_capacity(bucket_count);
8121 buffers.resize_with(bucket_count, Vec::new);
8122 Self {
8123 dir: dir.to_path_buf(),
8124 writers,
8125 buffers,
8126 phantom: PhantomData,
8127 }
8128 }
8129
8130 fn write_record(
8131 &mut self,
8132 bucket_id: usize,
8133 record: &VertexPathInfo<K>,
8134 ) -> Result<(), SerialCollationError> {
8135 if bucket_id >= self.writers.len() {
8136 return Err(SerialCollationError::MalformedCoordBucket(self.dir.clone()));
8137 }
8138 self.buffers[bucket_id].extend_from_slice(
8139 &encoded_vertex_path_info_record::<K>(record)[..vertex_path_info_record_len::<K>()],
8140 );
8141 if self.buffers[bucket_id].len() >= VERTEX_PATH_INFO_WRITE_BUFFER {
8142 self.write_buffer(bucket_id, false)?;
8143 }
8144 Ok(())
8145 }
8146
8147 fn flush_bucket(&mut self, bucket_id: usize) -> Result<(), SerialCollationError> {
8148 self.write_buffer(bucket_id, true)
8149 }
8150
8151 fn write_buffer(
8152 &mut self,
8153 bucket_id: usize,
8154 flush_writer: bool,
8155 ) -> Result<(), SerialCollationError> {
8156 if bucket_id >= self.writers.len() {
8157 return Err(SerialCollationError::MalformedCoordBucket(self.dir.clone()));
8158 }
8159 if self.buffers[bucket_id].is_empty() {
8160 if flush_writer && let Some(writer) = self.writers[bucket_id].as_mut() {
8161 writer.flush().map_err(|source| SerialCollationError::Io {
8162 path: vertex_path_info_bucket_path(&self.dir, bucket_id),
8163 source,
8164 })?;
8165 }
8166 return Ok(());
8167 }
8168 if self.writers[bucket_id].is_none() {
8169 let path = vertex_path_info_bucket_path(&self.dir, bucket_id);
8170 let file = OpenOptions::new()
8171 .create(true)
8172 .append(true)
8173 .open(&path)
8174 .map_err(|source| SerialCollationError::Io {
8175 path: path.clone(),
8176 source,
8177 })?;
8178 self.writers[bucket_id] = Some(BufWriter::with_capacity(
8179 VERTEX_PATH_INFO_WRITE_BUFFER,
8180 file,
8181 ));
8182 }
8183 let writer = self.writers[bucket_id]
8184 .as_mut()
8185 .expect("vertex path-info writer was just created");
8186 writer
8187 .write_all(&self.buffers[bucket_id])
8188 .map_err(|source| SerialCollationError::Io {
8189 path: vertex_path_info_bucket_path(&self.dir, bucket_id),
8190 source,
8191 })?;
8192 if flush_writer {
8193 writer.flush().map_err(|source| SerialCollationError::Io {
8194 path: vertex_path_info_bucket_path(&self.dir, bucket_id),
8195 source,
8196 })?;
8197 }
8198 self.buffers[bucket_id].clear();
8199 Ok(())
8200 }
8201}
8202
8203fn vertex_path_info_bucket_path(dir: &Path, bucket_id: usize) -> PathBuf {
8204 dir.join(format!("{bucket_id:05}.pv"))
8205}
8206
8207fn encoded_vertex_path_info_record<const K: usize>(record: &VertexPathInfo<K>) -> [u8; 41] {
8208 if K <= 31 {
8209 let mut bytes = [0u8; 41];
8210 bytes[..8].copy_from_slice(&(record.vertex.as_u128() as u64).to_le_bytes());
8211 bytes[8..16].copy_from_slice(&(record.info.path_id.as_u128() as u64).to_le_bytes());
8212 let flags =
8213 u64::from(record.info.exit_side == Side::Back) | (u64::from(record.info.is_cycle) << 1);
8214 bytes[16..24].copy_from_slice(&((record.info.rank << 2) | flags).to_le_bytes());
8215 return bytes;
8216 }
8217 let kmer_bytes = discontinuity_edge_kmer_bytes::<K>();
8218 let path_off = kmer_bytes;
8219 let rank_off = 2 * kmer_bytes;
8220 let flags_off = rank_off + 8;
8221 let mut bytes = [0u8; 41];
8222 bytes[..kmer_bytes].copy_from_slice(&record.vertex.as_u128().to_le_bytes()[..kmer_bytes]);
8223 bytes[path_off..rank_off]
8224 .copy_from_slice(&record.info.path_id.as_u128().to_le_bytes()[..kmer_bytes]);
8225 bytes[rank_off..flags_off].copy_from_slice(&record.info.rank.to_le_bytes());
8226 let mut flags = 0u8;
8227 if record.info.exit_side == Side::Back {
8228 flags |= 1;
8229 }
8230 if record.info.is_cycle {
8231 flags |= 2;
8232 }
8233 bytes[flags_off] = flags;
8234 bytes
8235}
8236
8237fn write_meta_vertex_bucket_parallel<const K: usize>(
8238 dir: &Path,
8239 bucket_id: usize,
8240 records: &[SerialMetaVertex<K>],
8241 pool: &ThreadPool,
8242) -> Result<(), SerialCollationError> {
8243 if records.is_empty() {
8244 return Ok(());
8245 }
8246 let path = vertex_path_info_bucket_path(dir, bucket_id);
8247 let file = File::create(&path).map_err(|source| SerialCollationError::Io {
8248 path: path.clone(),
8249 source,
8250 })?;
8251 let record_len = vertex_path_info_record_len::<K>();
8252 file.set_len((records.len() * record_len) as u64)
8253 .map_err(|source| SerialCollationError::Io {
8254 path: path.clone(),
8255 source,
8256 })?;
8257 let records_per_chunk = (1024 * 1024 / record_len).max(1);
8258 pool.install(|| {
8259 records
8260 .par_chunks(records_per_chunk)
8261 .enumerate()
8262 .try_for_each(|(chunk_id, chunk)| {
8263 let mut encoded = Vec::with_capacity(chunk.len() * record_len);
8264 for meta in chunk {
8265 let record = VertexPathInfo {
8266 vertex: meta.vertex,
8267 info: PathInfo {
8268 path_id: meta.vertex,
8269 rank: meta.weight,
8270 exit_side: meta.entry_side,
8271 is_cycle: meta.is_cycle,
8272 },
8273 };
8274 encoded
8275 .extend_from_slice(&encoded_vertex_path_info_record(&record)[..record_len]);
8276 }
8277 let offset = (chunk_id * records_per_chunk * record_len) as u64;
8278 file.write_all_at(&encoded, offset)
8279 .map_err(|source| SerialCollationError::Io {
8280 path: path.clone(),
8281 source,
8282 })
8283 })
8284 })
8285}
8286
8287#[inline(always)]
8288fn append_encoded_compact_vertex_path_info<const K: usize>(
8289 output: &mut Vec<u8>,
8290 vertex: Kmer<K>,
8291 info: CompactExpansionPathInfo,
8292) {
8293 debug_assert!(K <= 31);
8294 output.extend_from_slice(&(vertex.as_u128() as u64).to_le_bytes());
8295 output.extend_from_slice(&info.path_id.to_le_bytes());
8296 output.extend_from_slice(&info.rank_and_flags.to_le_bytes());
8297}
8298
8299fn decoded_vertex_path_info_record<const K: usize>(bytes: &[u8]) -> VertexPathInfo<K> {
8300 if K <= 31 {
8301 let vertex = u64::from_le_bytes(bytes[..8].try_into().expect("u64 vertex field"));
8302 let path_id = u64::from_le_bytes(bytes[8..16].try_into().expect("u64 path field"));
8303 let rank_and_flags = u64::from_le_bytes(bytes[16..24].try_into().expect("u64 rank field"));
8304 return VertexPathInfo {
8305 vertex: Kmer::from_bits(vertex as u128),
8306 info: PathInfo {
8307 path_id: Kmer::from_bits(path_id as u128),
8308 rank: rank_and_flags >> 2,
8309 exit_side: if rank_and_flags & 1 == 0 {
8310 Side::Front
8311 } else {
8312 Side::Back
8313 },
8314 is_cycle: rank_and_flags & 2 != 0,
8315 },
8316 };
8317 }
8318 let kmer_bytes = discontinuity_edge_kmer_bytes::<K>();
8319 let path_off = kmer_bytes;
8320 let rank_off = 2 * kmer_bytes;
8321 let flags_off = rank_off + 8;
8322 let mut vertex = [0u8; 16];
8323 vertex[..kmer_bytes].copy_from_slice(&bytes[..kmer_bytes]);
8324 let mut path_id = [0u8; 16];
8325 path_id[..kmer_bytes].copy_from_slice(&bytes[path_off..rank_off]);
8326 let mut rank = [0u8; 8];
8327 rank.copy_from_slice(&bytes[rank_off..flags_off]);
8328 let flags = bytes[flags_off];
8329 VertexPathInfo {
8330 vertex: Kmer::from_bits(u128::from_le_bytes(vertex)),
8331 info: PathInfo {
8332 path_id: Kmer::from_bits(u128::from_le_bytes(path_id)),
8333 rank: u64::from_le_bytes(rank),
8334 exit_side: if flags & 1 == 0 {
8335 Side::Front
8336 } else {
8337 Side::Back
8338 },
8339 is_cycle: flags & 2 != 0,
8340 },
8341 }
8342}
8343
8344fn read_vertex_path_info_bucket_into<const K: usize>(
8345 dir: &Path,
8346 bucket_id: usize,
8347 malformed_path: &Path,
8348 map: &ExpansionPathInfoTable<K>,
8349 pool: &ThreadPool,
8350 read_elapsed: &mut Duration,
8351 insert_elapsed: &mut Duration,
8352) -> Result<(), SerialCollationError> {
8353 let path = vertex_path_info_bucket_path(dir, bucket_id);
8354 if !path.exists() {
8355 return Ok(());
8356 }
8357 let record_len = vertex_path_info_record_len::<K>();
8358 let byte_len = fs::metadata(&path)
8359 .map_err(|source| SerialCollationError::Io {
8360 path: path.clone(),
8361 source,
8362 })?
8363 .len() as usize;
8364 if byte_len % record_len != 0 {
8365 return Err(SerialCollationError::MalformedCoordBucket(
8366 malformed_path.to_path_buf(),
8367 ));
8368 }
8369 if K <= 31 {
8370 let records_per_chunk = (1024 * 1024 / record_len).max(1);
8371 let record_count = byte_len / record_len;
8372 let next_record = AtomicUsize::new(0);
8373 let file = File::open(&path).map_err(|source| SerialCollationError::Io {
8374 path: path.clone(),
8375 source,
8376 })?;
8377 let worker_count = pool.current_num_threads().max(1).min(record_count.max(1));
8378 let worker_times = pool.install(|| {
8379 (0..worker_count)
8380 .into_par_iter()
8381 .map(|_| {
8382 let empty = CompactVertexPathInfoRecord {
8383 vertex: 0,
8384 path_id: 0,
8385 rank_and_flags: 0,
8386 };
8387 let mut records = vec![empty; records_per_chunk];
8388 let mut read_time = Duration::default();
8389 let mut insert_time = Duration::default();
8390 loop {
8391 let start = next_record.fetch_add(records_per_chunk, Ordering::Relaxed);
8392 if start >= record_count {
8393 break;
8394 }
8395 let count = records_per_chunk.min(record_count - start);
8396 let bytes = unsafe {
8397 std::slice::from_raw_parts_mut(
8398 records.as_mut_ptr().cast::<u8>(),
8399 count * record_len,
8400 )
8401 };
8402 let started = Instant::now();
8403 file.read_exact_at(bytes, (start * record_len) as u64)
8404 .map_err(|source| SerialCollationError::Io {
8405 path: path.clone(),
8406 source,
8407 })?;
8408 read_time += started.elapsed();
8409 let started = Instant::now();
8410 for &record in &records[..count] {
8411 map.insert_compact_record(record);
8412 }
8413 insert_time += started.elapsed();
8414 }
8415 Ok::<_, SerialCollationError>((read_time, insert_time))
8416 })
8417 .collect::<Result<Vec<_>, SerialCollationError>>()
8418 })?;
8419 let total_read = worker_times.iter().map(|times| times.0).sum::<Duration>();
8420 let total_insert = worker_times.iter().map(|times| times.1).sum::<Duration>();
8421 *read_elapsed += total_read / worker_count as u32;
8422 *insert_elapsed += total_insert / worker_count as u32;
8423 return Ok(());
8424 }
8425 let phase = Instant::now();
8426 let bytes = fs::read(&path).map_err(|source| SerialCollationError::Io {
8427 path: path.clone(),
8428 source,
8429 })?;
8430 *read_elapsed += phase.elapsed();
8431 let records_per_chunk = (1024 * 1024 / record_len).max(1);
8432 let phase = Instant::now();
8433 pool.install(|| {
8434 bytes
8435 .par_chunks(records_per_chunk * record_len)
8436 .for_each(|block| {
8437 for chunk in block.chunks_exact(record_len) {
8438 map.insert_encoded(chunk);
8439 }
8440 });
8441 });
8442 *insert_elapsed += phase.elapsed();
8443 Ok(())
8444}
8445
8446#[allow(clippy::too_many_arguments)]
8447fn push_edge_path_record_to_range_bucket_writer<const K: usize>(
8448 edge: &DiscontinuityEdge<K>,
8449 info: PathInfo<K>,
8450 ranges: &[ExternalLocalUnitigRange],
8451 range_index: &ExternalRangeIndex,
8452 ranges_per_bucket: usize,
8453 writers: &ConcurrentStitchedCoordWriters<'_>,
8454 phantom_records: &mut Vec<(StitchedCoordRecord, DiscontinuityEndpoint<K>)>,
8455 error_path: &Path,
8456) -> Result<(), SerialCollationError> {
8457 let record = stitched_record_from_edge_path_info(edge, info, error_path)?;
8458 if let Some(phantom) = edge.phantom_unitig {
8459 phantom_records.push((record, phantom));
8460 return Ok(());
8461 }
8462 let bucket_id = edge_path_info_bucket(edge, ranges, range_index, ranges_per_bucket)
8463 .ok_or_else(|| SerialCollationError::MalformedCoordBucket(error_path.to_path_buf()))?;
8464 writers.write_record(bucket_id, record)
8465}
8466
8467#[allow(clippy::too_many_arguments)]
8468#[allow(dead_code)]
8470fn push_edge_path_record_to_range_buffers<const K: usize>(
8471 edge: &DiscontinuityEdge<K>,
8472 info: PathInfo<K>,
8473 ranges: &[ExternalLocalUnitigRange],
8474 range_index: &ExternalRangeIndex,
8475 ranges_per_bucket: usize,
8476 buffers: &mut [Vec<StitchedCoordRecord>],
8477 phantom_records: &mut Vec<(StitchedCoordRecord, DiscontinuityEndpoint<K>)>,
8478 error_path: &Path,
8479) -> Result<(), SerialCollationError> {
8480 let record = stitched_record_from_edge_path_info(edge, info, error_path)?;
8481 if let Some(phantom) = edge.phantom_unitig {
8482 phantom_records.push((record, phantom));
8483 return Ok(());
8484 }
8485 let bucket_id = edge_path_info_bucket(edge, ranges, range_index, ranges_per_bucket)
8486 .ok_or_else(|| SerialCollationError::MalformedCoordBucket(error_path.to_path_buf()))?;
8487 buffers[bucket_id].push(record);
8488 Ok(())
8489}
8490
8491#[allow(clippy::too_many_arguments)]
8492#[allow(dead_code)]
8494fn push_compact_edge_path_record_to_range_buffers<const K: usize>(
8495 edge: &DiscontinuityEdge<K>,
8496 info: CompactExpansionPathInfo,
8497 ranges: &[ExternalLocalUnitigRange],
8498 range_index: &ExternalRangeIndex,
8499 ranges_per_bucket: usize,
8500 buffers: &mut [Vec<StitchedCoordRecord>],
8501 phantom_records: &mut Vec<(StitchedCoordRecord, DiscontinuityEndpoint<K>)>,
8502 error_path: &Path,
8503) -> Result<(), SerialCollationError> {
8504 let record = stitched_record_from_compact_edge_path_info(edge, info, error_path)?;
8505 if let Some(phantom) = edge.phantom_unitig {
8506 phantom_records.push((record, phantom));
8507 return Ok(());
8508 }
8509 let bucket_id = edge_path_info_bucket(edge, ranges, range_index, ranges_per_bucket)
8510 .ok_or_else(|| SerialCollationError::MalformedCoordBucket(error_path.to_path_buf()))?;
8511 buffers[bucket_id].push(record);
8512 Ok(())
8513}
8514
8515#[inline]
8516fn edge_path_info_bucket<const K: usize>(
8517 edge: &DiscontinuityEdge<K>,
8518 ranges: &[ExternalLocalUnitigRange],
8519 range_index: &ExternalRangeIndex,
8520 ranges_per_bucket: usize,
8521) -> Option<usize> {
8522 if edge.unitig_bucket != 0 {
8523 Some(edge.unitig_bucket as usize - 1)
8524 } else {
8525 range_index
8526 .find(ranges, edge.unitig_index)
8527 .map(|range_id| range_id / ranges_per_bucket)
8528 }
8529}
8530
8531fn stitched_record_from_edge_path_info<const K: usize>(
8532 edge: &DiscontinuityEdge<K>,
8533 info: PathInfo<K>,
8534 error_path: &Path,
8535) -> Result<StitchedCoordRecord, SerialCollationError> {
8536 let path_id = u64::try_from(info.path_id.as_u128())
8537 .map_err(|_| SerialCollationError::MalformedCoordBucket(error_path.to_path_buf()))?;
8538 let unitig_index = u32::try_from(edge.unitig_index)
8539 .map_err(|_| SerialCollationError::MalformedCoordBucket(error_path.to_path_buf()))?;
8540 Ok(StitchedCoordRecord {
8541 path_id,
8542 rank: info.rank,
8543 unitig_index,
8544 reverse: info.exit_side == Side::Front,
8545 is_cycle: info.is_cycle,
8546 })
8547}
8548
8549#[inline(always)]
8550fn stitched_record_from_compact_edge_path_info<const K: usize>(
8551 edge: &DiscontinuityEdge<K>,
8552 info: CompactExpansionPathInfo,
8553 error_path: &Path,
8554) -> Result<StitchedCoordRecord, SerialCollationError> {
8555 let unitig_index = u32::try_from(edge.unitig_index)
8556 .map_err(|_| SerialCollationError::MalformedCoordBucket(error_path.to_path_buf()))?;
8557 Ok(StitchedCoordRecord {
8558 path_id: info.path_id,
8559 rank: info.rank(),
8560 unitig_index,
8561 reverse: info.exit_side() == Side::Front,
8562 is_cycle: info.is_cycle(),
8563 })
8564}
8565
8566#[allow(dead_code)]
8568fn write_external_cpp_path_info_materialized_coord_buckets<const K: usize>(
8569 inputs: &ExternalDiscontinuityInputs<K>,
8570 coord_dir: &Path,
8571 threads: usize,
8572 expansion: &SerialExpansion<K>,
8573 final_buckets: &mut FinalUnitigBucketWriters,
8574) -> Result<ExternalCppPathInfoMaterialized, SerialCollationError> {
8575 let _ = final_buckets;
8576 write_external_cpp_path_info_materialized_coord_buckets_bucketed::<K>(
8577 inputs, coord_dir, threads, expansion,
8578 )
8579}
8580
8581#[allow(dead_code)]
8583fn write_external_cpp_path_info_materialized_coord_buckets_from_bucketed<const K: usize>(
8584 inputs: &ExternalDiscontinuityInputs<K>,
8585 coord_dir: &Path,
8586 threads: usize,
8587 ranges_per_bucket: usize,
8588 expansion: RangeBucketedExpansion<K>,
8589) -> Result<ExternalCppPathInfoMaterialized, SerialCollationError> {
8590 if coord_dir.exists() {
8591 fs::remove_dir_all(coord_dir).map_err(|source| SerialCollationError::Io {
8592 path: coord_dir.to_path_buf(),
8593 source,
8594 })?;
8595 }
8596 fs::create_dir_all(coord_dir).map_err(|source| SerialCollationError::Io {
8597 path: coord_dir.to_path_buf(),
8598 source,
8599 })?;
8600
8601 let max_unitig_bucket_count = materialized_stitched_coord_bucket_count(threads);
8602 let max_unitig_bucket_mask = max_unitig_bucket_count - 1;
8603 let mut phantom_writers = MaterializedStitchedCoordShardWriters::new(
8604 coord_dir,
8605 threads.max(1),
8606 max_unitig_bucket_count,
8607 );
8608 for (record, phantom) in expansion.phantom_records {
8609 let label = phantom_unitig_label(phantom);
8610 let bucket_id = stitched_coord_bucket(record.path_id, max_unitig_bucket_mask);
8611 phantom_writers.write_materialized_record(bucket_id, &record, &label)?;
8612 }
8613
8614 let path_info_records = expansion.records;
8615 let workers = threads.max(1).min(path_info_records.len().max(1));
8616 let mut manifest = if workers == 1 || path_info_records.len() < 2 {
8617 let mut writers =
8618 MaterializedStitchedCoordShardWriters::new(coord_dir, 0, max_unitig_bucket_count);
8619 for (bucket_id, records) in path_info_records.iter().enumerate() {
8620 materialize_external_stitched_coord_range_group_from_records::<K>(
8621 inputs,
8622 max_unitig_bucket_mask,
8623 ranges_per_bucket,
8624 bucket_id,
8625 records,
8626 &mut writers,
8627 )?;
8628 }
8629 writers.finish()?
8630 } else {
8631 let next_bucket = AtomicUsize::new(0);
8632 std::thread::scope(|scope| {
8633 let mut handles = Vec::new();
8634 for worker_id in 0..workers {
8635 let next_bucket = &next_bucket;
8636 let path_info_records = &path_info_records;
8637 handles.push(scope.spawn(move || {
8638 let mut writers = MaterializedStitchedCoordShardWriters::new(
8639 coord_dir,
8640 worker_id,
8641 max_unitig_bucket_count,
8642 );
8643 loop {
8644 let bucket_id = next_bucket.fetch_add(1, Ordering::Relaxed);
8645 let Some(records) = path_info_records.get(bucket_id) else {
8646 break;
8647 };
8648 materialize_external_stitched_coord_range_group_from_records::<K>(
8649 inputs,
8650 max_unitig_bucket_mask,
8651 ranges_per_bucket,
8652 bucket_id,
8653 records,
8654 &mut writers,
8655 )?;
8656 }
8657 writers.finish()
8658 }));
8659 }
8660
8661 let mut manifest = Vec::new();
8662 for handle in handles {
8663 manifest.extend(
8664 handle
8665 .join()
8666 .map_err(|_| SerialCollationError::WorkerPanic)??,
8667 );
8668 }
8669 Ok::<_, SerialCollationError>(manifest)
8670 })?
8671 };
8672 manifest.extend(phantom_writers.finish()?);
8673 manifest.sort_by(|left, right| {
8674 left.bucket_id
8675 .cmp(&right.bucket_id)
8676 .then_with(|| left.coord_path.cmp(&right.coord_path))
8677 });
8678 Ok(ExternalCppPathInfoMaterialized {
8679 manifest,
8680 retained: Vec::new(),
8681 direct_local_unitigs: 0,
8682 direct_local_unitigs_complete: false,
8683 })
8684}
8685
8686fn map_external_cpp_path_info_buckets_to_max_unitig_buckets<const K: usize>(
8687 inputs: &ExternalDiscontinuityInputs<K>,
8688 coord_dir: &Path,
8689 threads: usize,
8690 ranges_per_bucket: usize,
8691 expansion: RangeBucketedExpansion<K>,
8692 final_buckets: &mut FinalUnitigBucketWriters,
8693) -> Result<ExternalCppPathInfoMaterialized, SerialCollationError> {
8694 const DIRECT_FINAL_BATCH_RECORDS: usize = 4096;
8695 if coord_dir.exists() {
8696 fs::remove_dir_all(coord_dir).map_err(|source| SerialCollationError::Io {
8697 path: coord_dir.to_path_buf(),
8698 source,
8699 })?;
8700 }
8701 fs::create_dir_all(coord_dir).map_err(|source| SerialCollationError::Io {
8702 path: coord_dir.to_path_buf(),
8703 source,
8704 })?;
8705
8706 let colored = inputs.color_runs.is_some()
8707 || inputs
8708 .local_unitig_buckets
8709 .as_ref()
8710 .is_some_and(|buckets| buckets.iter().any(|bucket| bucket.colored));
8711 let (max_unitig_bucket_count, mapping_threads, open_writer_limit) =
8712 materialized_coordinate_plan(
8713 open_file_limit(),
8714 current_open_file_count(),
8715 threads,
8716 colored,
8717 inputs.stats.unitig_bases,
8718 );
8719 eprintln!(
8720 "cuttlefish: materializing final coordinates into {max_unitig_bucket_count} bucket(s) with {mapping_threads} mapping worker(s), {open_writer_limit} open writer(s); {} local unitig base(s)",
8721 inputs.stats.unitig_bases
8722 );
8723 let max_unitig_bucket_mask = max_unitig_bucket_count - 1;
8724 let shared_writers =
8725 SharedMaterializedWriters::new(coord_dir, max_unitig_bucket_count, open_writer_limit);
8726 let mut phantom_writers =
8727 SharedMaterializedBatch::new(&shared_writers, max_unitig_bucket_count);
8728 for (record, phantom) in expansion.phantom_records {
8729 let label = phantom_unitig_label(phantom);
8730 let bucket_id = stitched_coord_bucket(record.path_id, max_unitig_bucket_mask);
8731 phantom_writers.write_materialized_record(bucket_id, &record, &label)?;
8732 }
8733 phantom_writers.finish()?;
8734
8735 let mut manifest = Vec::new();
8736 let mut direct_local_unitigs = 0u64;
8737 let path_info_manifest = expansion.record_manifest;
8738 if let Some(local_buckets) = inputs.local_unitig_buckets.as_ref() {
8739 direct_local_unitigs = map_local_unitig_buckets_to_max_unitig_buckets::<K>(
8740 inputs,
8741 local_buckets,
8742 path_info_manifest,
8743 mapping_threads,
8744 max_unitig_bucket_mask,
8745 &shared_writers,
8746 final_buckets,
8747 )?;
8748 let materialized = shared_writers.finish()?;
8749 manifest = materialized.manifest;
8750 manifest.sort_by(|left, right| {
8751 left.bucket_id
8752 .cmp(&right.bucket_id)
8753 .then_with(|| left.coord_path.cmp(&right.coord_path))
8754 });
8755 return Ok(ExternalCppPathInfoMaterialized {
8756 manifest,
8757 retained: materialized.retained,
8758 direct_local_unitigs,
8759 direct_local_unitigs_complete: true,
8760 });
8761 }
8762 if !path_info_manifest.is_empty() {
8763 let path_info_bucket_count = inputs.ranges.len().div_ceil(ranges_per_bucket).max(1);
8764 let mut groups = vec![Vec::<StitchedCoordBucketEntry>::new(); path_info_bucket_count];
8765 for entry in path_info_manifest {
8766 if entry.bucket_id >= groups.len() {
8767 return Err(SerialCollationError::MalformedCoordBucket(entry.path));
8768 }
8769 groups[entry.bucket_id].push(entry);
8770 }
8771
8772 let workers = mapping_threads.max(1).min(groups.len().max(1));
8773 if workers == 1 || groups.len() < 2 {
8774 let mut writers =
8775 SharedMaterializedBatch::new(&shared_writers, max_unitig_bucket_count);
8776 for (bucket_id, group) in groups.iter().enumerate() {
8777 let records = read_stitched_coord_bucket_group(group)?;
8778 let mut emit_direct =
8779 |unitig: FinalUnitigRecord| -> Result<(), SerialCollationError> {
8780 if unitig.colors.is_empty() {
8781 final_buckets.write_label(&unitig.label)?;
8782 } else {
8783 final_buckets.write_colored_label(&unitig.label, &unitig.colors)?;
8784 }
8785 direct_local_unitigs += 1;
8786 Ok(())
8787 };
8788 map_external_cpp_path_info_range_bucket_owned::<K, _, _>(
8789 inputs,
8790 max_unitig_bucket_mask,
8791 ranges_per_bucket,
8792 bucket_id,
8793 records,
8794 &mut writers,
8795 &mut emit_direct,
8796 )?;
8797 }
8798 writers.finish()?;
8799 } else {
8800 let next_bucket = AtomicUsize::new(0);
8801 let next_direct_record = AtomicU64::new(final_buckets.direct_record_id_highwater + 1);
8802 let (tx, rx) =
8803 mpsc::sync_channel::<Result<EncodedFinalBatch, SerialCollationError>>(workers * 2);
8804 let mut first_error = None;
8805 std::thread::scope(|scope| {
8806 let mut handles = Vec::new();
8807 for _worker_id in 0..workers {
8808 let next_bucket = &next_bucket;
8809 let next_direct_record = &next_direct_record;
8810 let groups = &groups;
8811 let tx = tx.clone();
8812 let shared_writers = &shared_writers;
8813 handles.push(scope.spawn(move || {
8814 let mut writers =
8815 SharedMaterializedBatch::new(shared_writers, max_unitig_bucket_count);
8816 let mut direct_batch = Vec::with_capacity(DIRECT_FINAL_BATCH_RECORDS);
8817 let mut emit_direct =
8818 |unitig: FinalUnitigRecord| -> Result<(), SerialCollationError> {
8819 direct_batch.push(unitig);
8820 if direct_batch.len() >= DIRECT_FINAL_BATCH_RECORDS {
8821 let batch = std::mem::take(&mut direct_batch);
8822 let first_record = next_direct_record
8823 .fetch_add(batch.len() as u64, Ordering::Relaxed);
8824 tx.send(Ok(encode_final_unitig_batch(batch, first_record)))
8825 .map_err(|_| SerialCollationError::WorkerPanic)?;
8826 }
8827 Ok(())
8828 };
8829
8830 loop {
8831 let bucket_id = next_bucket.fetch_add(1, Ordering::Relaxed);
8832 let Some(group) = groups.get(bucket_id) else {
8833 break;
8834 };
8835 let records = read_stitched_coord_bucket_group(group)?;
8836 map_external_cpp_path_info_range_bucket_owned::<K, _, _>(
8837 inputs,
8838 max_unitig_bucket_mask,
8839 ranges_per_bucket,
8840 bucket_id,
8841 records,
8842 &mut writers,
8843 &mut emit_direct,
8844 )?;
8845 }
8846 if !direct_batch.is_empty() {
8847 let first_record = next_direct_record
8848 .fetch_add(direct_batch.len() as u64, Ordering::Relaxed);
8849 tx.send(Ok(encode_final_unitig_batch(direct_batch, first_record)))
8850 .map_err(|_| SerialCollationError::WorkerPanic)?;
8851 }
8852 writers.finish()
8853 }));
8854 }
8855 drop(tx);
8856
8857 for result in rx {
8858 match result {
8859 Ok(batch) if first_error.is_none() => {
8860 if let Err(err) = final_buckets.write_direct_batch(
8861 &batch.bytes,
8862 batch.records,
8863 batch.bases,
8864 ) {
8865 first_error = Some(err);
8866 } else {
8867 direct_local_unitigs += batch.records;
8868 }
8869 }
8870 Ok(_) => {}
8871 Err(err) if first_error.is_none() => first_error = Some(err),
8872 Err(_) => {}
8873 }
8874 }
8875
8876 for handle in handles {
8877 handle
8878 .join()
8879 .map_err(|_| SerialCollationError::WorkerPanic)??;
8880 }
8881 Ok::<_, SerialCollationError>(())
8882 })?;
8883 final_buckets.direct_record_id_highwater =
8884 next_direct_record.load(Ordering::Relaxed).saturating_sub(1);
8885 if let Some(err) = first_error {
8886 return Err(err);
8887 }
8888 }
8889
8890 let materialized = shared_writers.finish()?;
8891 manifest = materialized.manifest;
8892 manifest.sort_by(|left, right| {
8893 left.bucket_id
8894 .cmp(&right.bucket_id)
8895 .then_with(|| left.coord_path.cmp(&right.coord_path))
8896 });
8897 return Ok(ExternalCppPathInfoMaterialized {
8898 manifest,
8899 retained: materialized.retained,
8900 direct_local_unitigs,
8901 direct_local_unitigs_complete: true,
8902 });
8903 }
8904
8905 let path_info_records = expansion.records;
8906 let workers = threads.max(1).min(path_info_records.len().max(1));
8907
8908 if workers == 1 || path_info_records.len() < 2 {
8909 let mut writers =
8910 MaterializedStitchedCoordShardWriters::new(coord_dir, 0, max_unitig_bucket_count);
8911 for (bucket_id, records) in path_info_records.iter().enumerate() {
8912 let mut emit_direct = |unitig: FinalUnitigRecord| -> Result<(), SerialCollationError> {
8913 if unitig.colors.is_empty() {
8914 final_buckets.write_label(&unitig.label)?;
8915 } else {
8916 final_buckets.write_colored_label(&unitig.label, &unitig.colors)?;
8917 }
8918 direct_local_unitigs += 1;
8919 Ok(())
8920 };
8921 map_external_cpp_path_info_range_bucket::<K, _, _>(
8922 inputs,
8923 max_unitig_bucket_mask,
8924 ranges_per_bucket,
8925 bucket_id,
8926 records,
8927 &mut writers,
8928 &mut emit_direct,
8929 )?;
8930 }
8931 manifest.extend(writers.finish()?);
8932 } else {
8933 let next_bucket = AtomicUsize::new(0);
8934 let (tx, rx) =
8935 mpsc::sync_channel::<Result<Vec<FinalUnitigRecord>, SerialCollationError>>(workers * 2);
8936 let mut first_error = None;
8937 std::thread::scope(|scope| {
8938 let mut handles = Vec::new();
8939 for worker_id in 0..workers {
8940 let next_bucket = &next_bucket;
8941 let path_info_records = &path_info_records;
8942 let tx = tx.clone();
8943 handles.push(scope.spawn(move || {
8944 let mut writers = MaterializedStitchedCoordShardWriters::new(
8945 coord_dir,
8946 worker_id,
8947 max_unitig_bucket_count,
8948 );
8949 let mut direct_batch = Vec::with_capacity(256);
8950 let mut emit_direct =
8951 |unitig: FinalUnitigRecord| -> Result<(), SerialCollationError> {
8952 direct_batch.push(unitig);
8953 if direct_batch.len() >= 256 {
8954 let batch = std::mem::take(&mut direct_batch);
8955 tx.send(Ok(batch))
8956 .map_err(|_| SerialCollationError::WorkerPanic)?;
8957 }
8958 Ok(())
8959 };
8960
8961 loop {
8962 let bucket_id = next_bucket.fetch_add(1, Ordering::Relaxed);
8963 let Some(records) = path_info_records.get(bucket_id) else {
8964 break;
8965 };
8966 map_external_cpp_path_info_range_bucket::<K, _, _>(
8967 inputs,
8968 max_unitig_bucket_mask,
8969 ranges_per_bucket,
8970 bucket_id,
8971 records,
8972 &mut writers,
8973 &mut emit_direct,
8974 )?;
8975 }
8976 if !direct_batch.is_empty() {
8977 tx.send(Ok(direct_batch))
8978 .map_err(|_| SerialCollationError::WorkerPanic)?;
8979 }
8980 writers.finish()
8981 }));
8982 }
8983 drop(tx);
8984
8985 for result in rx {
8986 match result {
8987 Ok(unitigs) if first_error.is_none() => {
8988 for unitig in unitigs {
8989 let write = if unitig.colors.is_empty() {
8990 final_buckets.write_label(&unitig.label)
8991 } else {
8992 final_buckets.write_colored_label(&unitig.label, &unitig.colors)
8993 };
8994 if let Err(err) = write {
8995 first_error = Some(err);
8996 break;
8997 }
8998 direct_local_unitigs += 1;
8999 }
9000 }
9001 Ok(_) => {}
9002 Err(err) if first_error.is_none() => first_error = Some(err),
9003 Err(_) => {}
9004 }
9005 }
9006
9007 for handle in handles {
9008 manifest.extend(
9009 handle
9010 .join()
9011 .map_err(|_| SerialCollationError::WorkerPanic)??,
9012 );
9013 }
9014 Ok::<_, SerialCollationError>(())
9015 })?;
9016 if let Some(err) = first_error {
9017 return Err(err);
9018 }
9019 }
9020
9021 let materialized = shared_writers.finish()?;
9022 manifest.extend(materialized.manifest);
9023 manifest.sort_by(|left, right| {
9024 left.bucket_id
9025 .cmp(&right.bucket_id)
9026 .then_with(|| left.coord_path.cmp(&right.coord_path))
9027 });
9028 Ok(ExternalCppPathInfoMaterialized {
9029 manifest,
9030 retained: materialized.retained,
9031 direct_local_unitigs,
9032 direct_local_unitigs_complete: true,
9033 })
9034}
9035
9036fn map_local_unitig_buckets_to_max_unitig_buckets<const K: usize>(
9037 inputs: &ExternalDiscontinuityInputs<K>,
9038 local_buckets: &[LocalUnitigBucketEntry],
9039 path_info_manifest: Vec<StitchedCoordBucketEntry>,
9040 threads: usize,
9041 max_unitig_bucket_mask: usize,
9042 shared_writers: &SharedMaterializedWriters<'_>,
9043 final_buckets: &mut FinalUnitigBucketWriters,
9044) -> Result<u64, SerialCollationError> {
9045 const DIRECT_FINAL_BATCH_RECORDS: usize = 4096;
9046 let mut groups = (0..local_buckets.len())
9047 .map(|_| Vec::<StitchedCoordBucketEntry>::new())
9048 .collect::<Vec<_>>();
9049 for entry in path_info_manifest {
9050 let Some(group) = groups.get_mut(entry.bucket_id) else {
9051 return Err(SerialCollationError::MalformedCoordBucket(entry.path));
9052 };
9053 group.push(entry);
9054 }
9055
9056 let workers = threads.max(1).min(local_buckets.len().max(1));
9057 let next_bucket = AtomicUsize::new(0);
9058 let next_direct_record = AtomicU64::new(final_buckets.direct_record_id_highwater + 1);
9059 let (tx, rx) =
9060 mpsc::sync_channel::<Result<EncodedFinalBatch, SerialCollationError>>(workers * 2);
9061 let mut direct_local_unitigs = 0u64;
9062 let mut first_error = None;
9063 std::thread::scope(|scope| {
9064 let mut handles = Vec::new();
9065 for _ in 0..workers {
9066 let tx = tx.clone();
9067 let next_bucket = &next_bucket;
9068 let next_direct_record = &next_direct_record;
9069 let groups = &groups;
9070 handles.push(scope.spawn(move || {
9071 let mut writers =
9072 SharedMaterializedBatch::new(shared_writers, max_unitig_bucket_mask + 1);
9073 let mut path_info = Vec::new();
9074 let mut direct_batch = None::<DirectFinalBatchBuilder>;
9075 let mut emit_direct =
9076 |label: &[u8], colors: &[UnitigColor]| -> Result<(), SerialCollationError> {
9077 let batch = direct_batch.get_or_insert_with(|| {
9078 let first_record = next_direct_record
9079 .fetch_add(DIRECT_FINAL_BATCH_RECORDS as u64, Ordering::Relaxed);
9080 DirectFinalBatchBuilder::new(first_record)
9081 });
9082 batch.push(label, colors);
9083 if batch.records as usize >= DIRECT_FINAL_BATCH_RECORDS {
9084 tx.send(Ok(direct_batch.take().expect("batch exists").finish()))
9085 .map_err(|_| SerialCollationError::WorkerPanic)?;
9086 }
9087 Ok(())
9088 };
9089 loop {
9090 let bucket_index = next_bucket.fetch_add(1, Ordering::Relaxed);
9091 let Some(bucket) = local_buckets.get(bucket_index) else {
9092 break;
9093 };
9094 read_stitched_coord_bucket_group_dense(
9095 &groups[bucket_index],
9096 bucket.unitigs,
9097 &bucket.unitig_path,
9098 &mut path_info,
9099 )?;
9100 map_local_unitig_bucket::<K, _, _>(
9101 inputs,
9102 bucket,
9103 &path_info,
9104 max_unitig_bucket_mask,
9105 &mut writers,
9106 &mut emit_direct,
9107 )?;
9108 }
9109 if let Some(direct_batch) = direct_batch {
9110 tx.send(Ok(direct_batch.finish()))
9111 .map_err(|_| SerialCollationError::WorkerPanic)?;
9112 }
9113 writers.finish()
9114 }));
9115 }
9116 drop(tx);
9117 for result in rx {
9118 match result {
9119 Ok(batch) if first_error.is_none() => {
9120 if let Err(err) =
9121 final_buckets.write_direct_batch(&batch.bytes, batch.records, batch.bases)
9122 {
9123 first_error = Some(err);
9124 } else {
9125 direct_local_unitigs += batch.records;
9126 }
9127 }
9128 Ok(_) => {}
9129 Err(err) if first_error.is_none() => first_error = Some(err),
9130 Err(_) => {}
9131 }
9132 }
9133 for handle in handles {
9134 if let Err(err) = handle
9135 .join()
9136 .map_err(|_| SerialCollationError::WorkerPanic)?
9137 && first_error.is_none()
9138 {
9139 first_error = Some(err);
9140 }
9141 }
9142 Ok::<_, SerialCollationError>(())
9143 })?;
9144 final_buckets.direct_record_id_highwater =
9145 next_direct_record.load(Ordering::Relaxed).saturating_sub(1);
9146 if let Some(err) = first_error {
9147 return Err(err);
9148 }
9149 Ok(direct_local_unitigs)
9150}
9151
9152fn map_local_unitig_bucket<const K: usize, F, W>(
9153 inputs: &ExternalDiscontinuityInputs<K>,
9154 bucket: &LocalUnitigBucketEntry,
9155 path_info_by_unitig: &[DenseLocalPathInfo],
9156 max_unitig_bucket_mask: usize,
9157 writers: &mut W,
9158 emit_direct: &mut F,
9159) -> Result<(), SerialCollationError>
9160where
9161 F: FnMut(&[u8], &[UnitigColor]) -> Result<(), SerialCollationError>,
9162 W: MaterializedRecordSink,
9163{
9164 let unitig_file =
9165 File::open(&bucket.unitig_path).map_err(|source| SerialCollationError::Io {
9166 path: bucket.unitig_path.clone(),
9167 source,
9168 })?;
9169 let label_file = File::open(&bucket.label_path).map_err(|source| SerialCollationError::Io {
9170 path: bucket.label_path.clone(),
9171 source,
9172 })?;
9173 let mut unitig_input = BufReader::with_capacity(1024 * 1024, unitig_file);
9174 let mut label_input = BufReader::with_capacity(4 * 1024 * 1024, label_file);
9175 let mut label = Vec::new();
9176 let mut colors = Vec::new();
9177 let mut discard = vec![0u8; 64 * 1024];
9178
9179 for (unitig_index, &dense_record) in path_info_by_unitig.iter().enumerate() {
9180 let unitig = read_discontinuity_unitig_from_reader::<K>(
9181 &mut unitig_input,
9182 &bucket.unitig_path,
9183 true,
9184 )?;
9185 let has_colors = bucket.colored;
9186 if has_colors {
9187 read_unitig_color_runs(&mut unitig_input, &mut colors).map_err(|source| {
9188 SerialCollationError::Io {
9189 path: bucket.unitig_path.clone(),
9190 source,
9191 }
9192 })?;
9193 }
9194 if let Some(record) = dense_record.to_record(unitig_index) {
9195 label.resize(unitig.label_len as usize, 0);
9196 label_input
9197 .read_exact(&mut label)
9198 .map_err(|source| SerialCollationError::Io {
9199 path: bucket.label_path.clone(),
9200 source,
9201 })?;
9202 let max_bucket = stitched_coord_bucket(record.path_id, max_unitig_bucket_mask);
9203 write_external_materialized_record(
9204 writers,
9205 max_bucket,
9206 &record,
9207 &label,
9208 has_colors.then_some(colors.as_slice()),
9209 )?;
9210 } else if unitig.left_exit().is_none() && unitig.right_exit().is_none() {
9211 label.resize(unitig.label_len as usize, 0);
9212 label_input
9213 .read_exact(&mut label)
9214 .map_err(|source| SerialCollationError::Io {
9215 path: bucket.label_path.clone(),
9216 source,
9217 })?;
9218 let reverse = reverse_complement_is_less(&label);
9219 if reverse {
9220 let reversed_label = reverse_complement_label(&label);
9221 let reversed_colors = if has_colors && !colors.is_empty() {
9222 reverse_color_runs(&colors, (unitig.label_len as usize - K + 1) as u32)
9223 } else {
9224 Vec::new()
9225 };
9226 emit_direct(&reversed_label, &reversed_colors)?;
9227 } else {
9228 let direct_colors = if has_colors { colors.as_slice() } else { &[] };
9229 emit_direct(&label, direct_colors)?;
9230 }
9231 } else {
9232 read_and_discard_exact(
9233 &mut label_input,
9234 &bucket.label_path,
9235 &mut discard,
9236 unitig.label_len as u64,
9237 )?;
9238 }
9239 }
9240 let _ = inputs;
9241 Ok(())
9242}
9243
9244#[allow(dead_code)]
9246fn write_external_cpp_path_info_materialized_coord_buckets_bucketed<const K: usize>(
9247 inputs: &ExternalDiscontinuityInputs<K>,
9248 coord_dir: &Path,
9249 threads: usize,
9250 expansion: &SerialExpansion<K>,
9251) -> Result<ExternalCppPathInfoMaterialized, SerialCollationError> {
9252 if coord_dir.exists() {
9253 fs::remove_dir_all(coord_dir).map_err(|source| SerialCollationError::Io {
9254 path: coord_dir.to_path_buf(),
9255 source,
9256 })?;
9257 }
9258 fs::create_dir_all(coord_dir).map_err(|source| SerialCollationError::Io {
9259 path: coord_dir.to_path_buf(),
9260 source,
9261 })?;
9262
9263 let max_unitig_bucket_count = materialized_stitched_coord_bucket_count(threads);
9264 let max_unitig_bucket_mask = max_unitig_bucket_count - 1;
9265 let ranges_per_bucket = ranges_per_path_info_bucket(inputs.ranges.len(), 1);
9266 let path_info_bucket_count = inputs.ranges.len().div_ceil(ranges_per_bucket).max(1);
9267 let mut path_info_records = (0..path_info_bucket_count)
9268 .map(|_| Vec::<StitchedCoordRecord>::new())
9269 .collect::<Vec<_>>();
9270 let mut phantom_writers = MaterializedStitchedCoordShardWriters::new(
9271 coord_dir,
9272 threads.max(1),
9273 max_unitig_bucket_count,
9274 );
9275
9276 for edge in &expansion.edges {
9277 let path_id = u64::try_from(edge.info.path_id.as_u128())
9278 .map_err(|_| SerialCollationError::MalformedCoordBucket(inputs.unitig_path.clone()))?;
9279 let unitig_index = u32::try_from(edge.unitig_index)
9280 .map_err(|_| SerialCollationError::MalformedCoordBucket(inputs.unitig_path.clone()))?;
9281 let record = StitchedCoordRecord {
9282 path_id,
9283 rank: edge.info.rank,
9284 unitig_index,
9285 reverse: edge.info.exit_side == Side::Front,
9286 is_cycle: edge.info.is_cycle,
9287 };
9288
9289 if let Some(phantom) = edge.phantom_unitig {
9290 let label = phantom_unitig_label(phantom);
9291 let bucket_id = stitched_coord_bucket(path_id, max_unitig_bucket_mask);
9292 phantom_writers.write_materialized_record(bucket_id, &record, &label)?;
9293 continue;
9294 }
9295
9296 let range_id =
9297 external_range_id_for_unitig(&inputs.ranges, edge.unitig_index).ok_or_else(|| {
9298 SerialCollationError::MalformedCoordBucket(inputs.unitig_path.clone())
9299 })?;
9300 let bucket_id = range_id / ranges_per_bucket;
9301 path_info_records
9302 .get_mut(bucket_id)
9303 .ok_or_else(|| SerialCollationError::MalformedCoordBucket(inputs.unitig_path.clone()))?
9304 .push(record);
9305 }
9306
9307 let workers = threads.max(1).min(path_info_records.len().max(1));
9308 let mut manifest = if workers == 1 || path_info_records.len() < 2 {
9309 let mut writers =
9310 MaterializedStitchedCoordShardWriters::new(coord_dir, 0, max_unitig_bucket_count);
9311 for (bucket_id, records) in path_info_records.iter().enumerate() {
9312 materialize_external_stitched_coord_range_group_from_records::<K>(
9313 inputs,
9314 max_unitig_bucket_mask,
9315 ranges_per_bucket,
9316 bucket_id,
9317 records,
9318 &mut writers,
9319 )?;
9320 }
9321 writers.finish()?
9322 } else {
9323 let next_bucket = AtomicUsize::new(0);
9324 std::thread::scope(|scope| {
9325 let mut handles = Vec::new();
9326 for worker_id in 0..workers {
9327 let next_bucket = &next_bucket;
9328 let path_info_records = &path_info_records;
9329 handles.push(scope.spawn(move || {
9330 let mut writers = MaterializedStitchedCoordShardWriters::new(
9331 coord_dir,
9332 worker_id,
9333 max_unitig_bucket_count,
9334 );
9335 loop {
9336 let bucket_id = next_bucket.fetch_add(1, Ordering::Relaxed);
9337 let Some(records) = path_info_records.get(bucket_id) else {
9338 break;
9339 };
9340 materialize_external_stitched_coord_range_group_from_records::<K>(
9341 inputs,
9342 max_unitig_bucket_mask,
9343 ranges_per_bucket,
9344 bucket_id,
9345 records,
9346 &mut writers,
9347 )?;
9348 }
9349 writers.finish()
9350 }));
9351 }
9352
9353 let mut manifest = Vec::new();
9354 for handle in handles {
9355 manifest.extend(
9356 handle
9357 .join()
9358 .map_err(|_| SerialCollationError::WorkerPanic)??,
9359 );
9360 }
9361 Ok::<_, SerialCollationError>(manifest)
9362 })?
9363 };
9364 manifest.extend(phantom_writers.finish()?);
9365 manifest.sort_by(|left, right| {
9366 left.bucket_id
9367 .cmp(&right.bucket_id)
9368 .then_with(|| left.coord_path.cmp(&right.coord_path))
9369 });
9370 Ok(ExternalCppPathInfoMaterialized {
9371 manifest,
9372 retained: Vec::new(),
9373 direct_local_unitigs: 0,
9374 direct_local_unitigs_complete: false,
9375 })
9376}
9377
9378#[allow(dead_code)]
9380fn write_external_cpp_path_info_materialized_coord_buckets_legacy<const K: usize>(
9381 inputs: &ExternalDiscontinuityInputs<K>,
9382 coord_dir: &Path,
9383 threads: usize,
9384 expansion: &SerialExpansion<K>,
9385) -> Result<Vec<MaterializedStitchedCoordBucketEntry>, SerialCollationError> {
9386 if coord_dir.exists() {
9387 fs::remove_dir_all(coord_dir).map_err(|source| SerialCollationError::Io {
9388 path: coord_dir.to_path_buf(),
9389 source,
9390 })?;
9391 }
9392 fs::create_dir_all(coord_dir).map_err(|source| SerialCollationError::Io {
9393 path: coord_dir.to_path_buf(),
9394 source,
9395 })?;
9396
9397 let bucket_count = materialized_stitched_coord_bucket_count(threads);
9398 let bucket_mask = bucket_count - 1;
9399 let mut unitigs = Vec::with_capacity(inputs.unitig_count());
9400 let reader = ExternalDiscontinuityReader::open(inputs)?;
9401 for unitig in reader.iter()? {
9402 unitigs.push(unitig?);
9403 }
9404 let mut label = Vec::new();
9405 let mut writers = MaterializedStitchedCoordShardWriters::new(coord_dir, 0, bucket_count);
9406
9407 for edge in &expansion.edges {
9408 let path_id = u64::try_from(edge.info.path_id.as_u128())
9409 .map_err(|_| SerialCollationError::MalformedCoordBucket(inputs.unitig_path.clone()))?;
9410 let unitig_index = u32::try_from(edge.unitig_index)
9411 .map_err(|_| SerialCollationError::MalformedCoordBucket(inputs.unitig_path.clone()))?;
9412 let record = StitchedCoordRecord {
9413 path_id,
9414 rank: edge.info.rank,
9415 unitig_index,
9416 reverse: edge.info.exit_side == Side::Front,
9417 is_cycle: edge.info.is_cycle,
9418 };
9419 let bucket_id = stitched_coord_bucket(path_id, bucket_mask);
9420 if let Some(phantom) = edge.phantom_unitig {
9421 label = phantom_unitig_label(phantom);
9422 } else {
9423 let unitig = unitigs.get(edge.unitig_index).ok_or_else(|| {
9424 SerialCollationError::MalformedCoordBucket(inputs.unitig_path.clone())
9425 })?;
9426 reader.read_label(unitig, &mut label)?;
9427 }
9428 writers.write_materialized_record(bucket_id, &record, &label)?;
9429 }
9430
9431 writers.finish()
9432}
9433
9434#[allow(dead_code)]
9436fn write_external_cpp_path_info_materialized_coord_buckets_streaming<const K: usize>(
9437 inputs: &ExternalDiscontinuityInputs<K>,
9438 coord_dir: &Path,
9439 threads: usize,
9440 expansion: &SerialExpansion<K>,
9441 final_buckets: &mut FinalUnitigBucketWriters,
9442) -> Result<ExternalCppPathInfoMaterialized, SerialCollationError> {
9443 if coord_dir.exists() {
9444 fs::remove_dir_all(coord_dir).map_err(|source| SerialCollationError::Io {
9445 path: coord_dir.to_path_buf(),
9446 source,
9447 })?;
9448 }
9449 fs::create_dir_all(coord_dir).map_err(|source| SerialCollationError::Io {
9450 path: coord_dir.to_path_buf(),
9451 source,
9452 })?;
9453
9454 let bucket_count = materialized_stitched_coord_bucket_count(threads);
9455 let bucket_mask = bucket_count - 1;
9456 let mut path_info_by_unitig = vec![None::<StitchedCoordRecord>; inputs.unitig_count()];
9457 let mut writers = MaterializedStitchedCoordShardWriters::new(coord_dir, 0, bucket_count);
9458 let mut duplicate_path_info = false;
9459
9460 for edge in &expansion.edges {
9461 let path_id = u64::try_from(edge.info.path_id.as_u128())
9462 .map_err(|_| SerialCollationError::MalformedCoordBucket(inputs.unitig_path.clone()))?;
9463 let unitig_index = u32::try_from(edge.unitig_index)
9464 .map_err(|_| SerialCollationError::MalformedCoordBucket(inputs.unitig_path.clone()))?;
9465 let record = StitchedCoordRecord {
9466 path_id,
9467 rank: edge.info.rank,
9468 unitig_index,
9469 reverse: edge.info.exit_side == Side::Front,
9470 is_cycle: edge.info.is_cycle,
9471 };
9472
9473 if let Some(phantom) = edge.phantom_unitig {
9474 let label = phantom_unitig_label(phantom);
9475 let bucket_id = stitched_coord_bucket(path_id, bucket_mask);
9476 writers.write_materialized_record(bucket_id, &record, &label)?;
9477 continue;
9478 }
9479
9480 let Some(slot) = path_info_by_unitig.get_mut(edge.unitig_index) else {
9481 return Err(SerialCollationError::MalformedCoordBucket(
9482 inputs.unitig_path.clone(),
9483 ));
9484 };
9485 if slot.replace(record).is_some() {
9486 duplicate_path_info = true;
9487 }
9488 }
9489
9490 if duplicate_path_info {
9491 drop(writers);
9492 let manifest = write_external_cpp_path_info_materialized_coord_buckets_legacy::<K>(
9493 inputs, coord_dir, threads, expansion,
9494 )?;
9495 return Ok(ExternalCppPathInfoMaterialized {
9496 manifest,
9497 retained: Vec::new(),
9498 direct_local_unitigs: 0,
9499 direct_local_unitigs_complete: false,
9500 });
9501 }
9502
9503 let unitig_file =
9504 File::open(&inputs.unitig_path).map_err(|source| SerialCollationError::Io {
9505 path: inputs.unitig_path.clone(),
9506 source,
9507 })?;
9508 let mut unitigs: ExternalDiscontinuityIter<K> = ExternalDiscontinuityIter {
9509 input: BufReader::with_capacity(1024 * 1024, unitig_file),
9510 path: inputs.unitig_path.clone(),
9511 remaining: inputs.unitigs,
9512 compact_unitigs: inputs.compact_unitigs,
9513 };
9514 let label_file = File::open(&inputs.label_path).map_err(|source| SerialCollationError::Io {
9515 path: inputs.label_path.clone(),
9516 source,
9517 })?;
9518 let mut labels = BufReader::with_capacity(1024 * 1024, label_file);
9519 let mut label_pos = 0u64;
9520 let mut label = Vec::new();
9521
9522 let mut direct_local_unitigs = 0u64;
9523 for (unitig_index, record) in path_info_by_unitig.into_iter().enumerate() {
9524 let unitig = unitigs.next_unitig()?;
9525 if label_pos != unitig.label_start {
9526 labels
9527 .seek(SeekFrom::Start(unitig.label_start))
9528 .map_err(|source| SerialCollationError::Io {
9529 path: inputs.label_path.clone(),
9530 source,
9531 })?;
9532 label_pos = unitig.label_start;
9533 }
9534 label.resize(unitig.label_len as usize, 0);
9535 labels
9536 .read_exact(&mut label)
9537 .map_err(|source| SerialCollationError::Io {
9538 path: inputs.label_path.clone(),
9539 source,
9540 })?;
9541 label_pos += u64::from(unitig.label_len);
9542
9543 let Some(record) = record else {
9544 if unitig.left_exit().is_none() && unitig.right_exit().is_none() {
9545 let label = canonical_label(label.clone());
9546 final_buckets.write_label(&label)?;
9547 direct_local_unitigs += 1;
9548 }
9549 continue;
9550 };
9551 if usize::try_from(record.unitig_index).ok() != Some(unitig_index) {
9552 return Err(SerialCollationError::MalformedCoordBucket(
9553 inputs.unitig_path.clone(),
9554 ));
9555 }
9556
9557 let bucket_id = stitched_coord_bucket(record.path_id, bucket_mask);
9558 writers.write_materialized_record(bucket_id, &record, &label)?;
9559 }
9560
9561 Ok(ExternalCppPathInfoMaterialized {
9562 manifest: writers.finish()?,
9563 retained: Vec::new(),
9564 direct_local_unitigs,
9565 direct_local_unitigs_complete: true,
9566 })
9567}
9568
9569fn phantom_unitig_label<const K: usize>(endpoint: DiscontinuityEndpoint<K>) -> Vec<u8> {
9570 let label = endpoint.vertex.to_ascii_string().into_bytes();
9571 if endpoint.side == Side::Front {
9572 label
9573 } else {
9574 reverse_complement_label(&label)
9575 }
9576}
9577
9578fn emit_external_ordered_stitched_labels_to_final_buckets<const K: usize>(
9579 inputs: &ExternalDiscontinuityInputs<K>,
9580 coord_dir: &Path,
9581 label_refs: &[ExternalLabelRef],
9582 half_ends: &[HalfEnd],
9583 join_neighbor: &[u32],
9584 starts: &[usize],
9585 final_buckets: &mut FinalUnitigBucketWriters,
9586) -> Result<u64, SerialCollationError> {
9587 const NO_RECORD_INDEX: u32 = u32::MAX;
9588 let map_started = Instant::now();
9589 let mut record_index_by_unitig = vec![NO_RECORD_INDEX; inputs.unitig_count()];
9590 let mut path_offsets = Vec::with_capacity(starts.len() + 1);
9591 let mut records = Vec::<StitchedCoordRecord>::new();
9592 let mut path_records = Vec::new();
9593
9594 for (path_id, &start) in starts.iter().enumerate() {
9595 path_offsets.push(records.len());
9596 path_records.clear();
9597 walk_simple_stitched_component_coords(
9598 half_ends,
9599 join_neighbor,
9600 start,
9601 path_id as u64,
9602 &mut path_records,
9603 );
9604 for record in &path_records {
9605 let unitig_index = record.unitig_index as usize;
9606 let slot = record_index_by_unitig
9607 .get_mut(unitig_index)
9608 .ok_or_else(|| {
9609 SerialCollationError::MalformedCoordBucket(inputs.unitig_path.clone())
9610 })?;
9611 if *slot != NO_RECORD_INDEX {
9612 return Err(SerialCollationError::MalformedCoordBucket(
9613 inputs.unitig_path.clone(),
9614 ));
9615 }
9616 *slot = u32::try_from(records.len()).map_err(|_| {
9617 SerialCollationError::MalformedCoordBucket(inputs.unitig_path.clone())
9618 })?;
9619 records.push(*record);
9620 }
9621 }
9622 path_offsets.push(records.len());
9623 let map_elapsed = map_started.elapsed();
9624
9625 let offset_started = Instant::now();
9626 let mut record_label_offsets = Vec::with_capacity(records.len());
9627 let mut ordered_label_bytes = 0u64;
9628 for record in &records {
9629 let label_ref = label_refs
9630 .get(record.unitig_index as usize)
9631 .ok_or_else(|| SerialCollationError::MalformedCoordBucket(inputs.label_path.clone()))?;
9632 record_label_offsets.push(ordered_label_bytes);
9633 ordered_label_bytes += u64::from(label_ref.label_len);
9634 }
9635 let offset_elapsed = offset_started.elapsed();
9636
9637 let spill_started = Instant::now();
9638 fs::create_dir_all(coord_dir).map_err(|source| SerialCollationError::Io {
9639 path: coord_dir.to_path_buf(),
9640 source,
9641 })?;
9642 let ordered_label_path = coord_dir.join("ordered-stitch-labels.bin");
9643 let ordered_label_file =
9644 File::create(&ordered_label_path).map_err(|source| SerialCollationError::Io {
9645 path: ordered_label_path.clone(),
9646 source,
9647 })?;
9648 ordered_label_file
9649 .set_len(ordered_label_bytes)
9650 .map_err(|source| SerialCollationError::Io {
9651 path: ordered_label_path.clone(),
9652 source,
9653 })?;
9654
9655 let source_label_file =
9656 File::open(&inputs.label_path).map_err(|source| SerialCollationError::Io {
9657 path: inputs.label_path.clone(),
9658 source,
9659 })?;
9660 let mut source_labels = BufReader::with_capacity(1024 * 1024, source_label_file);
9661 let mut label = Vec::new();
9662 for (unitig_index, label_ref) in label_refs.iter().copied().enumerate() {
9663 label.resize(label_ref.label_len as usize, 0);
9664 source_labels
9665 .read_exact(&mut label)
9666 .map_err(|source| SerialCollationError::Io {
9667 path: inputs.label_path.clone(),
9668 source,
9669 })?;
9670 let record_index = record_index_by_unitig[unitig_index];
9671 if record_index == NO_RECORD_INDEX {
9672 continue;
9673 }
9674 ordered_label_file
9675 .write_all_at(&label, record_label_offsets[record_index as usize])
9676 .map_err(|source| SerialCollationError::Io {
9677 path: ordered_label_path.clone(),
9678 source,
9679 })?;
9680 }
9681 drop(ordered_label_file);
9682 let spill_elapsed = spill_started.elapsed();
9683
9684 let assemble_started = Instant::now();
9685 let ordered_label_file =
9686 File::open(&ordered_label_path).map_err(|source| SerialCollationError::Io {
9687 path: ordered_label_path.clone(),
9688 source,
9689 })?;
9690 let mut ordered_labels = BufReader::with_capacity(1024 * 1024, ordered_label_file);
9691 let mut emitted = 0u64;
9692 let mut unitig_label = Vec::new();
9693 let mut stitched_label = Vec::new();
9694 for bounds in path_offsets.windows(2) {
9695 let start = bounds[0];
9696 let end = bounds[1];
9697 if start == end {
9698 continue;
9699 }
9700 stitched_label.clear();
9701 let is_cycle = records[start].is_cycle;
9702 for record in &records[start..end] {
9703 let label_ref = label_refs
9704 .get(record.unitig_index as usize)
9705 .ok_or_else(|| {
9706 SerialCollationError::MalformedCoordBucket(inputs.label_path.clone())
9707 })?;
9708 unitig_label.resize(label_ref.label_len as usize, 0);
9709 ordered_labels
9710 .read_exact(&mut unitig_label)
9711 .map_err(|source| SerialCollationError::Io {
9712 path: ordered_label_path.clone(),
9713 source,
9714 })?;
9715 let mut reverse = record.reverse;
9716 if !stitched_label.is_empty()
9717 && !labels_overlap_oriented_fast::<K>(&stitched_label, &unitig_label, reverse)
9718 {
9719 let alternate = oriented_label(&unitig_label, !reverse);
9720 if labels_overlap::<K>(&stitched_label, &alternate) {
9721 reverse = !reverse;
9722 }
9723 }
9724 append_or_init_oriented_fast::<K>(&mut stitched_label, &unitig_label, reverse);
9725 }
9726
9727 if stitched_label.len() >= K {
9728 let label = if is_cycle {
9729 normalize_stitched_cycle::<K>(&stitched_label)
9730 } else {
9731 canonical_label(stitched_label.clone())
9732 };
9733 final_buckets.write_label(&label)?;
9734 emitted += 1;
9735 }
9736 }
9737 let assemble_elapsed = assemble_started.elapsed();
9738
9739 eprintln!(
9740 "cuttlefish: ordered external stitch labels: map {:.3}s, offsets {:.3}s, label spill {:.3}s, assemble {:.3}s, records {}, label bytes {}",
9741 map_elapsed.as_secs_f64(),
9742 offset_elapsed.as_secs_f64(),
9743 spill_elapsed.as_secs_f64(),
9744 assemble_elapsed.as_secs_f64(),
9745 records.len(),
9746 ordered_label_bytes
9747 );
9748 Ok(emitted)
9749}
9750
9751fn emit_external_stitched_labels_to_final_buckets<const K: usize>(
9752 inputs: &ExternalDiscontinuityInputs<K>,
9753 threads: usize,
9754 label_refs: &[ExternalLabelRef],
9755 half_ends: &[HalfEnd],
9756 join_neighbor: &[u32],
9757 starts: &[usize],
9758 final_buckets: &mut FinalUnitigBucketWriters,
9759) -> Result<u64, SerialCollationError> {
9760 let workers = threads.max(1).min(starts.len().max(1));
9761 if workers == 1 || starts.len() < 1024 {
9762 let label_file =
9763 File::open(&inputs.label_path).map_err(|source| SerialCollationError::Io {
9764 path: inputs.label_path.clone(),
9765 source,
9766 })?;
9767 let mut emitted = 0u64;
9768 for &start in starts {
9769 if let Some(label) = walk_external_simple_stitched_component_label::<K>(
9770 &label_file,
9771 &inputs.label_path,
9772 label_refs,
9773 half_ends,
9774 join_neighbor,
9775 start,
9776 )? {
9777 final_buckets.write_label(&label)?;
9778 emitted += 1;
9779 }
9780 }
9781 return Ok(emitted);
9782 }
9783
9784 const STITCH_LABEL_BATCH: usize = 256;
9785 let chunk_size = starts.len().div_ceil(workers);
9786 let (tx, rx) = mpsc::sync_channel::<Result<Vec<Vec<u8>>, SerialCollationError>>(workers * 2);
9787 let mut emitted = 0u64;
9788 std::thread::scope(|scope| {
9789 let mut handles = Vec::new();
9790 for chunk in starts.chunks(chunk_size) {
9791 let tx = tx.clone();
9792 handles.push(scope.spawn(move || {
9793 let label_file =
9794 File::open(&inputs.label_path).map_err(|source| SerialCollationError::Io {
9795 path: inputs.label_path.clone(),
9796 source,
9797 })?;
9798 let mut batch = Vec::with_capacity(STITCH_LABEL_BATCH);
9799 for &start in chunk {
9800 if let Some(label) = walk_external_simple_stitched_component_label::<K>(
9801 &label_file,
9802 &inputs.label_path,
9803 label_refs,
9804 half_ends,
9805 join_neighbor,
9806 start,
9807 )? {
9808 batch.push(label);
9809 if batch.len() == STITCH_LABEL_BATCH {
9810 let full = std::mem::take(&mut batch);
9811 if tx.send(Ok(full)).is_err() {
9812 return Ok(());
9813 }
9814 }
9815 }
9816 }
9817 if !batch.is_empty() {
9818 let _ = tx.send(Ok(batch));
9819 }
9820 Ok::<_, SerialCollationError>(())
9821 }));
9822 }
9823 drop(tx);
9824
9825 let mut first_error = None;
9826 for result in rx {
9827 match result {
9828 Ok(labels) if first_error.is_none() => {
9829 for label in labels {
9830 if let Err(err) = final_buckets.write_label(&label) {
9831 first_error = Some(err);
9832 break;
9833 }
9834 emitted += 1;
9835 }
9836 }
9837 Ok(_) => {}
9838 Err(err) if first_error.is_none() => first_error = Some(err),
9839 Err(_) => {}
9840 }
9841 }
9842
9843 for handle in handles {
9844 if let Err(err) = handle
9845 .join()
9846 .map_err(|_| SerialCollationError::WorkerPanic)?
9847 {
9848 if first_error.is_none() {
9849 first_error = Some(err);
9850 }
9851 }
9852 }
9853
9854 if let Some(err) = first_error {
9855 Err(err)
9856 } else {
9857 Ok(())
9858 }
9859 })?;
9860
9861 Ok(emitted)
9862}
9863
9864fn walk_external_simple_stitched_component_label<const K: usize>(
9865 label_file: &File,
9866 label_path: &Path,
9867 label_refs: &[ExternalLabelRef],
9868 half_ends: &[HalfEnd],
9869 join_neighbor: &[u32],
9870 start: usize,
9871) -> Result<Option<Vec<u8>>, SerialCollationError> {
9872 let mut current = start;
9873 let mut label = Vec::new();
9874 let mut unitig_label = Vec::new();
9875 let mut is_cycle = false;
9876
9877 loop {
9878 let unitig_index = half_ends[current].unitig_index();
9879 let label_ref = label_refs
9880 .get(unitig_index)
9881 .ok_or_else(|| SerialCollationError::MalformedCoordBucket(label_path.to_path_buf()))?;
9882 read_external_label_ref(label_file, label_path, *label_ref, &mut unitig_label)?;
9883 let reverse = reverse_for_stitch_node(current);
9884 let mut append_reverse = reverse;
9885 if !label.is_empty() && !labels_overlap_oriented_fast::<K>(&label, &unitig_label, reverse) {
9886 let alternate = oriented_label(&unitig_label, !reverse);
9887 if labels_overlap::<K>(&label, &alternate) {
9888 append_reverse = !reverse;
9889 }
9890 }
9891 append_or_init_oriented_fast::<K>(&mut label, &unitig_label, append_reverse);
9892
9893 let other = current ^ 1;
9894 let next = join_neighbor[other];
9895 if next == STITCH_NO_NODE {
9896 break;
9897 }
9898 if stitch_node_index(next) == start {
9899 is_cycle = true;
9900 break;
9901 }
9902 current = stitch_node_index(next);
9903 }
9904
9905 if label.len() < K {
9906 Ok(None)
9907 } else if is_cycle {
9908 Ok(Some(normalize_stitched_cycle::<K>(&label)))
9909 } else {
9910 Ok(Some(canonical_label(label)))
9911 }
9912}
9913
9914fn read_external_label_ref(
9915 label_file: &File,
9916 label_path: &Path,
9917 label_ref: ExternalLabelRef,
9918 scratch: &mut Vec<u8>,
9919) -> Result<(), SerialCollationError> {
9920 scratch.resize(label_ref.label_len as usize, 0);
9921 label_file
9922 .read_exact_at(scratch, label_ref.label_start)
9923 .map_err(|source| SerialCollationError::Io {
9924 path: label_path.to_path_buf(),
9925 source,
9926 })
9927}
9928
9929fn stitch_discontinuity_paths_impl<const K: usize>(
9930 inputs: &DiscontinuityInputs<K>,
9931 skip_unitigs: &[bool],
9932 threads: usize,
9933 coord_dir: Option<&Path>,
9934) -> Result<Vec<Vec<u8>>, SerialCollationError> {
9935 let debug_stitch = false;
9936 let use_adjacency_stitch = false;
9937 let endpoint_dir = coord_dir.map(|dir| dir.join("endpoints"));
9938 let mut endpoint_writers = if !use_adjacency_stitch {
9939 endpoint_dir
9940 .as_deref()
9941 .map(|dir| StitchEndpointBucketWriters::create(dir, threads))
9942 .transpose()?
9943 } else {
9944 None
9945 };
9946 let build_started = Instant::now();
9947 let mut half_ends = Vec::<HalfEnd>::new();
9948 let mut segment_edges = if use_adjacency_stitch {
9949 Some(Vec::<(usize, usize, usize)>::new())
9950 } else {
9951 None
9952 };
9953 let mut endpoint_records = if endpoint_writers.is_none() {
9954 Some(Vec::<StitchEndpointRecord<K>>::new())
9955 } else {
9956 None
9957 };
9958
9959 for (unitig_index, unitig) in inputs.unitigs.iter().enumerate() {
9960 if skip_unitigs.get(unitig_index).copied().unwrap_or(false) {
9961 continue;
9962 }
9963 if unitig.left_exit().is_none() && unitig.right_exit().is_none() {
9964 continue;
9965 }
9966
9967 let (left_endpoint, right_endpoint) = endpoints_by_label_end(unitig);
9968 let left_node = half_ends.len();
9969 half_ends.push(HalfEnd {
9970 unitig_index: u32::try_from(unitig_index).expect("unitig index exceeds u32"),
9971 });
9972 if let Some(endpoint) = left_endpoint {
9973 let record = StitchEndpointRecord {
9974 vertex: endpoint.vertex,
9975 side: endpoint.side,
9976 node: stitch_node(left_node),
9977 };
9978 if let Some(writers) = &mut endpoint_writers {
9979 writers.write_record(&record)?;
9980 } else {
9981 endpoint_records
9982 .as_mut()
9983 .expect("in-memory endpoint records are enabled")
9984 .push(record);
9985 }
9986 }
9987 let right_node = half_ends.len();
9988 half_ends.push(HalfEnd {
9989 unitig_index: u32::try_from(unitig_index).expect("unitig index exceeds u32"),
9990 });
9991 if let Some(endpoint) = right_endpoint {
9992 let record = StitchEndpointRecord {
9993 vertex: endpoint.vertex,
9994 side: endpoint.side,
9995 node: stitch_node(right_node),
9996 };
9997 if let Some(writers) = &mut endpoint_writers {
9998 writers.write_record(&record)?;
9999 } else {
10000 endpoint_records
10001 .as_mut()
10002 .expect("in-memory endpoint records are enabled")
10003 .push(record);
10004 }
10005 }
10006 if let Some(segment_edges) = &mut segment_edges {
10007 segment_edges.push((left_node, right_node, unitig_index));
10008 }
10009 }
10010 let half_end_elapsed = build_started.elapsed();
10011
10012 if debug_stitch {
10013 eprintln!("stitch inputs: {} unitigs", inputs.unitigs.len());
10014 for (index, unitig) in inputs.unitigs.iter().enumerate() {
10015 let (left_endpoint, right_endpoint) = endpoints_by_label_end(unitig);
10016 eprintln!(
10017 "unitig {index}: label={} left_exit={:?} right_exit={:?} label_left={:?} label_right={:?}",
10018 String::from_utf8_lossy(unitig.label(inputs)),
10019 unitig.left_exit().map(debug_endpoint),
10020 unitig.right_exit().map(debug_endpoint),
10021 left_endpoint.map(debug_endpoint),
10022 right_endpoint.map(debug_endpoint),
10023 );
10024 }
10025 }
10026
10027 if use_adjacency_stitch {
10028 return stitch_discontinuity_paths_with_adjacency(
10029 inputs,
10030 threads,
10031 coord_dir,
10032 debug_stitch,
10033 half_ends,
10034 segment_edges.expect("adjacency stitching requires segment edges"),
10035 endpoint_records.expect("adjacency stitching uses in-memory endpoint records"),
10036 half_end_elapsed,
10037 );
10038 }
10039
10040 let mut join_neighbor = vec![STITCH_NO_NODE; half_ends.len()];
10041
10042 let endpoint_sort_started = Instant::now();
10043 let endpoint_manifest = endpoint_writers
10044 .map(StitchEndpointBucketWriters::finish)
10045 .transpose()?;
10046 let mut endpoint_records = endpoint_records.unwrap_or_default();
10047 if endpoint_manifest.is_none() {
10048 endpoint_records.sort_unstable_by_key(|record| record.vertex.as_u128());
10049 }
10050 let endpoint_sort_elapsed = endpoint_sort_started.elapsed();
10051 let endpoint_join_started = Instant::now();
10052 if let Some(manifest) = &endpoint_manifest {
10053 join_neighbors_from_endpoint_buckets::<K>(manifest, &mut join_neighbor, threads)?;
10054 } else {
10055 join_neighbors_from_sorted_endpoints(&endpoint_records, &mut join_neighbor);
10056 }
10057 let endpoint_join_elapsed = endpoint_join_started.elapsed();
10058 drop(endpoint_records);
10059 trim_process_allocations();
10060
10061 if debug_stitch {
10062 for (node, half_end) in half_ends.iter().enumerate() {
10063 let (left_endpoint, right_endpoint) =
10064 endpoints_by_label_end(&inputs.unitigs[half_end.unitig_index()]);
10065 let label_end = label_end_for_node(node);
10066 let endpoint = match label_end {
10067 LabelEnd::Left => left_endpoint,
10068 LabelEnd::Right => right_endpoint,
10069 };
10070 eprintln!(
10071 "node {node}: unitig={} end={:?} endpoint={:?} join={}",
10072 half_end.unitig_index(),
10073 label_end,
10074 endpoint.map(debug_endpoint),
10075 join_neighbor[node],
10076 );
10077 }
10078 }
10079
10080 let component_started = Instant::now();
10081 let component_starts = stitch_component_starts_from_neighbors(&half_ends, &join_neighbor);
10082 let component_elapsed = component_started.elapsed();
10083 let walk_started = Instant::now();
10084 let mut unitigs = if let Some(coord_dir) = coord_dir {
10085 let manifest = write_stitched_coord_buckets(
10086 coord_dir,
10087 threads,
10088 &half_ends,
10089 &join_neighbor,
10090 &component_starts,
10091 )?;
10092 reduce_stitched_coord_bucket_files::<K>(inputs, &manifest, threads)?
10093 } else {
10094 stitch_simple_components_labels_with_threads(
10095 inputs,
10096 &half_ends,
10097 &join_neighbor,
10098 &component_starts,
10099 threads,
10100 )
10101 };
10102 let walk_elapsed = walk_started.elapsed();
10103
10104 let sort_started = Instant::now();
10105 unitigs = bucketed_maximal_unitig_reduce(unitigs, threads);
10106 let sort_elapsed = sort_started.elapsed();
10107 eprintln!(
10108 "cuttlefish: stitch detail: half-ends {:.3}s, endpoint sort {:.3}s, endpoint join {:.3}s, components {:.3}s, component walk {:.3}s, bucket reduce {:.3}s",
10109 half_end_elapsed.as_secs_f64(),
10110 endpoint_sort_elapsed.as_secs_f64(),
10111 endpoint_join_elapsed.as_secs_f64(),
10112 component_elapsed.as_secs_f64(),
10113 walk_elapsed.as_secs_f64(),
10114 sort_elapsed.as_secs_f64()
10115 );
10116 if debug_stitch {
10117 for unitig in &unitigs {
10118 eprintln!("stitched: {}", String::from_utf8_lossy(unitig));
10119 }
10120 }
10121 Ok(unitigs)
10122}
10123
10124#[allow(clippy::too_many_arguments)]
10125fn stitch_discontinuity_paths_with_adjacency<const K: usize>(
10126 inputs: &DiscontinuityInputs<K>,
10127 threads: usize,
10128 coord_dir: Option<&Path>,
10129 debug_stitch: bool,
10130 half_ends: Vec<HalfEnd>,
10131 segment_edges: Vec<(usize, usize, usize)>,
10132 mut endpoint_records: Vec<StitchEndpointRecord<K>>,
10133 half_end_elapsed: Duration,
10134) -> Result<Vec<Vec<u8>>, SerialCollationError> {
10135 let mut adjacency = vec![StitchAdjacencyList::default(); half_ends.len()];
10136 for &(left_node, right_node, unitig_index) in &segment_edges {
10137 push_stitch_edge(&mut adjacency, left_node, right_node, Some(unitig_index));
10138 }
10139
10140 let endpoint_sort_started = Instant::now();
10141 endpoint_records.sort_unstable_by_key(|record| record.vertex.as_u128());
10142 let endpoint_sort_elapsed = endpoint_sort_started.elapsed();
10143 let endpoint_join_started = Instant::now();
10144 let mut start = 0;
10145 while start < endpoint_records.len() {
10146 let vertex = endpoint_records[start].vertex;
10147 let mut end = start + 1;
10148 while end < endpoint_records.len() && endpoint_records[end].vertex == vertex {
10149 end += 1;
10150 }
10151
10152 let mut ends = StitchVertexEnds::default();
10153 for record in &endpoint_records[start..end] {
10154 ends.push(record.side, record.node);
10155 }
10156 if ends.front_len == 1 && ends.back_len == 1 {
10157 push_stitch_edge(
10158 &mut adjacency,
10159 stitch_node_index(ends.fronts[0]),
10160 stitch_node_index(ends.backs[0]),
10161 None,
10162 );
10163 } else if ends.total_len == 2 {
10164 push_stitch_edge(
10165 &mut adjacency,
10166 stitch_node_index(ends.nodes[0]),
10167 stitch_node_index(ends.nodes[1]),
10168 None,
10169 );
10170 }
10171 start = end;
10172 }
10173 let endpoint_join_elapsed = endpoint_join_started.elapsed();
10174 drop(endpoint_records);
10175 trim_process_allocations();
10176
10177 if debug_stitch {
10178 for (node, half_end) in half_ends.iter().enumerate() {
10179 let (left_endpoint, right_endpoint) =
10180 endpoints_by_label_end(&inputs.unitigs[half_end.unitig_index()]);
10181 let label_end = label_end_for_node(node);
10182 let endpoint = match label_end {
10183 LabelEnd::Left => left_endpoint,
10184 LabelEnd::Right => right_endpoint,
10185 };
10186 let edges = adjacency[node]
10187 .iter()
10188 .map(|edge| format!("{}:{:?}", edge.to, edge.unitig_index))
10189 .collect::<Vec<_>>()
10190 .join(",");
10191 eprintln!(
10192 "node {node}: unitig={} end={:?} endpoint={:?} edges=[{}]",
10193 half_end.unitig_index(),
10194 label_end,
10195 endpoint.map(debug_endpoint),
10196 edges,
10197 );
10198 }
10199 }
10200
10201 let mut unitigs = Vec::new();
10202 let component_started = Instant::now();
10203 let component_starts = stitch_component_starts_from_adjacency(&adjacency);
10204 let component_elapsed = component_started.elapsed();
10205 let walk_started = Instant::now();
10206 let used_simple_components = component_starts.is_some();
10207 if let Some(starts) = component_starts {
10208 unitigs = if let Some(coord_dir) = coord_dir {
10209 let manifest = write_stitched_coord_buckets_from_adjacency(
10210 coord_dir, threads, &adjacency, &starts,
10211 )?;
10212 reduce_stitched_coord_bucket_files::<K>(inputs, &manifest, threads)?
10213 } else {
10214 let records = map_simple_stitched_components_with_threads_from_adjacency(
10215 &adjacency, &starts, threads,
10216 );
10217 reduce_stitched_coord_buckets::<K>(inputs, records, threads)
10218 };
10219 } else {
10220 let mut visited_segments = vec![false; inputs.unitigs.len()];
10221
10222 for start in (0..half_ends.len()).filter(|&node| adjacency[node].len() == 1) {
10223 let Some(path) = walk_stitched_path(inputs, &adjacency, &mut visited_segments, start)
10224 else {
10225 continue;
10226 };
10227 push_stitched_path_label::<K>(&mut unitigs, path);
10228 }
10229
10230 for &(left_node, _, unitig_index) in &segment_edges {
10231 if visited_segments[unitig_index] {
10232 continue;
10233 }
10234 let Some(path) =
10235 walk_stitched_path(inputs, &adjacency, &mut visited_segments, left_node)
10236 else {
10237 continue;
10238 };
10239 push_stitched_path_label::<K>(&mut unitigs, path);
10240 }
10241 }
10242 let walk_elapsed = walk_started.elapsed();
10243
10244 let sort_started = Instant::now();
10245 unitigs = bucketed_maximal_unitig_reduce(unitigs, threads);
10246 let sort_elapsed = sort_started.elapsed();
10247 eprintln!(
10248 "cuttlefish: stitch detail: half-ends {:.3}s, endpoint sort {:.3}s, endpoint join {:.3}s, components {:.3}s, {} walk {:.3}s, bucket reduce {:.3}s",
10249 half_end_elapsed.as_secs_f64(),
10250 endpoint_sort_elapsed.as_secs_f64(),
10251 endpoint_join_elapsed.as_secs_f64(),
10252 component_elapsed.as_secs_f64(),
10253 if used_simple_components {
10254 "component"
10255 } else {
10256 "fallback"
10257 },
10258 walk_elapsed.as_secs_f64(),
10259 sort_elapsed.as_secs_f64()
10260 );
10261 if debug_stitch {
10262 for unitig in &unitigs {
10263 eprintln!("stitched: {}", String::from_utf8_lossy(unitig));
10264 }
10265 }
10266 Ok(unitigs)
10267}
10268
10269fn join_neighbors_from_endpoint_buckets<const K: usize>(
10270 manifest: &[StitchEndpointBucketEntry],
10271 join_neighbor: &mut [u32],
10272 threads: usize,
10273) -> Result<(), SerialCollationError> {
10274 let workers = threads.max(1).min(manifest.len().max(1));
10275 if workers == 1 || manifest.len() < 2 {
10276 for entry in manifest {
10277 let mut records = read_stitch_endpoint_bucket::<K>(entry)?;
10278 records.sort_unstable_by_key(|record| record.vertex.as_u128());
10279 join_neighbors_from_sorted_endpoint_assignments(&records, join_neighbor);
10280 }
10281 return Ok(());
10282 }
10283
10284 let next_entry = AtomicUsize::new(0);
10285 let join_writer = JoinNeighborWriter::new(join_neighbor);
10286 std::thread::scope(|scope| {
10287 let mut handles = Vec::new();
10288 for _ in 0..workers {
10289 let next_entry = &next_entry;
10290 let join_writer = &join_writer;
10291 handles.push(scope.spawn(move || {
10292 loop {
10293 let entry_idx = next_entry.fetch_add(1, Ordering::Relaxed);
10294 let Some(entry) = manifest.get(entry_idx) else {
10295 break;
10296 };
10297 let mut records = read_stitch_endpoint_bucket::<K>(entry)?;
10298 records.sort_unstable_by_key(|record| record.vertex.as_u128());
10299 join_writer.write_sorted_endpoint_assignments(&records);
10300 }
10301 Ok::<_, SerialCollationError>(())
10302 }));
10303 }
10304
10305 for handle in handles {
10306 handle
10307 .join()
10308 .map_err(|_| SerialCollationError::WorkerPanic)??;
10309 }
10310 Ok::<_, SerialCollationError>(())
10311 })
10312}
10313
10314struct JoinNeighborWriter {
10315 ptr: *mut u32,
10316 len: usize,
10317}
10318
10319unsafe impl Sync for JoinNeighborWriter {}
10323
10324impl JoinNeighborWriter {
10325 fn new(join_neighbor: &mut [u32]) -> Self {
10326 Self {
10327 ptr: join_neighbor.as_mut_ptr(),
10328 len: join_neighbor.len(),
10329 }
10330 }
10331
10332 fn write_sorted_endpoint_assignments<const K: usize>(
10333 &self,
10334 endpoint_records: &[StitchEndpointRecord<K>],
10335 ) {
10336 let mut start = 0;
10337 while start < endpoint_records.len() {
10338 let vertex = endpoint_records[start].vertex;
10339 let mut end = start + 1;
10340 while end < endpoint_records.len() && endpoint_records[end].vertex == vertex {
10341 end += 1;
10342 }
10343
10344 let mut ends = StitchVertexEnds::default();
10345 for record in &endpoint_records[start..end] {
10346 ends.push(record.side, record.node);
10347 }
10348 if ends.front_len == 1 && ends.back_len == 1 {
10349 self.write(ends.fronts[0], ends.backs[0]);
10350 self.write(ends.backs[0], ends.fronts[0]);
10351 } else if ends.total_len == 2 {
10352 self.write(ends.nodes[0], ends.nodes[1]);
10353 self.write(ends.nodes[1], ends.nodes[0]);
10354 }
10355 start = end;
10356 }
10357 }
10358
10359 fn write(&self, node: u32, neighbor: u32) {
10360 let index = stitch_node_index(node);
10361 debug_assert!(index < self.len);
10362 if index < self.len {
10363 unsafe {
10364 *self.ptr.add(index) = neighbor;
10365 }
10366 }
10367 }
10368}
10369
10370fn read_stitch_endpoint_bucket<const K: usize>(
10371 entry: &StitchEndpointBucketEntry,
10372) -> Result<Vec<StitchEndpointRecord<K>>, SerialCollationError> {
10373 let file = File::open(&entry.path).map_err(|source| SerialCollationError::Io {
10374 path: entry.path.clone(),
10375 source,
10376 })?;
10377 let mut input = BufReader::with_capacity(1024 * 1024, file);
10378 let mut records = Vec::with_capacity(entry.records as usize);
10379 let mut vertex_bytes = [0u8; 16];
10380 let mut side_byte = [0u8; 1];
10381 let mut node_bytes = [0u8; 4];
10382 for _ in 0..entry.records {
10383 input
10384 .read_exact(&mut vertex_bytes)
10385 .and_then(|_| input.read_exact(&mut side_byte))
10386 .and_then(|_| input.read_exact(&mut node_bytes))
10387 .map_err(|source| SerialCollationError::Io {
10388 path: entry.path.clone(),
10389 source,
10390 })?;
10391 let side = match side_byte[0] {
10392 0 => Side::Front,
10393 1 => Side::Back,
10394 _ => {
10395 return Err(SerialCollationError::MalformedCoordBucket(
10396 entry.path.clone(),
10397 ));
10398 }
10399 };
10400 records.push(StitchEndpointRecord {
10401 vertex: Kmer::from_bits(u128::from_le_bytes(vertex_bytes)),
10402 side,
10403 node: u32::from_le_bytes(node_bytes),
10404 });
10405 }
10406
10407 let mut trailing = [0u8; 1];
10408 if input
10409 .read(&mut trailing)
10410 .map_err(|source| SerialCollationError::Io {
10411 path: entry.path.clone(),
10412 source,
10413 })?
10414 != 0
10415 {
10416 return Err(SerialCollationError::MalformedCoordBucket(
10417 entry.path.clone(),
10418 ));
10419 }
10420
10421 Ok(records)
10422}
10423
10424fn join_neighbors_from_sorted_endpoints<const K: usize>(
10425 endpoint_records: &[StitchEndpointRecord<K>],
10426 join_neighbor: &mut [u32],
10427) {
10428 join_neighbors_from_sorted_endpoint_assignments(endpoint_records, join_neighbor);
10429}
10430
10431fn join_neighbors_from_sorted_endpoint_assignments<const K: usize>(
10432 endpoint_records: &[StitchEndpointRecord<K>],
10433 join_neighbor: &mut [u32],
10434) {
10435 let mut start = 0;
10436 while start < endpoint_records.len() {
10437 let vertex = endpoint_records[start].vertex;
10438 let mut end = start + 1;
10439 while end < endpoint_records.len() && endpoint_records[end].vertex == vertex {
10440 end += 1;
10441 }
10442
10443 let mut ends = StitchVertexEnds::default();
10444 for record in &endpoint_records[start..end] {
10445 ends.push(record.side, record.node);
10446 }
10447 if ends.front_len == 1 && ends.back_len == 1 {
10448 join_neighbor[stitch_node_index(ends.fronts[0])] = ends.backs[0];
10449 join_neighbor[stitch_node_index(ends.backs[0])] = ends.fronts[0];
10450 } else if ends.total_len == 2 {
10451 join_neighbor[stitch_node_index(ends.nodes[0])] = ends.nodes[1];
10452 join_neighbor[stitch_node_index(ends.nodes[1])] = ends.nodes[0];
10453 }
10454 start = end;
10455 }
10456}
10457
10458fn stitch_component_starts_from_adjacency(adjacency: &[StitchAdjacencyList]) -> Option<Vec<usize>> {
10459 let mut visited = vec![false; adjacency.len()];
10460 let mut starts = Vec::new();
10461
10462 if adjacency.iter().any(|neighbours| neighbours.len() > 2) {
10463 return None;
10464 }
10465
10466 for node in 0..adjacency.len() {
10467 if visited[node] || adjacency[node].len() != 1 {
10468 continue;
10469 }
10470 starts.push(node);
10471 mark_linear_stitch_component(adjacency, node, &mut visited);
10472 }
10473
10474 for node in 0..adjacency.len() {
10475 if visited[node] || adjacency[node].len() == 0 {
10476 continue;
10477 }
10478 starts.push(node);
10479 mark_linear_stitch_component(adjacency, node, &mut visited);
10480 }
10481
10482 Some(starts)
10483}
10484
10485fn mark_linear_stitch_component(
10486 adjacency: &[StitchAdjacencyList],
10487 start: usize,
10488 visited: &mut [bool],
10489) {
10490 let mut previous = usize::MAX;
10491 let mut current = start;
10492
10493 loop {
10494 if visited[current] {
10495 break;
10496 }
10497 visited[current] = true;
10498
10499 let next = adjacency[current]
10500 .iter()
10501 .map(|edge| edge.to)
10502 .find(|&next| next != previous && !visited[next]);
10503 let Some(next) = next else {
10504 break;
10505 };
10506 previous = current;
10507 current = next;
10508 }
10509}
10510
10511fn stitch_component_starts_from_neighbors(
10512 half_ends: &[HalfEnd],
10513 join_neighbor: &[u32],
10514) -> Vec<usize> {
10515 debug_assert_eq!(half_ends.len(), join_neighbor.len());
10516 debug_assert_eq!(half_ends.len() % 2, 0);
10517 let mut visited = vec![0u8; half_ends.len().div_ceil(2)];
10518 let mut starts = Vec::new();
10519
10520 for node in 0..join_neighbor.len() {
10521 if join_neighbor[node] != STITCH_NO_NODE {
10522 continue;
10523 }
10524 let unitig_pair = node >> 1;
10525 if visited[unitig_pair] != 0 {
10526 continue;
10527 }
10528 mark_neighbor_component(join_neighbor, node, &mut visited);
10529 starts.push(node);
10530 }
10531
10532 for node in (0..join_neighbor.len()).step_by(2) {
10533 let unitig_pair = node >> 1;
10534 if visited[unitig_pair] != 0 {
10535 continue;
10536 }
10537 mark_neighbor_component(join_neighbor, node, &mut visited);
10538 starts.push(node);
10539 }
10540
10541 starts
10542}
10543
10544fn mark_neighbor_component(join_neighbor: &[u32], start: usize, visited: &mut [u8]) {
10545 let mut current = start;
10546 loop {
10547 let unitig_pair = current >> 1;
10548 if visited[unitig_pair] != 0 {
10549 break;
10550 }
10551 visited[unitig_pair] = 1;
10552
10553 let other = current ^ 1;
10554 let next = join_neighbor[other];
10555 if next == STITCH_NO_NODE || stitch_node_index(next) == start {
10556 break;
10557 }
10558 current = stitch_node_index(next);
10559 }
10560}
10561
10562fn push_stitched_path_label<const K: usize>(unitigs: &mut Vec<Vec<u8>>, path: StitchedPath) {
10563 let label = path.label;
10564 if label.len() >= K {
10565 if path.is_cycle {
10566 unitigs.push(normalize_stitched_cycle::<K>(&label));
10567 } else {
10568 unitigs.push(canonical_label(label));
10569 }
10570 }
10571}
10572
10573fn stitch_simple_components_labels_with_threads<const K: usize>(
10574 inputs: &DiscontinuityInputs<K>,
10575 half_ends: &[HalfEnd],
10576 join_neighbor: &[u32],
10577 starts: &[usize],
10578 threads: usize,
10579) -> Vec<Vec<u8>> {
10580 let workers = threads.max(1).min(starts.len().max(1));
10581 if workers == 1 || starts.len() < 1024 {
10582 let mut unitigs = Vec::new();
10583 for &start in starts {
10584 if let Some(label) =
10585 walk_simple_stitched_component_label(inputs, half_ends, join_neighbor, start)
10586 {
10587 unitigs.push(label);
10588 }
10589 }
10590 return unitigs;
10591 }
10592
10593 let chunk_size = starts.len().div_ceil(workers);
10594 let worker_unitigs = std::thread::scope(|scope| {
10595 let mut handles = Vec::new();
10596 for chunk in starts.chunks(chunk_size) {
10597 handles.push(scope.spawn(move || {
10598 let mut unitigs = Vec::new();
10599 for &start in chunk {
10600 if let Some(label) = walk_simple_stitched_component_label(
10601 inputs,
10602 half_ends,
10603 join_neighbor,
10604 start,
10605 ) {
10606 unitigs.push(label);
10607 }
10608 }
10609 unitigs
10610 }));
10611 }
10612
10613 let mut worker_unitigs = Vec::new();
10614 for handle in handles {
10615 worker_unitigs.push(handle.join().expect("stitch label worker panicked"));
10616 }
10617 worker_unitigs
10618 });
10619
10620 let total = worker_unitigs.iter().map(Vec::len).sum();
10621 let mut unitigs = Vec::with_capacity(total);
10622 for mut worker in worker_unitigs {
10623 unitigs.append(&mut worker);
10624 }
10625 unitigs
10626}
10627
10628fn walk_simple_stitched_component_label<const K: usize>(
10629 inputs: &DiscontinuityInputs<K>,
10630 half_ends: &[HalfEnd],
10631 join_neighbor: &[u32],
10632 start: usize,
10633) -> Option<Vec<u8>> {
10634 let mut current = start;
10635 let mut label = Vec::new();
10636 let mut is_cycle = false;
10637
10638 loop {
10639 let unitig = &inputs.unitigs[half_ends[current].unitig_index()];
10640 let unitig_label = unitig.label(inputs);
10641 let reverse = reverse_for_stitch_node(current);
10642 let mut append_reverse = reverse;
10643 if !label.is_empty() && !labels_overlap_oriented_fast::<K>(&label, unitig_label, reverse) {
10644 let alternate = oriented_label(unitig_label, !reverse);
10645 if labels_overlap::<K>(&label, &alternate) {
10646 append_reverse = !reverse;
10647 }
10648 }
10649 append_or_init_oriented_fast::<K>(&mut label, unitig_label, append_reverse);
10650
10651 let other = current ^ 1;
10652 let next = join_neighbor[other];
10653 if next == STITCH_NO_NODE {
10654 break;
10655 }
10656 if stitch_node_index(next) == start {
10657 is_cycle = true;
10658 break;
10659 }
10660 current = stitch_node_index(next);
10661 }
10662
10663 if label.len() < K {
10664 None
10665 } else if is_cycle {
10666 Some(normalize_stitched_cycle::<K>(&label))
10667 } else {
10668 Some(canonical_label(label))
10669 }
10670}
10671
10672fn map_simple_stitched_components_with_threads_from_adjacency(
10673 adjacency: &[StitchAdjacencyList],
10674 starts: &[usize],
10675 threads: usize,
10676) -> Vec<StitchedCoordRecord> {
10677 let workers = threads.max(1).min(starts.len().max(1));
10678 if workers == 1 || starts.len() < 1024 {
10679 let mut records = Vec::new();
10680 for (path_id, &start) in starts.iter().enumerate() {
10681 walk_simple_stitched_component_coords_from_adjacency(
10682 adjacency,
10683 start,
10684 path_id as u64,
10685 &mut records,
10686 );
10687 }
10688 return records;
10689 }
10690
10691 let chunk_size = starts.len().div_ceil(workers);
10692 std::thread::scope(|scope| {
10693 let mut handles = Vec::new();
10694 for (chunk_index, chunk) in starts.chunks(chunk_size).enumerate() {
10695 let base_path_id = (chunk_index * chunk_size) as u64;
10696 handles.push(scope.spawn(move || {
10697 let mut records = Vec::new();
10698 for (offset, &start) in chunk.iter().enumerate() {
10699 walk_simple_stitched_component_coords_from_adjacency(
10700 adjacency,
10701 start,
10702 base_path_id + offset as u64,
10703 &mut records,
10704 );
10705 }
10706 records
10707 }));
10708 }
10709
10710 let mut records = Vec::with_capacity(starts.len());
10711 for handle in handles {
10712 records.extend(handle.join().expect("stitch mapper worker panicked"));
10713 }
10714 records
10715 })
10716}
10717
10718fn write_stitched_coord_buckets(
10719 coord_dir: &Path,
10720 threads: usize,
10721 half_ends: &[HalfEnd],
10722 join_neighbor: &[u32],
10723 starts: &[usize],
10724) -> Result<Vec<StitchedCoordBucketEntry>, SerialCollationError> {
10725 if coord_dir.exists() {
10726 fs::remove_dir_all(coord_dir).map_err(|source| SerialCollationError::Io {
10727 path: coord_dir.to_path_buf(),
10728 source,
10729 })?;
10730 }
10731 fs::create_dir_all(coord_dir).map_err(|source| SerialCollationError::Io {
10732 path: coord_dir.to_path_buf(),
10733 source,
10734 })?;
10735
10736 let bucket_count = stitched_coord_bucket_count(threads);
10737 let bucket_mask = bucket_count - 1;
10738 let workers = threads.max(1).min(starts.len().max(1));
10739
10740 let mut manifest = if workers == 1 || starts.len() < 1024 {
10741 write_neighbor_stitched_coord_shards(
10742 coord_dir,
10743 0,
10744 bucket_count,
10745 bucket_mask,
10746 half_ends,
10747 join_neighbor,
10748 starts,
10749 0,
10750 )?
10751 } else {
10752 let chunk_size = starts.len().div_ceil(workers);
10753 std::thread::scope(|scope| {
10754 let mut handles = Vec::new();
10755 for (chunk_index, chunk) in starts.chunks(chunk_size).enumerate() {
10756 let base_path_id = (chunk_index * chunk_size) as u64;
10757 handles.push(scope.spawn(move || {
10758 write_neighbor_stitched_coord_shards(
10759 coord_dir,
10760 chunk_index,
10761 bucket_count,
10762 bucket_mask,
10763 half_ends,
10764 join_neighbor,
10765 chunk,
10766 base_path_id,
10767 )
10768 }));
10769 }
10770
10771 let mut manifest = Vec::new();
10772 for handle in handles {
10773 manifest.extend(
10774 handle
10775 .join()
10776 .map_err(|_| SerialCollationError::WorkerPanic)??,
10777 );
10778 }
10779 Ok::<_, SerialCollationError>(manifest)
10780 })?
10781 };
10782 manifest.sort_by(|left, right| {
10783 left.bucket_id
10784 .cmp(&right.bucket_id)
10785 .then_with(|| left.path.cmp(&right.path))
10786 });
10787
10788 Ok(manifest)
10789}
10790
10791fn write_stitched_coord_buckets_from_adjacency(
10792 coord_dir: &Path,
10793 threads: usize,
10794 adjacency: &[StitchAdjacencyList],
10795 starts: &[usize],
10796) -> Result<Vec<StitchedCoordBucketEntry>, SerialCollationError> {
10797 if coord_dir.exists() {
10798 fs::remove_dir_all(coord_dir).map_err(|source| SerialCollationError::Io {
10799 path: coord_dir.to_path_buf(),
10800 source,
10801 })?;
10802 }
10803 fs::create_dir_all(coord_dir).map_err(|source| SerialCollationError::Io {
10804 path: coord_dir.to_path_buf(),
10805 source,
10806 })?;
10807
10808 let bucket_count = stitched_coord_bucket_count(threads);
10809 let bucket_mask = bucket_count - 1;
10810 let workers = threads.max(1).min(starts.len().max(1));
10811
10812 let mut manifest = if workers == 1 || starts.len() < 1024 {
10813 write_adjacency_stitched_coord_shards(
10814 coord_dir,
10815 0,
10816 bucket_count,
10817 bucket_mask,
10818 adjacency,
10819 starts,
10820 0,
10821 )?
10822 } else {
10823 let chunk_size = starts.len().div_ceil(workers);
10824 std::thread::scope(|scope| {
10825 let mut handles = Vec::new();
10826 for (chunk_index, chunk) in starts.chunks(chunk_size).enumerate() {
10827 let base_path_id = (chunk_index * chunk_size) as u64;
10828 handles.push(scope.spawn(move || {
10829 write_adjacency_stitched_coord_shards(
10830 coord_dir,
10831 chunk_index,
10832 bucket_count,
10833 bucket_mask,
10834 adjacency,
10835 chunk,
10836 base_path_id,
10837 )
10838 }));
10839 }
10840
10841 let mut manifest = Vec::new();
10842 for handle in handles {
10843 manifest.extend(
10844 handle
10845 .join()
10846 .map_err(|_| SerialCollationError::WorkerPanic)??,
10847 );
10848 }
10849 Ok::<_, SerialCollationError>(manifest)
10850 })?
10851 };
10852 manifest.sort_by(|left, right| {
10853 left.bucket_id
10854 .cmp(&right.bucket_id)
10855 .then_with(|| left.path.cmp(&right.path))
10856 });
10857 Ok(manifest)
10858}
10859
10860#[allow(clippy::too_many_arguments)]
10861fn write_neighbor_stitched_coord_shards(
10862 coord_dir: &Path,
10863 worker_id: usize,
10864 bucket_count: usize,
10865 bucket_mask: usize,
10866 half_ends: &[HalfEnd],
10867 join_neighbor: &[u32],
10868 starts: &[usize],
10869 base_path_id: u64,
10870) -> Result<Vec<StitchedCoordBucketEntry>, SerialCollationError> {
10871 let mut writers = StitchedCoordShardWriters::new(coord_dir, worker_id, bucket_count);
10872 let mut records = Vec::new();
10873 for (offset, &start) in starts.iter().enumerate() {
10874 records.clear();
10875 let path_id = base_path_id + offset as u64;
10876 walk_simple_stitched_component_coords(
10877 half_ends,
10878 join_neighbor,
10879 start,
10880 path_id,
10881 &mut records,
10882 );
10883 writers.write_path_records(stitched_coord_bucket(path_id, bucket_mask), &records)?;
10884 }
10885 writers.finish()
10886}
10887
10888fn write_adjacency_stitched_coord_shards(
10889 coord_dir: &Path,
10890 worker_id: usize,
10891 bucket_count: usize,
10892 bucket_mask: usize,
10893 adjacency: &[StitchAdjacencyList],
10894 starts: &[usize],
10895 base_path_id: u64,
10896) -> Result<Vec<StitchedCoordBucketEntry>, SerialCollationError> {
10897 let mut writers = StitchedCoordShardWriters::new(coord_dir, worker_id, bucket_count);
10898 let mut records = Vec::new();
10899 for (offset, &start) in starts.iter().enumerate() {
10900 records.clear();
10901 let path_id = base_path_id + offset as u64;
10902 walk_simple_stitched_component_coords_from_adjacency(
10903 adjacency,
10904 start,
10905 path_id,
10906 &mut records,
10907 );
10908 writers.write_path_records(stitched_coord_bucket(path_id, bucket_mask), &records)?;
10909 }
10910 writers.finish()
10911}
10912
10913fn write_materialized_stitched_coord_buckets<const K: usize>(
10914 inputs: &DiscontinuityInputs<K>,
10915 coord_dir: &Path,
10916 threads: usize,
10917 half_ends: &[HalfEnd],
10918 join_neighbor: &[u32],
10919 starts: &[usize],
10920) -> Result<Vec<MaterializedStitchedCoordBucketEntry>, SerialCollationError> {
10921 if coord_dir.exists() {
10922 fs::remove_dir_all(coord_dir).map_err(|source| SerialCollationError::Io {
10923 path: coord_dir.to_path_buf(),
10924 source,
10925 })?;
10926 }
10927 fs::create_dir_all(coord_dir).map_err(|source| SerialCollationError::Io {
10928 path: coord_dir.to_path_buf(),
10929 source,
10930 })?;
10931
10932 let bucket_count = materialized_stitched_coord_bucket_count(threads);
10933 let bucket_mask = bucket_count - 1;
10934 let workers = threads.max(1).min(starts.len().max(1));
10935
10936 let mut manifest = if workers == 1 || starts.len() < 1024 {
10937 write_materialized_neighbor_stitched_coord_shards(
10938 inputs,
10939 coord_dir,
10940 0,
10941 bucket_count,
10942 bucket_mask,
10943 half_ends,
10944 join_neighbor,
10945 starts,
10946 0,
10947 )?
10948 } else {
10949 let chunk_size = starts.len().div_ceil(workers);
10950 std::thread::scope(|scope| {
10951 let mut handles = Vec::new();
10952 for (chunk_index, chunk) in starts.chunks(chunk_size).enumerate() {
10953 let base_path_id = (chunk_index * chunk_size) as u64;
10954 handles.push(scope.spawn(move || {
10955 write_materialized_neighbor_stitched_coord_shards(
10956 inputs,
10957 coord_dir,
10958 chunk_index,
10959 bucket_count,
10960 bucket_mask,
10961 half_ends,
10962 join_neighbor,
10963 chunk,
10964 base_path_id,
10965 )
10966 }));
10967 }
10968
10969 let mut manifest = Vec::new();
10970 for handle in handles {
10971 manifest.extend(
10972 handle
10973 .join()
10974 .map_err(|_| SerialCollationError::WorkerPanic)??,
10975 );
10976 }
10977 Ok::<_, SerialCollationError>(manifest)
10978 })?
10979 };
10980 manifest.sort_by(|left, right| {
10981 left.bucket_id
10982 .cmp(&right.bucket_id)
10983 .then_with(|| left.coord_path.cmp(&right.coord_path))
10984 });
10985 Ok(manifest)
10986}
10987
10988fn write_external_materialized_stitched_coord_buckets_from_neighbors<const K: usize>(
10989 inputs: &ExternalDiscontinuityInputs<K>,
10990 coord_dir: &Path,
10991 threads: usize,
10992 half_ends: &[HalfEnd],
10993 join_neighbor: &[u32],
10994) -> Result<Vec<MaterializedStitchedCoordBucketEntry>, SerialCollationError> {
10995 let estimated_path_info_bytes = (half_ends.len() / 2)
10996 .checked_mul(std::mem::size_of::<StitchedCoordRecord>())
10997 .ok_or_else(|| SerialCollationError::MalformedCoordBucket(coord_dir.to_path_buf()))?;
10998 let inmem_limit = in_memory_stitch_path_info_limit();
10999 if estimated_path_info_bytes <= inmem_limit {
11000 return write_external_materialized_stitched_coord_buckets_from_neighbor_path_info::<K>(
11001 inputs,
11002 coord_dir,
11003 threads,
11004 half_ends,
11005 join_neighbor,
11006 estimated_path_info_bytes,
11007 );
11008 }
11009
11010 eprintln!(
11011 "cuttlefish: external stitch path-info estimate {:.1} MiB exceeds in-memory limit {:.1} MiB; using disk path-info buckets",
11012 estimated_path_info_bytes as f64 / (1024.0 * 1024.0),
11013 inmem_limit as f64 / (1024.0 * 1024.0)
11014 );
11015
11016 let component_starts = stitch_component_starts_from_neighbors(half_ends, join_neighbor);
11017 write_external_materialized_stitched_coord_buckets(
11018 inputs,
11019 coord_dir,
11020 threads,
11021 half_ends,
11022 join_neighbor,
11023 &component_starts,
11024 )
11025}
11026
11027fn write_external_materialized_stitched_coord_buckets<const K: usize>(
11028 inputs: &ExternalDiscontinuityInputs<K>,
11029 coord_dir: &Path,
11030 threads: usize,
11031 half_ends: &[HalfEnd],
11032 join_neighbor: &[u32],
11033 starts: &[usize],
11034) -> Result<Vec<MaterializedStitchedCoordBucketEntry>, SerialCollationError> {
11035 let estimated_path_info_bytes = (half_ends.len() / 2)
11036 .checked_mul(std::mem::size_of::<StitchedCoordRecord>())
11037 .ok_or_else(|| SerialCollationError::MalformedCoordBucket(coord_dir.to_path_buf()))?;
11038 let inmem_limit = in_memory_stitch_path_info_limit();
11039 if estimated_path_info_bytes <= inmem_limit {
11040 return write_external_materialized_stitched_coord_buckets_in_memory::<K>(
11041 inputs,
11042 coord_dir,
11043 threads,
11044 half_ends,
11045 join_neighbor,
11046 starts,
11047 estimated_path_info_bytes,
11048 );
11049 }
11050
11051 eprintln!(
11052 "cuttlefish: external stitch path-info estimate {:.1} MiB exceeds in-memory limit {:.1} MiB; using disk path-info buckets",
11053 estimated_path_info_bytes as f64 / (1024.0 * 1024.0),
11054 inmem_limit as f64 / (1024.0 * 1024.0)
11055 );
11056
11057 if coord_dir.exists() {
11058 fs::remove_dir_all(coord_dir).map_err(|source| SerialCollationError::Io {
11059 path: coord_dir.to_path_buf(),
11060 source,
11061 })?;
11062 }
11063 fs::create_dir_all(coord_dir).map_err(|source| SerialCollationError::Io {
11064 path: coord_dir.to_path_buf(),
11065 source,
11066 })?;
11067
11068 let path_info_started = Instant::now();
11069 let path_info_manifest = write_external_stitched_path_info_buckets::<K>(
11070 inputs,
11071 coord_dir,
11072 threads,
11073 half_ends,
11074 join_neighbor,
11075 starts,
11076 )?;
11077 let path_info_elapsed = path_info_started.elapsed();
11078
11079 let materialize_started = Instant::now();
11080 let manifest = materialize_external_stitched_coord_buckets_from_path_info::<K>(
11081 inputs,
11082 coord_dir,
11083 threads,
11084 &path_info_manifest,
11085 )?;
11086 eprintln!(
11087 "cuttlefish: external materialized stitch path-info {:.3}s, label materialize {:.3}s",
11088 path_info_elapsed.as_secs_f64(),
11089 materialize_started.elapsed().as_secs_f64()
11090 );
11091 Ok(manifest)
11092}
11093
11094#[allow(clippy::too_many_arguments)]
11095fn write_external_materialized_stitched_coord_buckets_in_memory<const K: usize>(
11096 inputs: &ExternalDiscontinuityInputs<K>,
11097 coord_dir: &Path,
11098 threads: usize,
11099 half_ends: &[HalfEnd],
11100 join_neighbor: &[u32],
11101 starts: &[usize],
11102 estimated_path_info_bytes: usize,
11103) -> Result<Vec<MaterializedStitchedCoordBucketEntry>, SerialCollationError> {
11104 if coord_dir.exists() {
11105 fs::remove_dir_all(coord_dir).map_err(|source| SerialCollationError::Io {
11106 path: coord_dir.to_path_buf(),
11107 source,
11108 })?;
11109 }
11110 fs::create_dir_all(coord_dir).map_err(|source| SerialCollationError::Io {
11111 path: coord_dir.to_path_buf(),
11112 source,
11113 })?;
11114
11115 let ranges_per_bucket = ranges_per_path_info_bucket(inputs.ranges.len(), threads.max(1));
11116 let path_info_bucket_count = inputs.ranges.len().div_ceil(ranges_per_bucket).max(1);
11117 let path_info_started = Instant::now();
11118 let records_by_range_bucket = collect_external_stitched_path_info_by_range_bucket(
11119 inputs,
11120 threads,
11121 half_ends,
11122 join_neighbor,
11123 starts,
11124 ranges_per_bucket,
11125 path_info_bucket_count,
11126 )?;
11127 let path_info_elapsed = path_info_started.elapsed();
11128
11129 let materialize_started = Instant::now();
11130 let manifest = materialize_external_stitched_coord_buckets_from_memory::<K>(
11131 inputs,
11132 coord_dir,
11133 threads,
11134 ranges_per_bucket,
11135 records_by_range_bucket,
11136 )?;
11137 eprintln!(
11138 "cuttlefish: external in-memory stitch path-info {:.3}s ({:.1} MiB est), label materialize {:.3}s",
11139 path_info_elapsed.as_secs_f64(),
11140 estimated_path_info_bytes as f64 / (1024.0 * 1024.0),
11141 materialize_started.elapsed().as_secs_f64()
11142 );
11143 Ok(manifest)
11144}
11145
11146fn write_external_materialized_stitched_coord_buckets_from_neighbor_path_info<const K: usize>(
11147 inputs: &ExternalDiscontinuityInputs<K>,
11148 coord_dir: &Path,
11149 threads: usize,
11150 half_ends: &[HalfEnd],
11151 join_neighbor: &[u32],
11152 estimated_path_info_bytes: usize,
11153) -> Result<Vec<MaterializedStitchedCoordBucketEntry>, SerialCollationError> {
11154 if coord_dir.exists() {
11155 fs::remove_dir_all(coord_dir).map_err(|source| SerialCollationError::Io {
11156 path: coord_dir.to_path_buf(),
11157 source,
11158 })?;
11159 }
11160 fs::create_dir_all(coord_dir).map_err(|source| SerialCollationError::Io {
11161 path: coord_dir.to_path_buf(),
11162 source,
11163 })?;
11164
11165 let ranges_per_bucket = ranges_per_path_info_bucket(inputs.ranges.len(), threads.max(1));
11166 let path_info_bucket_count = inputs.ranges.len().div_ceil(ranges_per_bucket).max(1);
11167 let path_info_started = Instant::now();
11168 let records_by_range_bucket = collect_external_stitched_path_info_from_neighbors(
11169 inputs,
11170 half_ends,
11171 join_neighbor,
11172 ranges_per_bucket,
11173 path_info_bucket_count,
11174 )?;
11175 let path_info_elapsed = path_info_started.elapsed();
11176
11177 let materialize_started = Instant::now();
11178 let manifest = materialize_external_stitched_coord_buckets_from_memory::<K>(
11179 inputs,
11180 coord_dir,
11181 threads,
11182 ranges_per_bucket,
11183 records_by_range_bucket,
11184 )?;
11185 eprintln!(
11186 "cuttlefish: external neighbor path-info {:.3}s ({:.1} MiB est), label materialize {:.3}s",
11187 path_info_elapsed.as_secs_f64(),
11188 estimated_path_info_bytes as f64 / (1024.0 * 1024.0),
11189 materialize_started.elapsed().as_secs_f64()
11190 );
11191 Ok(manifest)
11192}
11193
11194fn in_memory_stitch_path_info_limit() -> usize {
11195 if std::env::var_os("CF3_RS_DISABLE_INMEM_STITCH_PATH_INFO").is_some() {
11196 return 0;
11197 }
11198 std::env::var("CF3_RS_INMEM_STITCH_PATH_INFO_BYTES")
11199 .ok()
11200 .and_then(|value| value.parse::<usize>().ok())
11201 .unwrap_or(DEFAULT_INMEM_STITCH_PATH_INFO_LIMIT)
11202}
11203
11204fn collect_external_stitched_path_info_by_range_bucket<const K: usize>(
11205 inputs: &ExternalDiscontinuityInputs<K>,
11206 threads: usize,
11207 half_ends: &[HalfEnd],
11208 join_neighbor: &[u32],
11209 starts: &[usize],
11210 ranges_per_bucket: usize,
11211 path_info_bucket_count: usize,
11212) -> Result<Vec<Vec<StitchedCoordRecord>>, SerialCollationError> {
11213 let workers = threads.max(1).min(starts.len().max(1));
11214 if workers == 1 || starts.len() < 1024 {
11215 let mut records_by_bucket = empty_stitched_record_buckets(path_info_bucket_count);
11216 collect_external_stitched_path_info_range_into_memory(
11217 &inputs.ranges,
11218 &inputs.unitig_path,
11219 ranges_per_bucket,
11220 half_ends,
11221 join_neighbor,
11222 starts,
11223 0,
11224 &mut records_by_bucket,
11225 )?;
11226 return Ok(records_by_bucket);
11227 }
11228
11229 let chunk_size = starts.len().div_ceil(workers);
11230 std::thread::scope(|scope| {
11231 let mut handles = Vec::new();
11232 for (chunk_index, chunk) in starts.chunks(chunk_size).enumerate() {
11233 let base_path_id = (chunk_index * chunk_size) as u64;
11234 handles.push(scope.spawn(move || {
11235 let mut records_by_bucket = empty_stitched_record_buckets(path_info_bucket_count);
11236 collect_external_stitched_path_info_range_into_memory(
11237 &inputs.ranges,
11238 &inputs.unitig_path,
11239 ranges_per_bucket,
11240 half_ends,
11241 join_neighbor,
11242 chunk,
11243 base_path_id,
11244 &mut records_by_bucket,
11245 )?;
11246 Ok::<_, SerialCollationError>(records_by_bucket)
11247 }));
11248 }
11249
11250 let mut merged = empty_stitched_record_buckets(path_info_bucket_count);
11251 for handle in handles {
11252 let mut worker_buckets = handle
11253 .join()
11254 .map_err(|_| SerialCollationError::WorkerPanic)??;
11255 for (bucket_id, bucket) in worker_buckets.iter_mut().enumerate() {
11256 merged[bucket_id].append(bucket);
11257 }
11258 }
11259 Ok::<_, SerialCollationError>(merged)
11260 })
11261}
11262
11263fn empty_stitched_record_buckets(bucket_count: usize) -> Vec<Vec<StitchedCoordRecord>> {
11264 let mut buckets = Vec::with_capacity(bucket_count);
11265 buckets.resize_with(bucket_count, Vec::new);
11266 buckets
11267}
11268
11269fn collect_external_stitched_path_info_from_neighbors<const K: usize>(
11270 inputs: &ExternalDiscontinuityInputs<K>,
11271 half_ends: &[HalfEnd],
11272 join_neighbor: &[u32],
11273 ranges_per_bucket: usize,
11274 path_info_bucket_count: usize,
11275) -> Result<Vec<Vec<StitchedCoordRecord>>, SerialCollationError> {
11276 debug_assert_eq!(half_ends.len(), join_neighbor.len());
11277 debug_assert_eq!(half_ends.len() % 2, 0);
11278 let mut visited = vec![0u8; half_ends.len().div_ceil(2)];
11279 let mut records_by_bucket = empty_stitched_record_buckets(path_info_bucket_count);
11280 let mut path_records = Vec::new();
11281 let unitig_bucket = unitig_path_info_bucket_map(
11282 &inputs.ranges,
11283 ranges_per_bucket,
11284 path_info_bucket_count,
11285 inputs.unitig_count(),
11286 &inputs.unitig_path,
11287 )?;
11288 let mut path_id = 0u64;
11289
11290 for node in 0..join_neighbor.len() {
11291 if join_neighbor[node] != STITCH_NO_NODE {
11292 continue;
11293 }
11294 let unitig_pair = node >> 1;
11295 if visited[unitig_pair] != 0 {
11296 continue;
11297 }
11298 path_records.clear();
11299 walk_simple_stitched_component_coords_marked(
11300 half_ends,
11301 join_neighbor,
11302 node,
11303 path_id,
11304 &mut visited,
11305 &mut path_records,
11306 );
11307 if !path_records.is_empty() {
11308 push_stitched_records_to_range_buckets(
11309 &unitig_bucket,
11310 &inputs.unitig_path,
11311 &path_records,
11312 &mut records_by_bucket,
11313 )?;
11314 path_id += 1;
11315 }
11316 }
11317
11318 for node in (0..join_neighbor.len()).step_by(2) {
11319 let unitig_pair = node >> 1;
11320 if visited[unitig_pair] != 0 {
11321 continue;
11322 }
11323 path_records.clear();
11324 walk_simple_stitched_component_coords_marked(
11325 half_ends,
11326 join_neighbor,
11327 node,
11328 path_id,
11329 &mut visited,
11330 &mut path_records,
11331 );
11332 if !path_records.is_empty() {
11333 push_stitched_records_to_range_buckets(
11334 &unitig_bucket,
11335 &inputs.unitig_path,
11336 &path_records,
11337 &mut records_by_bucket,
11338 )?;
11339 path_id += 1;
11340 }
11341 }
11342
11343 Ok(records_by_bucket)
11344}
11345
11346fn unitig_path_info_bucket_map(
11347 ranges: &[ExternalLocalUnitigRange],
11348 ranges_per_bucket: usize,
11349 path_info_bucket_count: usize,
11350 unitig_count: usize,
11351 malformed_path: &Path,
11352) -> Result<Vec<u32>, SerialCollationError> {
11353 let bucket_limit = u32::try_from(path_info_bucket_count)
11354 .map_err(|_| SerialCollationError::MalformedCoordBucket(malformed_path.to_path_buf()))?;
11355 let mut unitig_bucket = vec![bucket_limit; unitig_count];
11356 for (range_id, range) in ranges.iter().enumerate() {
11357 let bucket_id = u32::try_from(range_id / ranges_per_bucket).map_err(|_| {
11358 SerialCollationError::MalformedCoordBucket(malformed_path.to_path_buf())
11359 })?;
11360 let end = range
11361 .start_unitig
11362 .checked_add(range.unitigs)
11363 .ok_or_else(|| {
11364 SerialCollationError::MalformedCoordBucket(malformed_path.to_path_buf())
11365 })?;
11366 if end > unitig_bucket.len() {
11367 return Err(SerialCollationError::MalformedCoordBucket(
11368 malformed_path.to_path_buf(),
11369 ));
11370 }
11371 unitig_bucket[range.start_unitig..end].fill(bucket_id);
11372 }
11373 Ok(unitig_bucket)
11374}
11375
11376fn walk_simple_stitched_component_coords_marked(
11377 half_ends: &[HalfEnd],
11378 join_neighbor: &[u32],
11379 start: usize,
11380 path_id: u64,
11381 visited: &mut [u8],
11382 records: &mut Vec<StitchedCoordRecord>,
11383) {
11384 let record_start = records.len();
11385 let mut current = start;
11386 let mut rank = 0u64;
11387 let mut is_cycle = false;
11388
11389 loop {
11390 let unitig_pair = current >> 1;
11391 if visited[unitig_pair] != 0 {
11392 break;
11393 }
11394 visited[unitig_pair] = 1;
11395
11396 let unitig_index = half_ends[current].unitig_index;
11397 let reverse = reverse_for_stitch_node(current);
11398 records.push(StitchedCoordRecord {
11399 path_id,
11400 rank,
11401 unitig_index,
11402 reverse,
11403 is_cycle: false,
11404 });
11405 rank += 1;
11406
11407 let other = current ^ 1;
11408 let next = join_neighbor[other];
11409 if next == STITCH_NO_NODE {
11410 break;
11411 }
11412 if stitch_node_index(next) == start {
11413 is_cycle = true;
11414 break;
11415 }
11416 current = stitch_node_index(next);
11417 }
11418
11419 if rank == 0 {
11420 records.truncate(record_start);
11421 return;
11422 }
11423
11424 if is_cycle {
11425 for record in &mut records[record_start..] {
11426 record.is_cycle = true;
11427 }
11428 }
11429}
11430
11431fn push_stitched_records_to_range_buckets(
11432 unitig_bucket: &[u32],
11433 malformed_path: &Path,
11434 records: &[StitchedCoordRecord],
11435 records_by_bucket: &mut [Vec<StitchedCoordRecord>],
11436) -> Result<(), SerialCollationError> {
11437 for &record in records {
11438 let bucket_id = unitig_bucket
11439 .get(record.unitig_index as usize)
11440 .copied()
11441 .ok_or_else(|| {
11442 SerialCollationError::MalformedCoordBucket(malformed_path.to_path_buf())
11443 })?;
11444 let bucket_id = usize::try_from(bucket_id).map_err(|_| {
11445 SerialCollationError::MalformedCoordBucket(malformed_path.to_path_buf())
11446 })?;
11447 if bucket_id >= records_by_bucket.len() {
11448 return Err(SerialCollationError::MalformedCoordBucket(
11449 malformed_path.to_path_buf(),
11450 ));
11451 }
11452 records_by_bucket[bucket_id].push(record);
11453 }
11454 Ok(())
11455}
11456
11457#[allow(clippy::too_many_arguments)]
11458fn collect_external_stitched_path_info_range_into_memory(
11459 ranges: &[ExternalLocalUnitigRange],
11460 malformed_path: &Path,
11461 ranges_per_bucket: usize,
11462 half_ends: &[HalfEnd],
11463 join_neighbor: &[u32],
11464 starts: &[usize],
11465 base_path_id: u64,
11466 records_by_bucket: &mut [Vec<StitchedCoordRecord>],
11467) -> Result<(), SerialCollationError> {
11468 let mut records = Vec::new();
11469 for (offset, &start) in starts.iter().enumerate() {
11470 records.clear();
11471 let path_id = base_path_id + offset as u64;
11472 walk_simple_stitched_component_coords(
11473 half_ends,
11474 join_neighbor,
11475 start,
11476 path_id,
11477 &mut records,
11478 );
11479 for &record in &records {
11480 let range_id = external_range_id_for_unitig(ranges, record.unitig_index as usize)
11481 .ok_or_else(|| {
11482 SerialCollationError::MalformedCoordBucket(malformed_path.to_path_buf())
11483 })?;
11484 let bucket_id = range_id / ranges_per_bucket;
11485 records_by_bucket[bucket_id].push(record);
11486 }
11487 }
11488 Ok(())
11489}
11490
11491fn materialize_external_stitched_coord_buckets_from_memory<const K: usize>(
11492 inputs: &ExternalDiscontinuityInputs<K>,
11493 coord_dir: &Path,
11494 threads: usize,
11495 ranges_per_bucket: usize,
11496 records_by_range_bucket: Vec<Vec<StitchedCoordRecord>>,
11497) -> Result<Vec<MaterializedStitchedCoordBucketEntry>, SerialCollationError> {
11498 let bucket_count = materialized_stitched_coord_bucket_count(threads);
11499 let bucket_mask = bucket_count - 1;
11500 let non_empty_buckets = records_by_range_bucket
11501 .iter()
11502 .enumerate()
11503 .filter_map(|(bucket_id, records)| (!records.is_empty()).then_some(bucket_id))
11504 .collect::<Vec<_>>();
11505
11506 let workers = threads.max(1).min(non_empty_buckets.len().max(1));
11507 let mut manifest = if workers == 1 || non_empty_buckets.len() < 2 {
11508 let mut writers = MaterializedStitchedCoordShardWriters::new(coord_dir, 0, bucket_count);
11509 for bucket_id in non_empty_buckets {
11510 materialize_external_stitched_coord_range_group_from_records::<K>(
11511 inputs,
11512 bucket_mask,
11513 ranges_per_bucket,
11514 bucket_id,
11515 &records_by_range_bucket[bucket_id],
11516 &mut writers,
11517 )?;
11518 }
11519 writers.finish()?
11520 } else {
11521 let next_group = AtomicUsize::new(0);
11522 std::thread::scope(|scope| {
11523 let mut handles = Vec::new();
11524 for worker_id in 0..workers {
11525 let next_group = &next_group;
11526 let non_empty_buckets = &non_empty_buckets;
11527 let records_by_range_bucket = &records_by_range_bucket;
11528 handles.push(scope.spawn(move || {
11529 let mut writers = MaterializedStitchedCoordShardWriters::new(
11530 coord_dir,
11531 worker_id,
11532 bucket_count,
11533 );
11534 loop {
11535 let group_idx = next_group.fetch_add(1, Ordering::Relaxed);
11536 let Some(&bucket_id) = non_empty_buckets.get(group_idx) else {
11537 break;
11538 };
11539 materialize_external_stitched_coord_range_group_from_records::<K>(
11540 inputs,
11541 bucket_mask,
11542 ranges_per_bucket,
11543 bucket_id,
11544 &records_by_range_bucket[bucket_id],
11545 &mut writers,
11546 )?;
11547 }
11548 writers.finish()
11549 }));
11550 }
11551
11552 let mut manifest = Vec::new();
11553 for handle in handles {
11554 manifest.extend(
11555 handle
11556 .join()
11557 .map_err(|_| SerialCollationError::WorkerPanic)??,
11558 );
11559 }
11560 Ok::<_, SerialCollationError>(manifest)
11561 })?
11562 };
11563
11564 manifest.sort_by(|left, right| {
11565 left.bucket_id
11566 .cmp(&right.bucket_id)
11567 .then_with(|| left.coord_path.cmp(&right.coord_path))
11568 });
11569 Ok(manifest)
11570}
11571
11572#[allow(clippy::too_many_arguments)]
11573fn write_external_stitched_path_info_buckets<const K: usize>(
11574 inputs: &ExternalDiscontinuityInputs<K>,
11575 coord_dir: &Path,
11576 threads: usize,
11577 half_ends: &[HalfEnd],
11578 join_neighbor: &[u32],
11579 starts: &[usize],
11580) -> Result<Vec<StitchedCoordBucketEntry>, SerialCollationError> {
11581 let path_info_dir = coord_dir.join("path-info");
11582 fs::create_dir_all(&path_info_dir).map_err(|source| SerialCollationError::Io {
11583 path: path_info_dir.clone(),
11584 source,
11585 })?;
11586 let ranges_per_bucket = ranges_per_path_info_bucket(inputs.ranges.len(), threads.max(1));
11587 let path_info_bucket_count = inputs.ranges.len().div_ceil(ranges_per_bucket).max(1);
11588 let workers = threads.max(1).min(starts.len().max(1));
11589
11590 let mut manifest = if workers == 1 || starts.len() < 1024 {
11591 write_external_stitched_path_info_range(
11592 &path_info_dir,
11593 0,
11594 &inputs.ranges,
11595 ranges_per_bucket,
11596 path_info_bucket_count,
11597 half_ends,
11598 join_neighbor,
11599 starts,
11600 0,
11601 )?
11602 } else {
11603 let chunk_size = starts.len().div_ceil(workers);
11604 std::thread::scope(|scope| {
11605 let mut handles = Vec::new();
11606 for (chunk_index, chunk) in starts.chunks(chunk_size).enumerate() {
11607 let base_path_id = (chunk_index * chunk_size) as u64;
11608 let worker_dir = path_info_dir.clone();
11609 handles.push(scope.spawn(move || {
11610 write_external_stitched_path_info_range(
11611 &worker_dir,
11612 chunk_index,
11613 &inputs.ranges,
11614 ranges_per_bucket,
11615 path_info_bucket_count,
11616 half_ends,
11617 join_neighbor,
11618 chunk,
11619 base_path_id,
11620 )
11621 }));
11622 }
11623
11624 let mut manifest = Vec::new();
11625 for handle in handles {
11626 manifest.extend(
11627 handle
11628 .join()
11629 .map_err(|_| SerialCollationError::WorkerPanic)??,
11630 );
11631 }
11632 Ok::<_, SerialCollationError>(manifest)
11633 })?
11634 };
11635
11636 manifest.sort_by(|left, right| {
11637 left.bucket_id
11638 .cmp(&right.bucket_id)
11639 .then_with(|| left.path.cmp(&right.path))
11640 });
11641 Ok(manifest)
11642}
11643
11644#[allow(clippy::too_many_arguments)]
11645fn write_external_stitched_path_info_range(
11646 path_info_dir: &Path,
11647 worker_id: usize,
11648 ranges: &[ExternalLocalUnitigRange],
11649 ranges_per_bucket: usize,
11650 path_info_bucket_count: usize,
11651 half_ends: &[HalfEnd],
11652 join_neighbor: &[u32],
11653 starts: &[usize],
11654 base_path_id: u64,
11655) -> Result<Vec<StitchedCoordBucketEntry>, SerialCollationError> {
11656 let mut writers =
11657 StitchedCoordShardWriters::new(path_info_dir, worker_id, path_info_bucket_count);
11658 let mut records = Vec::new();
11659 for (offset, &start) in starts.iter().enumerate() {
11660 records.clear();
11661 let path_id = base_path_id + offset as u64;
11662 walk_simple_stitched_component_coords(
11663 half_ends,
11664 join_neighbor,
11665 start,
11666 path_id,
11667 &mut records,
11668 );
11669 for record in &records {
11670 let range_id = external_range_id_for_unitig(ranges, record.unitig_index as usize)
11671 .ok_or_else(|| {
11672 SerialCollationError::MalformedCoordBucket(path_info_dir.to_path_buf())
11673 })?;
11674 let bucket_id = range_id / ranges_per_bucket;
11675 writers.write_record(bucket_id, *record)?;
11676 }
11677 }
11678 writers.finish()
11679}
11680
11681fn ranges_per_path_info_bucket(range_count: usize, workers: usize) -> usize {
11682 let target_buckets = range_count.div_ceil(STITCH_PATH_INFO_BUCKET_TARGET).max(1);
11683 let _ = workers;
11688 let bucket_count = target_buckets.clamp(1, MAX_OPEN_STITCH_PATH_INFO_WRITERS);
11689 range_count.div_ceil(bucket_count).max(1)
11690}
11691
11692fn external_range_id_for_unitig(
11693 ranges: &[ExternalLocalUnitigRange],
11694 unitig_index: usize,
11695) -> Option<usize> {
11696 let range_id = ranges.partition_point(|range| range.start_unitig <= unitig_index);
11697 let range_id = range_id.checked_sub(1)?;
11698 let range = ranges.get(range_id)?;
11699 (unitig_index < range.start_unitig + range.unitigs).then_some(range_id)
11700}
11701
11702struct ExternalRangeIndex {
11703 page_starts: Vec<usize>,
11704}
11705
11706impl ExternalRangeIndex {
11707 const PAGE_SHIFT: usize = 16;
11708
11709 fn new(ranges: &[ExternalLocalUnitigRange]) -> Self {
11710 let unitigs = ranges
11711 .last()
11712 .map_or(0, |range| range.start_unitig + range.unitigs);
11713 let page_count = unitigs.div_ceil(1 << Self::PAGE_SHIFT);
11714 let mut page_starts = Vec::with_capacity(page_count);
11715 let mut range_id = 0;
11716 for page in 0..page_count {
11717 let unitig = page << Self::PAGE_SHIFT;
11718 while range_id + 1 < ranges.len() && ranges[range_id + 1].start_unitig <= unitig {
11719 range_id += 1;
11720 }
11721 page_starts.push(range_id);
11722 }
11723 Self { page_starts }
11724 }
11725
11726 #[inline(always)]
11727 fn find(&self, ranges: &[ExternalLocalUnitigRange], unitig_index: usize) -> Option<usize> {
11728 let mut range_id = *self.page_starts.get(unitig_index >> Self::PAGE_SHIFT)?;
11729 while range_id + 1 < ranges.len() && ranges[range_id + 1].start_unitig <= unitig_index {
11730 range_id += 1;
11731 }
11732 let range = ranges.get(range_id)?;
11733 (unitig_index < range.start_unitig + range.unitigs).then_some(range_id)
11734 }
11735}
11736
11737fn materialize_external_stitched_coord_buckets_from_path_info<const K: usize>(
11738 inputs: &ExternalDiscontinuityInputs<K>,
11739 coord_dir: &Path,
11740 threads: usize,
11741 path_info_manifest: &[StitchedCoordBucketEntry],
11742) -> Result<Vec<MaterializedStitchedCoordBucketEntry>, SerialCollationError> {
11743 let bucket_count = materialized_stitched_coord_bucket_count(threads);
11744 let bucket_mask = bucket_count - 1;
11745 let ranges_per_bucket = ranges_per_path_info_bucket(inputs.ranges.len(), threads.max(1));
11746
11747 let mut group_bounds = Vec::new();
11748 let mut group_start = 0;
11749 while group_start < path_info_manifest.len() {
11750 let bucket_id = path_info_manifest[group_start].bucket_id;
11751 let mut group_end = group_start + 1;
11752 while group_end < path_info_manifest.len()
11753 && path_info_manifest[group_end].bucket_id == bucket_id
11754 {
11755 group_end += 1;
11756 }
11757 group_bounds.push((bucket_id, group_start, group_end));
11758 group_start = group_end;
11759 }
11760
11761 let workers = threads.max(1).min(group_bounds.len().max(1));
11762 let mut manifest = if workers == 1 || group_bounds.len() < 2 {
11763 let mut writers = MaterializedStitchedCoordShardWriters::new(coord_dir, 0, bucket_count);
11764 for &(bucket_id, start, end) in &group_bounds {
11765 materialize_external_stitched_coord_range_group::<K>(
11766 inputs,
11767 bucket_mask,
11768 ranges_per_bucket,
11769 bucket_id,
11770 &path_info_manifest[start..end],
11771 &mut writers,
11772 )?;
11773 }
11774 writers.finish()?
11775 } else {
11776 let next_group = AtomicUsize::new(0);
11777 std::thread::scope(|scope| {
11778 let mut handles = Vec::new();
11779 for worker_id in 0..workers {
11780 let next_group = &next_group;
11781 let group_bounds = &group_bounds;
11782 handles.push(scope.spawn(move || {
11783 let mut writers = MaterializedStitchedCoordShardWriters::new(
11784 coord_dir,
11785 worker_id,
11786 bucket_count,
11787 );
11788 loop {
11789 let group_idx = next_group.fetch_add(1, Ordering::Relaxed);
11790 let Some(&(bucket_id, start, end)) = group_bounds.get(group_idx) else {
11791 break;
11792 };
11793 materialize_external_stitched_coord_range_group::<K>(
11794 inputs,
11795 bucket_mask,
11796 ranges_per_bucket,
11797 bucket_id,
11798 &path_info_manifest[start..end],
11799 &mut writers,
11800 )?;
11801 }
11802 writers.finish()
11803 }));
11804 }
11805
11806 let mut manifest = Vec::new();
11807 for handle in handles {
11808 manifest.extend(
11809 handle
11810 .join()
11811 .map_err(|_| SerialCollationError::WorkerPanic)??,
11812 );
11813 }
11814 Ok::<_, SerialCollationError>(manifest)
11815 })?
11816 };
11817
11818 manifest.sort_by(|left, right| {
11819 left.bucket_id
11820 .cmp(&right.bucket_id)
11821 .then_with(|| left.coord_path.cmp(&right.coord_path))
11822 });
11823 Ok(manifest)
11824}
11825
11826fn materialize_external_stitched_coord_range_group<const K: usize>(
11827 inputs: &ExternalDiscontinuityInputs<K>,
11828 bucket_mask: usize,
11829 ranges_per_bucket: usize,
11830 path_info_bucket_id: usize,
11831 path_info_entries: &[StitchedCoordBucketEntry],
11832 writers: &mut MaterializedStitchedCoordShardWriters<'_>,
11833) -> Result<(), SerialCollationError> {
11834 if path_info_entries.is_empty() {
11835 return Ok(());
11836 }
11837 let range_start = path_info_bucket_id
11838 .checked_mul(ranges_per_bucket)
11839 .ok_or_else(|| SerialCollationError::MalformedCoordBucket(inputs.unitig_path.clone()))?;
11840 let range_end = (range_start + ranges_per_bucket).min(inputs.ranges.len());
11841 let Some(first_range) = inputs.ranges.get(range_start) else {
11842 return Err(SerialCollationError::MalformedCoordBucket(
11843 inputs.unitig_path.clone(),
11844 ));
11845 };
11846 let last_range = inputs
11847 .ranges
11848 .get(range_end.saturating_sub(1))
11849 .ok_or_else(|| SerialCollationError::MalformedCoordBucket(inputs.unitig_path.clone()))?;
11850 let start_unitig = first_range.start_unitig;
11851 let end_unitig = last_range.start_unitig + last_range.unitigs;
11852 let unitig_span = end_unitig
11853 .checked_sub(start_unitig)
11854 .ok_or_else(|| SerialCollationError::MalformedCoordBucket(inputs.unitig_path.clone()))?;
11855
11856 let mut path_info: Vec<Option<StitchedCoordRecord>> = vec![None; unitig_span];
11857 for entry in path_info_entries {
11858 for record in read_stitched_coord_bucket_file(entry)? {
11859 let unitig_index = record.unitig_index as usize;
11860 if unitig_index < start_unitig || unitig_index >= end_unitig {
11861 return Err(SerialCollationError::MalformedCoordBucket(
11862 entry.path.clone(),
11863 ));
11864 }
11865 let slot = &mut path_info[unitig_index - start_unitig];
11866 if slot.replace(record).is_some() {
11867 return Err(SerialCollationError::MalformedCoordBucket(
11868 entry.path.clone(),
11869 ));
11870 }
11871 }
11872 }
11873
11874 let unitig_file =
11875 File::open(&inputs.unitig_path).map_err(|source| SerialCollationError::Io {
11876 path: inputs.unitig_path.clone(),
11877 source,
11878 })?;
11879 let mut unitig_input = BufReader::with_capacity(1024 * 1024, unitig_file);
11880 unitig_input
11881 .seek(SeekFrom::Start(
11882 (start_unitig * external_unitig_record_len::<K>(inputs.compact_unitigs)) as u64,
11883 ))
11884 .map_err(|source| SerialCollationError::Io {
11885 path: inputs.unitig_path.clone(),
11886 source,
11887 })?;
11888 let label_file = File::open(&inputs.label_path).map_err(|source| SerialCollationError::Io {
11889 path: inputs.label_path.clone(),
11890 source,
11891 })?;
11892 let mut label_input = BufReader::with_capacity(1024 * 1024, label_file);
11893 label_input
11894 .seek(SeekFrom::Start(first_range.label_start))
11895 .map_err(|source| SerialCollationError::Io {
11896 path: inputs.label_path.clone(),
11897 source,
11898 })?;
11899 let mut label = Vec::new();
11900 let mut discard = vec![0u8; 64 * 1024];
11901
11902 for record in &path_info {
11903 let unitig: DiscontinuityUnitig<K> = read_discontinuity_unitig_from_reader(
11904 &mut unitig_input,
11905 &inputs.unitig_path,
11906 inputs.compact_unitigs,
11907 )?;
11908 if let Some(record) = record {
11909 label.resize(unitig.label_len as usize, 0);
11910 label_input
11911 .read_exact(&mut label)
11912 .map_err(|source| SerialCollationError::Io {
11913 path: inputs.label_path.clone(),
11914 source,
11915 })?;
11916 let bucket_id = stitched_coord_bucket(record.path_id, bucket_mask);
11917 writers.write_materialized_record(bucket_id, record, &label)?;
11918 } else {
11919 read_and_discard_exact(
11920 &mut label_input,
11921 &inputs.label_path,
11922 &mut discard,
11923 unitig.label_len as u64,
11924 )?;
11925 }
11926 }
11927 Ok(())
11928}
11929
11930fn materialize_external_stitched_coord_range_group_from_records<const K: usize>(
11931 inputs: &ExternalDiscontinuityInputs<K>,
11932 bucket_mask: usize,
11933 ranges_per_bucket: usize,
11934 path_info_bucket_id: usize,
11935 path_info_records: &[StitchedCoordRecord],
11936 writers: &mut MaterializedStitchedCoordShardWriters<'_>,
11937) -> Result<(), SerialCollationError> {
11938 if path_info_records.is_empty() {
11939 return Ok(());
11940 }
11941 let range_start = path_info_bucket_id
11942 .checked_mul(ranges_per_bucket)
11943 .ok_or_else(|| SerialCollationError::MalformedCoordBucket(inputs.unitig_path.clone()))?;
11944 let range_end = (range_start + ranges_per_bucket).min(inputs.ranges.len());
11945 let Some(first_range) = inputs.ranges.get(range_start) else {
11946 return Err(SerialCollationError::MalformedCoordBucket(
11947 inputs.unitig_path.clone(),
11948 ));
11949 };
11950 let last_range = inputs
11951 .ranges
11952 .get(range_end.saturating_sub(1))
11953 .ok_or_else(|| SerialCollationError::MalformedCoordBucket(inputs.unitig_path.clone()))?;
11954 let start_unitig = first_range.start_unitig;
11955 let end_unitig = last_range.start_unitig + last_range.unitigs;
11956 let unitig_span = end_unitig
11957 .checked_sub(start_unitig)
11958 .ok_or_else(|| SerialCollationError::MalformedCoordBucket(inputs.unitig_path.clone()))?;
11959
11960 let mut path_info = path_info_records.to_vec();
11961 let empty_record = StitchedCoordRecord {
11962 path_id: 0,
11963 rank: 0,
11964 unitig_index: u32::MAX,
11965 reverse: false,
11966 is_cycle: false,
11967 };
11968 let mut path_info_by_unitig = vec![empty_record; unitig_span];
11969 let expected_record_count = path_info.len();
11970 for record in path_info.drain(..) {
11971 let unitig_index = record.unitig_index as usize;
11972 if unitig_index < start_unitig || unitig_index >= end_unitig {
11973 return Err(SerialCollationError::MalformedCoordBucket(
11974 inputs.unitig_path.clone(),
11975 ));
11976 }
11977 let slot = &mut path_info_by_unitig[unitig_index - start_unitig];
11978 if slot.unitig_index != u32::MAX {
11979 return Err(SerialCollationError::MalformedCoordBucket(
11980 inputs.unitig_path.clone(),
11981 ));
11982 }
11983 *slot = record;
11984 }
11985
11986 let unitig_file =
11987 File::open(&inputs.unitig_path).map_err(|source| SerialCollationError::Io {
11988 path: inputs.unitig_path.clone(),
11989 source,
11990 })?;
11991 let mut unitig_input = BufReader::with_capacity(1024 * 1024, unitig_file);
11992 unitig_input
11993 .seek(SeekFrom::Start(
11994 (start_unitig * external_unitig_record_len::<K>(inputs.compact_unitigs)) as u64,
11995 ))
11996 .map_err(|source| SerialCollationError::Io {
11997 path: inputs.unitig_path.clone(),
11998 source,
11999 })?;
12000 let label_file = File::open(&inputs.label_path).map_err(|source| SerialCollationError::Io {
12001 path: inputs.label_path.clone(),
12002 source,
12003 })?;
12004 let mut label_input = BufReader::with_capacity(1024 * 1024, label_file);
12005 label_input
12006 .seek(SeekFrom::Start(first_range.label_start))
12007 .map_err(|source| SerialCollationError::Io {
12008 path: inputs.label_path.clone(),
12009 source,
12010 })?;
12011 let mut label = Vec::new();
12012 let mut discard = vec![0u8; 64 * 1024];
12013 let mut record_count = 0usize;
12014 for unitig_index in start_unitig..end_unitig {
12015 let unitig: DiscontinuityUnitig<K> = read_discontinuity_unitig_from_reader(
12016 &mut unitig_input,
12017 &inputs.unitig_path,
12018 inputs.compact_unitigs,
12019 )?;
12020 let dense_record = &path_info_by_unitig[unitig_index - start_unitig];
12021 let record = (dense_record.unitig_index != u32::MAX).then_some(dense_record);
12022 if let Some(record) = record {
12023 label.resize(unitig.label_len as usize, 0);
12024 label_input
12025 .read_exact(&mut label)
12026 .map_err(|source| SerialCollationError::Io {
12027 path: inputs.label_path.clone(),
12028 source,
12029 })?;
12030 let bucket_id = stitched_coord_bucket(record.path_id, bucket_mask);
12031 writers.write_materialized_record(bucket_id, record, &label)?;
12032 record_count += 1;
12033 } else {
12034 read_and_discard_exact(
12035 &mut label_input,
12036 &inputs.label_path,
12037 &mut discard,
12038 unitig.label_len as u64,
12039 )?;
12040 }
12041 }
12042 if record_count != expected_record_count {
12043 return Err(SerialCollationError::MalformedCoordBucket(
12044 inputs.unitig_path.clone(),
12045 ));
12046 }
12047 debug_assert_eq!(unitig_span, end_unitig - start_unitig);
12048 Ok(())
12049}
12050
12051fn map_external_cpp_path_info_range_bucket<const K: usize, F, W>(
12052 inputs: &ExternalDiscontinuityInputs<K>,
12053 bucket_mask: usize,
12054 ranges_per_bucket: usize,
12055 path_info_bucket_id: usize,
12056 path_info_records: &[StitchedCoordRecord],
12057 writers: &mut W,
12058 emit_direct_local: &mut F,
12059) -> Result<(), SerialCollationError>
12060where
12061 F: FnMut(FinalUnitigRecord) -> Result<(), SerialCollationError>,
12062 W: MaterializedRecordSink,
12063{
12064 map_external_cpp_path_info_range_bucket_owned(
12065 inputs,
12066 bucket_mask,
12067 ranges_per_bucket,
12068 path_info_bucket_id,
12069 path_info_records.to_vec(),
12070 writers,
12071 emit_direct_local,
12072 )
12073}
12074
12075fn map_external_cpp_path_info_range_bucket_owned<const K: usize, F, W>(
12076 inputs: &ExternalDiscontinuityInputs<K>,
12077 bucket_mask: usize,
12078 ranges_per_bucket: usize,
12079 path_info_bucket_id: usize,
12080 mut path_info: Vec<StitchedCoordRecord>,
12081 writers: &mut W,
12082 emit_direct_local: &mut F,
12083) -> Result<(), SerialCollationError>
12084where
12085 F: FnMut(FinalUnitigRecord) -> Result<(), SerialCollationError>,
12086 W: MaterializedRecordSink,
12087{
12088 let range_start = path_info_bucket_id
12089 .checked_mul(ranges_per_bucket)
12090 .ok_or_else(|| SerialCollationError::MalformedCoordBucket(inputs.unitig_path.clone()))?;
12091 let range_end = (range_start + ranges_per_bucket).min(inputs.ranges.len());
12092 let Some(first_range) = inputs.ranges.get(range_start) else {
12093 if path_info.is_empty() {
12094 return Ok(());
12095 }
12096 return Err(SerialCollationError::MalformedCoordBucket(
12097 inputs.unitig_path.clone(),
12098 ));
12099 };
12100 let last_range = inputs
12101 .ranges
12102 .get(range_end.saturating_sub(1))
12103 .ok_or_else(|| SerialCollationError::MalformedCoordBucket(inputs.unitig_path.clone()))?;
12104 let start_unitig = first_range.start_unitig;
12105 let end_unitig = last_range.start_unitig + last_range.unitigs;
12106 let unitig_span = end_unitig
12107 .checked_sub(start_unitig)
12108 .ok_or_else(|| SerialCollationError::MalformedCoordBucket(inputs.unitig_path.clone()))?;
12109
12110 path_info.sort_unstable_by_key(|record| record.unitig_index);
12111 for pair in path_info.windows(2) {
12112 if pair[0].unitig_index == pair[1].unitig_index {
12113 return Err(SerialCollationError::MalformedCoordBucket(
12114 inputs.unitig_path.clone(),
12115 ));
12116 }
12117 }
12118 for record in &path_info {
12119 let unitig_index = record.unitig_index as usize;
12120 if unitig_index < start_unitig || unitig_index >= end_unitig {
12121 return Err(SerialCollationError::MalformedCoordBucket(
12122 inputs.unitig_path.clone(),
12123 ));
12124 }
12125 }
12126
12127 let unitig_file =
12128 File::open(&inputs.unitig_path).map_err(|source| SerialCollationError::Io {
12129 path: inputs.unitig_path.clone(),
12130 source,
12131 })?;
12132 let mut unitig_input = BufReader::with_capacity(1024 * 1024, unitig_file);
12133 unitig_input
12134 .seek(SeekFrom::Start(
12135 (start_unitig * external_unitig_record_len::<K>(inputs.compact_unitigs)) as u64,
12136 ))
12137 .map_err(|source| SerialCollationError::Io {
12138 path: inputs.unitig_path.clone(),
12139 source,
12140 })?;
12141 let label_file = File::open(&inputs.label_path).map_err(|source| SerialCollationError::Io {
12142 path: inputs.label_path.clone(),
12143 source,
12144 })?;
12145 let mut label_input = BufReader::with_capacity(1024 * 1024, label_file);
12146 label_input
12147 .seek(SeekFrom::Start(first_range.label_start))
12148 .map_err(|source| SerialCollationError::Io {
12149 path: inputs.label_path.clone(),
12150 source,
12151 })?;
12152 let mut label = Vec::new();
12153 let mut discard = vec![0u8; 64 * 1024];
12154 let mut record_idx = 0usize;
12155 let mut color_reader = inputs
12156 .color_runs
12157 .as_ref()
12158 .map(|sidecar| sidecar.reader_at(first_range.color_start))
12159 .transpose()?;
12160 let mut colors = Vec::new();
12161 let mut current_range = range_start;
12162
12163 for unitig_index in start_unitig..end_unitig {
12164 if inputs
12165 .ranges
12166 .get(current_range)
12167 .is_some_and(|range| range.start_unitig == unitig_index)
12168 {
12169 let range = &inputs.ranges[current_range];
12170 label_input
12171 .seek(SeekFrom::Start(range.label_start))
12172 .map_err(|source| SerialCollationError::Io {
12173 path: inputs.label_path.clone(),
12174 source,
12175 })?;
12176 color_reader = inputs
12177 .color_runs
12178 .as_ref()
12179 .map(|sidecar| sidecar.reader_at(range.color_start))
12180 .transpose()?;
12181 current_range += 1;
12182 }
12183 let unitig: DiscontinuityUnitig<K> = read_discontinuity_unitig_from_reader(
12184 &mut unitig_input,
12185 &inputs.unitig_path,
12186 inputs.compact_unitigs,
12187 )?;
12188 let has_colors = if let Some(reader) = color_reader.as_mut() {
12189 reader.read_next_into(&mut colors)?;
12190 true
12191 } else {
12192 false
12193 };
12194 let record = path_info
12195 .get(record_idx)
12196 .filter(|record| record.unitig_index as usize == unitig_index);
12197 if let Some(record) = record {
12198 label.resize(unitig.label_len as usize, 0);
12199 label_input
12200 .read_exact(&mut label)
12201 .map_err(|source| SerialCollationError::Io {
12202 path: inputs.label_path.clone(),
12203 source,
12204 })?;
12205 let bucket_id = stitched_coord_bucket(record.path_id, bucket_mask);
12206 write_external_materialized_record(
12207 writers,
12208 bucket_id,
12209 record,
12210 &label,
12211 has_colors.then_some(colors.as_slice()),
12212 )?;
12213 record_idx += 1;
12214 } else if unitig.left_exit().is_none() && unitig.right_exit().is_none() {
12215 label.resize(unitig.label_len as usize, 0);
12216 label_input
12217 .read_exact(&mut label)
12218 .map_err(|source| SerialCollationError::Io {
12219 path: inputs.label_path.clone(),
12220 source,
12221 })?;
12222 let reverse = reverse_complement_is_less(&label);
12223 let label = if reverse {
12224 reverse_complement_label(&label)
12225 } else {
12226 label.clone()
12227 };
12228 let mut direct_colors = if has_colors {
12229 std::mem::take(&mut colors)
12230 } else {
12231 Vec::new()
12232 };
12233 if reverse && !direct_colors.is_empty() {
12234 direct_colors =
12235 reverse_color_runs(&direct_colors, (unitig.label_len as usize - K + 1) as u32);
12236 }
12237 emit_direct_local(FinalUnitigRecord {
12238 label,
12239 colors: direct_colors,
12240 })?;
12241 } else {
12242 read_and_discard_exact(
12243 &mut label_input,
12244 &inputs.label_path,
12245 &mut discard,
12246 unitig.label_len as u64,
12247 )?;
12248 }
12249 }
12250 if record_idx != path_info.len() {
12251 return Err(SerialCollationError::MalformedCoordBucket(
12252 inputs.unitig_path.clone(),
12253 ));
12254 }
12255 debug_assert_eq!(unitig_span, end_unitig - start_unitig);
12256 Ok(())
12257}
12258
12259fn write_external_materialized_record<W: MaterializedRecordSink>(
12260 writers: &mut W,
12261 bucket_id: usize,
12262 record: &StitchedCoordRecord,
12263 label: &[u8],
12264 colors: Option<&[UnitigColor]>,
12265) -> Result<(), SerialCollationError> {
12266 if let Some(colors) = colors {
12267 writers.write_materialized_colored_record(bucket_id, record, label, colors)
12268 } else {
12269 writers.write_materialized_record(bucket_id, record, label)
12270 }
12271}
12272
12273fn read_and_discard_exact(
12274 input: &mut BufReader<File>,
12275 path: &Path,
12276 scratch: &mut [u8],
12277 mut len: u64,
12278) -> Result<(), SerialCollationError> {
12279 while len != 0 {
12280 let chunk_len = scratch.len().min(len as usize);
12281 input
12282 .read_exact(&mut scratch[..chunk_len])
12283 .map_err(|source| SerialCollationError::Io {
12284 path: path.to_path_buf(),
12285 source,
12286 })?;
12287 len -= chunk_len as u64;
12288 }
12289 Ok(())
12290}
12291
12292#[allow(clippy::too_many_arguments)]
12293fn write_materialized_neighbor_stitched_coord_shards<const K: usize>(
12294 inputs: &DiscontinuityInputs<K>,
12295 coord_dir: &Path,
12296 worker_id: usize,
12297 bucket_count: usize,
12298 bucket_mask: usize,
12299 half_ends: &[HalfEnd],
12300 join_neighbor: &[u32],
12301 starts: &[usize],
12302 base_path_id: u64,
12303) -> Result<Vec<MaterializedStitchedCoordBucketEntry>, SerialCollationError> {
12304 let mut writers =
12305 MaterializedStitchedCoordShardWriters::new(coord_dir, worker_id, bucket_count);
12306 let mut records = Vec::new();
12307 for (offset, &start) in starts.iter().enumerate() {
12308 records.clear();
12309 let path_id = base_path_id + offset as u64;
12310 walk_simple_stitched_component_coords(
12311 half_ends,
12312 join_neighbor,
12313 start,
12314 path_id,
12315 &mut records,
12316 );
12317 writers.write_path_records(
12318 inputs,
12319 stitched_coord_bucket(path_id, bucket_mask),
12320 &records,
12321 )?;
12322 }
12323 writers.finish()
12324}
12325
12326struct MaterializedStitchedCoordShardWriters<'a> {
12327 coord_dir: &'a Path,
12328 worker_id: usize,
12329 writers: Vec<Option<MaterializedStitchedCoordShardWriter>>,
12330 open_writers: usize,
12331}
12332
12333trait MaterializedRecordSink {
12334 fn write_materialized_record(
12335 &mut self,
12336 bucket_id: usize,
12337 record: &StitchedCoordRecord,
12338 label: &[u8],
12339 ) -> Result<(), SerialCollationError>;
12340
12341 fn write_materialized_colored_record(
12342 &mut self,
12343 bucket_id: usize,
12344 record: &StitchedCoordRecord,
12345 label: &[u8],
12346 colors: &[UnitigColor],
12347 ) -> Result<(), SerialCollationError>;
12348}
12349
12350impl MaterializedRecordSink for MaterializedStitchedCoordShardWriters<'_> {
12351 fn write_materialized_record(
12352 &mut self,
12353 bucket_id: usize,
12354 record: &StitchedCoordRecord,
12355 label: &[u8],
12356 ) -> Result<(), SerialCollationError> {
12357 MaterializedStitchedCoordShardWriters::write_materialized_record(
12358 self, bucket_id, record, label,
12359 )
12360 }
12361
12362 fn write_materialized_colored_record(
12363 &mut self,
12364 bucket_id: usize,
12365 record: &StitchedCoordRecord,
12366 label: &[u8],
12367 colors: &[UnitigColor],
12368 ) -> Result<(), SerialCollationError> {
12369 self.ensure_writer(bucket_id)?;
12370 self.writers[bucket_id]
12371 .as_mut()
12372 .expect("bucket writer was just created")
12373 .write_colored_record(record, label, colors)
12374 }
12375}
12376
12377struct SharedMaterializedWriters<'a> {
12378 coord_dir: &'a Path,
12379 writers: Vec<Mutex<Option<MaterializedStitchedCoordShardWriter>>>,
12380 retained: Vec<Mutex<Vec<PendingMaterializedBucket>>>,
12381 open_cache: Mutex<OpenMaterializedWriterCache>,
12382}
12383
12384struct OpenMaterializedWriterCache {
12385 open: usize,
12386 limit: usize,
12387 eviction_cursor: usize,
12388}
12389
12390fn materialized_coordinate_plan(
12400 file_limit: usize,
12401 open_files: usize,
12402 threads: usize,
12403 colored: bool,
12404 unitig_bases: u64,
12405) -> (usize, usize, usize) {
12406 let threads = threads.max(1);
12407 let writer_files = if colored { 3 } else { 2 };
12408 let mut buckets = coordinate_bucket_target(threads, unitig_bases);
12416 loop {
12417 let writer_descriptors = buckets.saturating_mul(writer_files);
12418 let mapping_workers = file_limit
12419 .saturating_sub(open_files.saturating_add(writer_descriptors))
12420 .checked_div(2)
12421 .unwrap_or(0)
12422 .min(threads);
12423 if mapping_workers == threads || buckets == 1 {
12424 return (buckets, mapping_workers.max(1), buckets);
12425 }
12426 buckets /= 2;
12427 }
12428}
12429
12430fn coordinate_bucket_target(threads: usize, unitig_bases: u64) -> usize {
12446 if let Some(buckets) = std::env::var("CF3_RS_MCOORD_BUCKETS")
12447 .ok()
12448 .and_then(|value| value.parse::<usize>().ok())
12449 .filter(|value| value.is_power_of_two())
12450 {
12451 return buckets;
12452 }
12453 if threads < HIGH_THREAD_COORD_BUCKET_THRESHOLD {
12454 return DEFAULT_MAX_UNITIG_COORD_BUCKETS;
12455 }
12456 let narrow_bucket_bases = unitig_bases / HIGH_THREAD_MAX_UNITIG_COORD_BUCKETS as u64;
12459 if narrow_bucket_bases <= MAX_NARROW_COORD_BUCKET_BASES {
12460 HIGH_THREAD_MAX_UNITIG_COORD_BUCKETS
12461 } else {
12462 DEFAULT_MAX_UNITIG_COORD_BUCKETS
12463 }
12464}
12465
12466const LOCAL_UNITIG_BUCKETS_PER_WORKER: usize = 8;
12476const MAX_LOCAL_UNITIG_BUCKETS: usize = 1024;
12483
12484fn local_unitig_bucket_plan(file_limit: usize, open_files: usize, workers: usize) -> usize {
12485 if workers <= 1 {
12486 return 0;
12487 }
12488 let concurrent_worker_files = workers.saturating_mul(2);
12491 let affordable =
12492 file_limit.saturating_sub(open_files.saturating_add(concurrent_worker_files)) / 2;
12493 let per_worker =
12496 (MAX_LOCAL_UNITIG_BUCKETS / workers.max(1)).clamp(1, LOCAL_UNITIG_BUCKETS_PER_WORKER);
12497 let desired = workers.saturating_mul(per_worker);
12498 if affordable >= desired {
12499 return desired;
12500 }
12501 let per_worker = (affordable / workers.max(1)).max(1);
12502 affordable.min(workers.saturating_mul(per_worker))
12503}
12504
12505#[derive(Default)]
12506struct PendingMaterializedBucket {
12507 records: Vec<LoadedMaterializedStitchedCoordRecord>,
12508 labels: Vec<u8>,
12509 colors: Vec<UnitigColor>,
12510}
12511
12512impl<'a> SharedMaterializedWriters<'a> {
12513 fn new(coord_dir: &'a Path, bucket_count: usize, open_limit: usize) -> Self {
12514 Self {
12515 coord_dir,
12516 writers: (0..bucket_count).map(|_| Mutex::new(None)).collect(),
12517 retained: (0..bucket_count).map(|_| Mutex::new(Vec::new())).collect(),
12518 open_cache: Mutex::new(OpenMaterializedWriterCache {
12519 open: 0,
12520 limit: open_limit.clamp(1, bucket_count.max(1)),
12521 eviction_cursor: 0,
12522 }),
12523 }
12524 }
12525
12526 fn ensure_writer_open(
12527 &self,
12528 bucket_id: usize,
12529 writer: &mut Option<MaterializedStitchedCoordShardWriter>,
12530 ) -> Result<(), SerialCollationError> {
12531 if writer.as_ref().is_some_and(|writer| writer.is_open()) {
12532 return Ok(());
12533 }
12534 loop {
12535 let mut cache = self
12536 .open_cache
12537 .lock()
12538 .map_err(|_| SerialCollationError::WorkerPanic)?;
12539 if cache.open < cache.limit {
12540 if let Some(writer) = writer.as_mut() {
12541 writer.ensure_open()?;
12542 } else {
12543 *writer = Some(MaterializedStitchedCoordShardWriter::create(
12544 self.coord_dir,
12545 0,
12546 bucket_id,
12547 )?);
12548 }
12549 cache.open += 1;
12550 return Ok(());
12551 }
12552
12553 let start = cache.eviction_cursor;
12554 let mut evicted = false;
12555 for offset in 0..self.writers.len() {
12556 let candidate_id = (start + offset) % self.writers.len();
12557 if candidate_id == bucket_id {
12558 continue;
12559 }
12560 let Ok(mut candidate) = self.writers[candidate_id].try_lock() else {
12561 continue;
12562 };
12563 if let Some(candidate) = candidate.as_mut()
12564 && candidate.is_open()
12565 {
12566 candidate.flush_and_close()?;
12567 cache.open -= 1;
12568 cache.eviction_cursor = (candidate_id + 1) % self.writers.len();
12569 evicted = true;
12570 break;
12571 }
12572 }
12573 drop(cache);
12574 if !evicted {
12575 std::thread::yield_now();
12576 }
12577 }
12578 }
12579
12580 fn retain_batch(
12581 &self,
12582 bucket_id: usize,
12583 batch: &mut PendingMaterializedBucket,
12584 ) -> Result<(), SerialCollationError> {
12585 if batch.records.is_empty() {
12586 return Ok(());
12587 }
12588 self.retained[bucket_id]
12589 .lock()
12590 .map_err(|_| SerialCollationError::WorkerPanic)?
12591 .push(std::mem::take(batch));
12592 Ok(())
12593 }
12594
12595 fn write_batch(
12596 &self,
12597 bucket_id: usize,
12598 batch: &mut PendingMaterializedBucket,
12599 ) -> Result<(), SerialCollationError> {
12600 if batch.records.is_empty() {
12601 return Ok(());
12602 }
12603 let mut writer = self.writers[bucket_id]
12604 .lock()
12605 .map_err(|_| SerialCollationError::WorkerPanic)?;
12606 self.ensure_writer_open(bucket_id, &mut writer)?;
12607 let writer = writer.as_mut().expect("global writer was just created");
12608 writer.write_pending_batch(batch)?;
12609 Ok(())
12610 }
12611
12612 #[allow(dead_code)]
12614 fn write_colored_record(
12615 &self,
12616 bucket_id: usize,
12617 record: &StitchedCoordRecord,
12618 label: &[u8],
12619 colors: &[UnitigColor],
12620 ) -> Result<(), SerialCollationError> {
12621 let mut writer = self.writers[bucket_id]
12622 .lock()
12623 .map_err(|_| SerialCollationError::WorkerPanic)?;
12624 self.ensure_writer_open(bucket_id, &mut writer)?;
12625 let writer = writer.as_mut().expect("global writer was just created");
12626 writer.write_colored_record(record, label, colors)?;
12627 Ok(())
12628 }
12629
12630 fn finish(self) -> Result<MaterializedStitchedCoordBuckets, SerialCollationError> {
12631 let mut manifest = Vec::new();
12632 for writer in self.writers {
12633 if let Some(writer) = writer
12634 .into_inner()
12635 .map_err(|_| SerialCollationError::WorkerPanic)?
12636 {
12637 manifest.push(writer.finish()?);
12638 }
12639 }
12640 let retained = self
12641 .retained
12642 .into_iter()
12643 .map(|bucket| {
12644 bucket
12645 .into_inner()
12646 .map_err(|_| SerialCollationError::WorkerPanic)
12647 })
12648 .collect::<Result<Vec<_>, _>>()?;
12649 Ok(MaterializedStitchedCoordBuckets { manifest, retained })
12650 }
12651}
12652
12653struct MaterializedStitchedCoordBuckets {
12654 manifest: Vec<MaterializedStitchedCoordBucketEntry>,
12655 retained: Vec<Vec<PendingMaterializedBucket>>,
12656}
12657
12658struct SharedMaterializedBatch<'a, 'b> {
12659 shared: &'a SharedMaterializedWriters<'b>,
12660 buckets: Vec<PendingMaterializedBucket>,
12661}
12662
12663impl<'a, 'b> SharedMaterializedBatch<'a, 'b> {
12664 const FLUSH_BYTES: usize = 8 * 1024;
12667
12668 fn new(shared: &'a SharedMaterializedWriters<'b>, bucket_count: usize) -> Self {
12669 Self {
12670 shared,
12671 buckets: (0..bucket_count)
12672 .map(|_| PendingMaterializedBucket::default())
12673 .collect(),
12674 }
12675 }
12676
12677 fn finish(mut self) -> Result<(), SerialCollationError> {
12678 for bucket_id in 0..self.buckets.len() {
12679 self.shared
12680 .retain_batch(bucket_id, &mut self.buckets[bucket_id])?;
12681 }
12682 Ok(())
12683 }
12684
12685 fn finish_bucket(&mut self, bucket_id: usize) -> Result<(), SerialCollationError> {
12686 self.shared
12687 .write_batch(bucket_id, &mut self.buckets[bucket_id])?;
12688 Ok(())
12689 }
12690
12691 #[inline]
12692 fn uncolored_bucket_ready(bucket: &PendingMaterializedBucket) -> bool {
12693 bucket.records.len() >= Self::FLUSH_BYTES / STITCH_COORD_RECORD_LEN as usize
12694 && bucket.labels.len() >= Self::FLUSH_BYTES
12695 }
12696
12697 #[inline]
12698 fn colored_bucket_ready(bucket: &PendingMaterializedBucket) -> bool {
12699 Self::uncolored_bucket_ready(bucket)
12700 && bucket.colors.len() * std::mem::size_of::<UnitigColor>() >= Self::FLUSH_BYTES
12701 }
12702}
12703
12704impl MaterializedRecordSink for SharedMaterializedBatch<'_, '_> {
12705 fn write_materialized_record(
12706 &mut self,
12707 bucket_id: usize,
12708 record: &StitchedCoordRecord,
12709 label: &[u8],
12710 ) -> Result<(), SerialCollationError> {
12711 let bucket = &mut self.buckets[bucket_id];
12712 let label_offset = u32::try_from(bucket.labels.len()).map_err(|_| {
12713 SerialCollationError::MalformedCoordBucket(self.shared.coord_dir.to_path_buf())
12714 })?;
12715 bucket.labels.extend_from_slice(label);
12716 bucket
12717 .records
12718 .push(LoadedMaterializedStitchedCoordRecord::new(
12719 record.path_id,
12720 record.rank,
12721 label_offset,
12722 label.len() as u32,
12723 record.reverse,
12724 record.is_cycle,
12725 u32::MAX,
12726 0,
12727 ));
12728 if Self::uncolored_bucket_ready(bucket) {
12729 self.shared
12730 .write_batch(bucket_id, &mut self.buckets[bucket_id])?;
12731 }
12732 Ok(())
12733 }
12734
12735 fn write_materialized_colored_record(
12736 &mut self,
12737 bucket_id: usize,
12738 record: &StitchedCoordRecord,
12739 label: &[u8],
12740 colors: &[UnitigColor],
12741 ) -> Result<(), SerialCollationError> {
12742 let color_count = u32::try_from(colors.len()).map_err(|_| {
12743 SerialCollationError::MalformedCoordBucket(self.shared.coord_dir.to_path_buf())
12744 })?;
12745 let bucket = &mut self.buckets[bucket_id];
12746 let label_offset = u32::try_from(bucket.labels.len()).map_err(|_| {
12747 SerialCollationError::MalformedCoordBucket(self.shared.coord_dir.to_path_buf())
12748 })?;
12749 let color_start = bucket.colors.len() as u32;
12750 bucket.labels.extend_from_slice(label);
12751 bucket.colors.extend_from_slice(colors);
12752 bucket
12753 .records
12754 .push(LoadedMaterializedStitchedCoordRecord::new(
12755 record.path_id,
12756 record.rank,
12757 label_offset,
12758 label.len() as u32,
12759 record.reverse,
12760 record.is_cycle,
12761 color_start,
12762 color_count,
12763 ));
12764 if Self::colored_bucket_ready(bucket) {
12765 self.finish_bucket(bucket_id)?;
12766 }
12767 Ok(())
12768 }
12769}
12770
12771impl<'a> MaterializedStitchedCoordShardWriters<'a> {
12772 fn new(coord_dir: &'a Path, worker_id: usize, bucket_count: usize) -> Self {
12773 let mut writers = Vec::with_capacity(bucket_count);
12774 writers.resize_with(bucket_count, || None);
12775 Self {
12776 coord_dir,
12777 worker_id,
12778 writers,
12779 open_writers: 0,
12780 }
12781 }
12782
12783 fn write_path_records<const K: usize>(
12784 &mut self,
12785 inputs: &DiscontinuityInputs<K>,
12786 bucket_id: usize,
12787 records: &[StitchedCoordRecord],
12788 ) -> Result<(), SerialCollationError> {
12789 if records.is_empty() {
12790 return Ok(());
12791 }
12792 self.ensure_writer(bucket_id)?;
12793 let writer = self.writers[bucket_id]
12794 .as_mut()
12795 .expect("bucket writer was just created");
12796 writer.ensure_open()?;
12797 for record in records {
12798 let unitig = &inputs.unitigs[record.unitig_index as usize];
12799 writer.write_record(record, unitig.label(inputs))?;
12800 }
12801 Ok(())
12802 }
12803
12804 fn ensure_writer(&mut self, bucket_id: usize) -> Result<(), SerialCollationError> {
12805 if self.writers[bucket_id].is_none() {
12806 self.evict_writer_if_needed(bucket_id)?;
12807 self.writers[bucket_id] = Some(MaterializedStitchedCoordShardWriter::create(
12808 self.coord_dir,
12809 self.worker_id,
12810 bucket_id,
12811 )?);
12812 self.open_writers += 1;
12813 } else if self
12814 .writers
12815 .get(bucket_id)
12816 .and_then(Option::as_ref)
12817 .is_some_and(|writer| !writer.is_open())
12818 {
12819 self.evict_writer_if_needed(bucket_id)?;
12820 self.writers[bucket_id]
12821 .as_mut()
12822 .expect("checked that writer exists")
12823 .ensure_open()?;
12824 self.open_writers += 1;
12825 }
12826 Ok(())
12827 }
12828
12829 fn evict_writer_if_needed(
12830 &mut self,
12831 requested_bucket_id: usize,
12832 ) -> Result<(), SerialCollationError> {
12833 if self.open_writers < MAX_OPEN_MATERIALIZED_STITCH_WRITERS_PER_SHARD {
12834 return Ok(());
12835 }
12836 let evict_bucket_id = self
12837 .writers
12838 .iter()
12839 .enumerate()
12840 .find_map(|(bucket_id, writer)| {
12841 (bucket_id != requested_bucket_id
12842 && writer
12843 .as_ref()
12844 .is_some_and(MaterializedStitchedCoordShardWriter::is_open))
12845 .then_some(bucket_id)
12846 })
12847 .unwrap_or(requested_bucket_id);
12848 if let Some(writer) = self.writers[evict_bucket_id].as_mut()
12849 && writer.is_open()
12850 {
12851 writer.flush_and_close()?;
12852 self.open_writers -= 1;
12853 }
12854 Ok(())
12855 }
12856
12857 fn write_materialized_record(
12858 &mut self,
12859 bucket_id: usize,
12860 record: &StitchedCoordRecord,
12861 label: &[u8],
12862 ) -> Result<(), SerialCollationError> {
12863 self.ensure_writer(bucket_id)?;
12864 self.writers[bucket_id]
12865 .as_mut()
12866 .expect("bucket writer was just created")
12867 .write_record(record, label)
12868 }
12869
12870 fn finish(self) -> Result<Vec<MaterializedStitchedCoordBucketEntry>, SerialCollationError> {
12871 let mut manifest = Vec::new();
12872 for writer in self.writers.into_iter().flatten() {
12873 manifest.push(writer.finish()?);
12874 }
12875 Ok(manifest)
12876 }
12877}
12878
12879struct MaterializedStitchedCoordShardWriter {
12880 bucket_id: usize,
12881 coord_path: PathBuf,
12882 label_path: PathBuf,
12883 color_path: PathBuf,
12884 coord_out: Option<BufWriter<File>>,
12885 label_out: Option<BufWriter<File>>,
12886 color_out: Option<BufWriter<File>>,
12887 record_buffer: Vec<u8>,
12888 records: u64,
12889 label_bytes: u64,
12890 color_runs: u64,
12891}
12892
12893#[inline]
12894fn unitig_colors_as_bytes(colors: &[UnitigColor]) -> &[u8] {
12895 unsafe {
12898 std::slice::from_raw_parts(colors.as_ptr().cast::<u8>(), std::mem::size_of_val(colors))
12899 }
12900}
12901
12902impl MaterializedStitchedCoordShardWriter {
12903 fn create(
12904 coord_dir: &Path,
12905 worker_id: usize,
12906 bucket_id: usize,
12907 ) -> Result<Self, SerialCollationError> {
12908 let coord_path = coord_dir.join(format!("{bucket_id:05}.{worker_id:03}.mcoord"));
12909 let label_path = coord_dir.join(format!("{bucket_id:05}.{worker_id:03}.mlabel"));
12910 let color_path = coord_dir.join(format!("{bucket_id:05}.{worker_id:03}.mcolor"));
12911 let coord_file = File::create(&coord_path).map_err(|source| SerialCollationError::Io {
12912 path: coord_path.clone(),
12913 source,
12914 })?;
12915 let label_file = File::create(&label_path).map_err(|source| SerialCollationError::Io {
12916 path: label_path.clone(),
12917 source,
12918 })?;
12919 let mut coord_out =
12920 BufWriter::with_capacity(MATERIALIZED_STITCH_COORD_SHARD_WRITE_BUFFER, coord_file);
12921 coord_out
12922 .write_all(MATERIALIZED_STITCH_COORD_MAGIC)
12923 .and_then(|_| coord_out.write_all(&(bucket_id as u64).to_le_bytes()))
12924 .and_then(|_| coord_out.write_all(&0u64.to_le_bytes()))
12925 .and_then(|_| coord_out.write_all(&0u64.to_le_bytes()))
12926 .map_err(|source| SerialCollationError::Io {
12927 path: coord_path.clone(),
12928 source,
12929 })?;
12930 Ok(Self {
12931 bucket_id,
12932 coord_path,
12933 label_path,
12934 color_path,
12935 coord_out: Some(coord_out),
12936 label_out: Some(BufWriter::with_capacity(1024 * 1024, label_file)),
12937 color_out: None,
12938 record_buffer: Vec::with_capacity(STITCH_COORD_RECORD_WRITE_BUFFER),
12939 records: 0,
12940 label_bytes: 0,
12941 color_runs: 0,
12942 })
12943 }
12944
12945 fn is_open(&self) -> bool {
12946 self.coord_out.is_some() || self.label_out.is_some() || self.color_out.is_some()
12947 }
12948
12949 fn ensure_open(&mut self) -> Result<(), SerialCollationError> {
12950 if self.coord_out.is_none() {
12951 let coord_file = OpenOptions::new()
12952 .append(true)
12953 .open(&self.coord_path)
12954 .map_err(|source| SerialCollationError::Io {
12955 path: self.coord_path.clone(),
12956 source,
12957 })?;
12958 self.coord_out = Some(BufWriter::with_capacity(
12959 MATERIALIZED_STITCH_COORD_SHARD_WRITE_BUFFER,
12960 coord_file,
12961 ));
12962 }
12963 if self.label_out.is_none() {
12964 let label_file = OpenOptions::new()
12965 .append(true)
12966 .open(&self.label_path)
12967 .map_err(|source| SerialCollationError::Io {
12968 path: self.label_path.clone(),
12969 source,
12970 })?;
12971 self.label_out = Some(BufWriter::with_capacity(1024 * 1024, label_file));
12972 }
12973 if self.color_runs != 0 && self.color_out.is_none() {
12974 let color_file = OpenOptions::new()
12975 .append(true)
12976 .open(&self.color_path)
12977 .map_err(|source| SerialCollationError::Io {
12978 path: self.color_path.clone(),
12979 source,
12980 })?;
12981 self.color_out = Some(BufWriter::with_capacity(1024 * 1024, color_file));
12982 }
12983 Ok(())
12984 }
12985
12986 fn write_record(
12987 &mut self,
12988 record: &StitchedCoordRecord,
12989 label: &[u8],
12990 ) -> Result<(), SerialCollationError> {
12991 self.write_record_with_color_index(record, label, u32::MAX, 0)
12992 }
12993
12994 fn write_record_with_color_index(
12995 &mut self,
12996 record: &StitchedCoordRecord,
12997 label: &[u8],
12998 color_index: u32,
12999 color_count: u32,
13000 ) -> Result<(), SerialCollationError> {
13001 let label_len = u32::try_from(label.len())
13002 .map_err(|_| SerialCollationError::MalformedCoordBucket(self.coord_path.clone()))?;
13003 self.record_buffer
13004 .extend_from_slice(&encoded_materialized_stitched_coord_record(
13005 MaterializedStitchedCoordRecord {
13006 path_id: record.path_id,
13007 rank: record.rank,
13008 label_offset: self.label_bytes,
13009 label_len,
13010 reverse: record.reverse,
13011 is_cycle: record.is_cycle,
13012 color_index,
13013 color_count,
13014 },
13015 ));
13016 if self.record_buffer.len() >= STITCH_COORD_RECORD_WRITE_BUFFER {
13017 self.flush_record_buffer()?;
13018 }
13019 self.label_out
13020 .as_mut()
13021 .expect("label writer is open")
13022 .write_all(label)
13023 .map_err(|source| SerialCollationError::Io {
13024 path: self.label_path.clone(),
13025 source,
13026 })?;
13027 self.records += 1;
13028 self.label_bytes += u64::from(label_len);
13029 Ok(())
13030 }
13031
13032 fn write_colored_record(
13033 &mut self,
13034 record: &StitchedCoordRecord,
13035 label: &[u8],
13036 colors: &[UnitigColor],
13037 ) -> Result<(), SerialCollationError> {
13038 let count = u32::try_from(colors.len())
13039 .map_err(|_| SerialCollationError::MalformedCoordBucket(self.color_path.clone()))?;
13040 if self.color_out.is_none() {
13041 let file = OpenOptions::new()
13042 .create(true)
13043 .append(true)
13044 .open(&self.color_path)
13045 .map_err(|source| SerialCollationError::Io {
13046 path: self.color_path.clone(),
13047 source,
13048 })?;
13049 self.color_out = Some(BufWriter::with_capacity(1024 * 1024, file));
13050 }
13051 let color_out = self
13052 .color_out
13053 .as_mut()
13054 .expect("color writer was just created");
13055 color_out
13056 .write_all(unitig_colors_as_bytes(colors))
13057 .map_err(|source| SerialCollationError::Io {
13058 path: self.color_path.clone(),
13059 source,
13060 })?;
13061 let color_index = u32::try_from(self.color_runs)
13062 .map_err(|_| SerialCollationError::MalformedCoordBucket(self.color_path.clone()))?;
13063 if color_index >= 0x3fff_ffff {
13064 return Err(SerialCollationError::MalformedCoordBucket(
13065 self.color_path.clone(),
13066 ));
13067 }
13068 self.color_runs += u64::from(count);
13069 self.write_record_with_color_index(record, label, color_index, count)
13070 }
13071
13072 #[allow(dead_code)]
13074 fn write_encoded_colored_record(
13075 &mut self,
13076 record: &StitchedCoordRecord,
13077 label: &[u8],
13078 count: u32,
13079 encoded_colors: &[u8],
13080 ) -> Result<(), SerialCollationError> {
13081 if encoded_colors.len() != count as usize * std::mem::size_of::<u64>() {
13082 return Err(SerialCollationError::MalformedCoordBucket(
13083 self.color_path.clone(),
13084 ));
13085 }
13086 if self.color_out.is_none() {
13087 let file = OpenOptions::new()
13088 .create(true)
13089 .append(true)
13090 .open(&self.color_path)
13091 .map_err(|source| SerialCollationError::Io {
13092 path: self.color_path.clone(),
13093 source,
13094 })?;
13095 self.color_out = Some(BufWriter::with_capacity(1024 * 1024, file));
13096 }
13097 let color_out = self
13098 .color_out
13099 .as_mut()
13100 .expect("color writer was just created");
13101 color_out
13102 .write_all(encoded_colors)
13103 .map_err(|source| SerialCollationError::Io {
13104 path: self.color_path.clone(),
13105 source,
13106 })?;
13107 let color_index = u32::try_from(self.color_runs)
13108 .map_err(|_| SerialCollationError::MalformedCoordBucket(self.color_path.clone()))?;
13109 if color_index >= 0x3fff_ffff {
13110 return Err(SerialCollationError::MalformedCoordBucket(
13111 self.color_path.clone(),
13112 ));
13113 }
13114 self.color_runs += u64::from(count);
13115 self.write_record_with_color_index(record, label, color_index, count)
13116 }
13117
13118 fn write_pending_batch(
13119 &mut self,
13120 batch: &mut PendingMaterializedBucket,
13121 ) -> Result<(), SerialCollationError> {
13122 let color_base = u32::try_from(self.color_runs)
13123 .map_err(|_| SerialCollationError::MalformedCoordBucket(self.color_path.clone()))?;
13124 for pending in &batch.records {
13125 let color_index = if pending.color_start == u32::MAX {
13126 u32::MAX
13127 } else {
13128 color_base.checked_add(pending.color_start).ok_or_else(|| {
13129 SerialCollationError::MalformedCoordBucket(self.color_path.clone())
13130 })?
13131 };
13132 if color_index != u32::MAX && color_index >= 0x3fff_ffff {
13133 return Err(SerialCollationError::MalformedCoordBucket(
13134 self.color_path.clone(),
13135 ));
13136 }
13137 self.record_buffer
13138 .extend_from_slice(&encoded_materialized_stitched_coord_record(
13139 MaterializedStitchedCoordRecord {
13140 path_id: pending.path_id,
13141 rank: u64::from(pending.rank),
13142 label_offset: self.label_bytes + u64::from(pending.label_offset),
13143 label_len: u32::from(pending.label_len),
13144 reverse: pending.reverse(),
13145 is_cycle: pending.is_cycle(),
13146 color_index,
13147 color_count: pending.color_count(),
13148 },
13149 ));
13150 }
13151
13152 self.label_out
13153 .as_mut()
13154 .expect("label writer is open")
13155 .write_all(&batch.labels)
13156 .map_err(|source| SerialCollationError::Io {
13157 path: self.label_path.clone(),
13158 source,
13159 })?;
13160 if !batch.colors.is_empty() {
13161 if self.color_out.is_none() {
13162 let file = OpenOptions::new()
13163 .create(true)
13164 .append(true)
13165 .open(&self.color_path)
13166 .map_err(|source| SerialCollationError::Io {
13167 path: self.color_path.clone(),
13168 source,
13169 })?;
13170 self.color_out = Some(BufWriter::with_capacity(1024 * 1024, file));
13171 }
13172 let color_out = self
13173 .color_out
13174 .as_mut()
13175 .expect("color writer was just created");
13176 color_out
13177 .write_all(unitig_colors_as_bytes(&batch.colors))
13178 .map_err(|source| SerialCollationError::Io {
13179 path: self.color_path.clone(),
13180 source,
13181 })?;
13182 }
13183 self.records += batch.records.len() as u64;
13184 self.label_bytes += batch.labels.len() as u64;
13185 self.color_runs += batch.colors.len() as u64;
13186 batch.records.clear();
13187 batch.labels.clear();
13188 batch.colors.clear();
13189 if self.record_buffer.len() >= STITCH_COORD_RECORD_WRITE_BUFFER {
13190 self.flush_record_buffer()?;
13191 }
13192 Ok(())
13193 }
13194
13195 fn flush_record_buffer(&mut self) -> Result<(), SerialCollationError> {
13196 if self.record_buffer.is_empty() {
13197 return Ok(());
13198 }
13199 self.coord_out
13200 .as_mut()
13201 .expect("coord writer is open")
13202 .write_all(&self.record_buffer)
13203 .map_err(|source| SerialCollationError::Io {
13204 path: self.coord_path.clone(),
13205 source,
13206 })?;
13207 self.record_buffer.clear();
13208 Ok(())
13209 }
13210
13211 fn flush_and_close(&mut self) -> Result<(), SerialCollationError> {
13212 self.flush_record_buffer()?;
13213 if let Some(mut coord_out) = self.coord_out.take() {
13214 coord_out
13215 .flush()
13216 .map_err(|source| SerialCollationError::Io {
13217 path: self.coord_path.clone(),
13218 source,
13219 })?;
13220 }
13221 if let Some(mut label_out) = self.label_out.take() {
13222 label_out
13223 .flush()
13224 .map_err(|source| SerialCollationError::Io {
13225 path: self.label_path.clone(),
13226 source,
13227 })?;
13228 }
13229 if let Some(mut color_out) = self.color_out.take() {
13230 color_out
13231 .flush()
13232 .map_err(|source| SerialCollationError::Io {
13233 path: self.color_path.clone(),
13234 source,
13235 })?;
13236 }
13237 Ok(())
13238 }
13239
13240 fn finish(mut self) -> Result<MaterializedStitchedCoordBucketEntry, SerialCollationError> {
13241 self.flush_and_close()?;
13242 let mut coord_out = BufWriter::with_capacity(
13243 MATERIALIZED_STITCH_COORD_SHARD_WRITE_BUFFER,
13244 OpenOptions::new()
13245 .write(true)
13246 .open(&self.coord_path)
13247 .map_err(|source| SerialCollationError::Io {
13248 path: self.coord_path.clone(),
13249 source,
13250 })?,
13251 );
13252 coord_out
13253 .seek(SeekFrom::Start(16))
13254 .map_err(|source| SerialCollationError::Io {
13255 path: self.coord_path.clone(),
13256 source,
13257 })?;
13258 coord_out
13259 .write_all(&self.records.to_le_bytes())
13260 .and_then(|_| coord_out.write_all(&self.label_bytes.to_le_bytes()))
13261 .and_then(|_| coord_out.flush())
13262 .map_err(|source| SerialCollationError::Io {
13263 path: self.coord_path.clone(),
13264 source,
13265 })?;
13266 Ok(MaterializedStitchedCoordBucketEntry {
13267 bucket_id: self.bucket_id,
13268 records: self.records,
13269 label_bytes: self.label_bytes,
13270 coord_path: self.coord_path,
13271 label_path: self.label_path,
13272 color_path: (self.color_runs != 0).then_some(self.color_path),
13273 color_runs: self.color_runs,
13274 })
13275 }
13276}
13277
13278fn stitched_coord_worker_buffer_records(bucket_count: usize, workers: usize) -> usize {
13287 const DEFAULT_STAGING_BUDGET: usize = 2 * 1024 * 1024 * 1024;
13288 const MIN_BUFFER_BYTES: usize = 8 * 1024;
13289 let budget = std::env::var("CF3_RS_STITCH_STAGING_BYTES")
13290 .ok()
13291 .and_then(|value| value.parse::<usize>().ok())
13292 .filter(|value| *value > 0)
13293 .unwrap_or(DEFAULT_STAGING_BUDGET);
13294 let pairs = bucket_count.max(1).saturating_mul(workers.max(1) + 1);
13295 let bytes = (budget / pairs).clamp(MIN_BUFFER_BYTES, EDGE_PATH_INFO_WORKER_BUFFER);
13296 bytes
13297 .div_ceil(std::mem::size_of::<StitchedCoordRecord>())
13298 .max(1)
13299}
13300
13301struct StitchedCoordShardWriters<'a> {
13302 coord_dir: &'a Path,
13303 worker_id: usize,
13304 writers: Vec<Option<StitchedCoordShardWriter>>,
13305}
13306
13307struct ConcurrentStitchedCoordWriters<'a> {
13308 coord_dir: &'a Path,
13309 writers: Vec<Mutex<Option<StitchedCoordShardWriter>>>,
13310 worker_buffers: Vec<UnsafeCell<Vec<Option<Vec<StitchedCoordRecord>>>>>,
13311 worker_buffer_capacity: usize,
13313}
13314
13315unsafe impl Sync for ConcurrentStitchedCoordWriters<'_> {}
13318
13319impl<'a> ConcurrentStitchedCoordWriters<'a> {
13320 fn new(coord_dir: &'a Path, bucket_count: usize, workers: usize) -> Self {
13321 Self {
13322 coord_dir,
13323 writers: (0..bucket_count).map(|_| Mutex::new(None)).collect(),
13324 worker_buffers: (0..workers.max(1) + 1)
13325 .map(|_| {
13326 let mut buckets = Vec::with_capacity(bucket_count);
13327 buckets.resize_with(bucket_count, || None);
13328 UnsafeCell::new(buckets)
13329 })
13330 .collect(),
13331 worker_buffer_capacity: stitched_coord_worker_buffer_records(bucket_count, workers),
13332 }
13333 }
13334
13335 fn flush_records(
13336 &self,
13337 bucket_id: usize,
13338 records: &[StitchedCoordRecord],
13339 ) -> Result<(), SerialCollationError> {
13340 let mut writer = self.writers[bucket_id]
13341 .lock()
13342 .map_err(|_| SerialCollationError::WorkerPanic)?;
13343 if writer.is_none() {
13344 *writer = Some(StitchedCoordShardWriter::create(
13345 self.coord_dir,
13346 0,
13347 bucket_id,
13348 )?);
13349 }
13350 let writer = writer.as_mut().expect("bucket writer was just created");
13351 for &record in records {
13352 writer.write_record(record)?;
13353 }
13354 Ok(())
13355 }
13356
13357 fn write_path_records(
13358 &self,
13359 bucket_id: usize,
13360 records: &[StitchedCoordRecord],
13361 ) -> Result<(), SerialCollationError> {
13362 if records.is_empty() {
13363 return Ok(());
13364 }
13365 let serial_slot = self.worker_buffers.len() - 1;
13366 let worker_id = rayon::current_thread_index()
13367 .filter(|&worker_id| worker_id < serial_slot)
13368 .unwrap_or(serial_slot);
13369 let buckets = unsafe { &mut *self.worker_buffers[worker_id].get() };
13372 let capacity = self.worker_buffer_capacity;
13373 let buffer = buckets[bucket_id].get_or_insert_with(|| Vec::with_capacity(capacity));
13374 let mut remaining = records;
13375 while !remaining.is_empty() {
13376 let take = (capacity - buffer.len()).min(remaining.len());
13377 buffer.extend_from_slice(&remaining[..take]);
13378 remaining = &remaining[take..];
13379 if buffer.len() == capacity {
13380 self.flush_records(bucket_id, buffer)?;
13381 buffer.clear();
13382 }
13383 }
13384 Ok(())
13385 }
13386
13387 fn write_record(
13388 &self,
13389 bucket_id: usize,
13390 record: StitchedCoordRecord,
13391 ) -> Result<(), SerialCollationError> {
13392 self.write_path_records(bucket_id, std::slice::from_ref(&record))
13393 }
13394
13395 fn finish(
13396 self,
13397 pool: &ThreadPool,
13398 ) -> Result<Vec<StitchedCoordBucketEntry>, SerialCollationError> {
13399 let bucket_count = self.writers.len();
13400 let mut pending_by_bucket = (0..bucket_count)
13401 .map(|_| Vec::<Vec<StitchedCoordRecord>>::new())
13402 .collect::<Vec<_>>();
13403 for worker in self.worker_buffers {
13404 for (bucket_id, records) in worker.into_inner().into_iter().enumerate() {
13405 if let Some(records) = records
13406 && !records.is_empty()
13407 {
13408 pending_by_bucket[bucket_id].push(records);
13409 }
13410 }
13411 }
13412 let writers = self
13413 .writers
13414 .into_iter()
13415 .map(|writer| {
13416 writer
13417 .into_inner()
13418 .map_err(|_| SerialCollationError::WorkerPanic)
13419 })
13420 .collect::<Result<Vec<_>, _>>()?;
13421 let coord_dir = self.coord_dir;
13422 let entries = pool.install(|| {
13423 writers
13424 .into_par_iter()
13425 .zip(pending_by_bucket.into_par_iter())
13426 .enumerate()
13427 .map(|(bucket_id, (writer, pending))| {
13428 if writer.is_none() && pending.is_empty() {
13429 return Ok(None);
13430 }
13431 let mut writer = match writer {
13432 Some(writer) => writer,
13433 None => StitchedCoordShardWriter::create(coord_dir, 0, bucket_id)?,
13434 };
13435 for records in pending {
13436 for record in records {
13437 writer.write_record(record)?;
13438 }
13439 }
13440 writer.finish().map(Some)
13441 })
13442 .collect::<Result<Vec<_>, SerialCollationError>>()
13443 })?;
13444 Ok(entries.into_iter().flatten().collect())
13445 }
13446}
13447
13448impl<'a> StitchedCoordShardWriters<'a> {
13449 fn new(coord_dir: &'a Path, worker_id: usize, bucket_count: usize) -> Self {
13450 let mut writers = Vec::with_capacity(bucket_count);
13451 writers.resize_with(bucket_count, || None);
13452 Self {
13453 coord_dir,
13454 worker_id,
13455 writers,
13456 }
13457 }
13458
13459 fn write_path_records(
13460 &mut self,
13461 bucket_id: usize,
13462 records: &[StitchedCoordRecord],
13463 ) -> Result<(), SerialCollationError> {
13464 if records.is_empty() {
13465 return Ok(());
13466 }
13467 if self.writers[bucket_id].is_none() {
13468 self.writers[bucket_id] = Some(StitchedCoordShardWriter::create(
13469 self.coord_dir,
13470 self.worker_id,
13471 bucket_id,
13472 )?);
13473 }
13474 let writer = self.writers[bucket_id]
13475 .as_mut()
13476 .expect("bucket writer was just created");
13477 for &record in records {
13478 writer.write_record(record)?;
13479 }
13480 Ok(())
13481 }
13482
13483 fn write_record(
13484 &mut self,
13485 bucket_id: usize,
13486 record: StitchedCoordRecord,
13487 ) -> Result<(), SerialCollationError> {
13488 if self.writers[bucket_id].is_none() {
13489 self.writers[bucket_id] = Some(StitchedCoordShardWriter::create(
13490 self.coord_dir,
13491 self.worker_id,
13492 bucket_id,
13493 )?);
13494 }
13495 self.writers[bucket_id]
13496 .as_mut()
13497 .expect("bucket writer was just created")
13498 .write_record(record)
13499 }
13500
13501 fn finish(self) -> Result<Vec<StitchedCoordBucketEntry>, SerialCollationError> {
13502 let mut manifest = Vec::new();
13503 for writer in self.writers.into_iter().flatten() {
13504 manifest.push(writer.finish()?);
13505 }
13506 Ok(manifest)
13507 }
13508}
13509
13510struct StitchedCoordShardWriter {
13511 bucket_id: usize,
13512 path: PathBuf,
13513 out: BufWriter<File>,
13514 record_buffer: Vec<u8>,
13515 records: u64,
13516}
13517
13518impl StitchedCoordShardWriter {
13519 fn create(
13520 coord_dir: &Path,
13521 worker_id: usize,
13522 bucket_id: usize,
13523 ) -> Result<Self, SerialCollationError> {
13524 let path = coord_dir.join(format!("{bucket_id:05}.{worker_id:03}.scb"));
13525 let file = File::create(&path).map_err(|source| SerialCollationError::Io {
13526 path: path.clone(),
13527 source,
13528 })?;
13529 let mut out = BufWriter::with_capacity(STITCH_COORD_SHARD_WRITE_BUFFER, file);
13530 out.write_all(STITCH_COORD_MAGIC)
13531 .map_err(|source| SerialCollationError::Io {
13532 path: path.clone(),
13533 source,
13534 })?;
13535 out.write_all(&(bucket_id as u64).to_le_bytes())
13536 .map_err(|source| SerialCollationError::Io {
13537 path: path.clone(),
13538 source,
13539 })?;
13540 out.write_all(&0u64.to_le_bytes())
13541 .map_err(|source| SerialCollationError::Io {
13542 path: path.clone(),
13543 source,
13544 })?;
13545 out.write_all(&[0u8; 8])
13546 .map_err(|source| SerialCollationError::Io {
13547 path: path.clone(),
13548 source,
13549 })?;
13550 Ok(Self {
13551 bucket_id,
13552 path,
13553 out,
13554 record_buffer: Vec::with_capacity(STITCH_COORD_RECORD_WRITE_BUFFER),
13555 records: 0,
13556 })
13557 }
13558
13559 fn write_record(&mut self, record: StitchedCoordRecord) -> Result<(), SerialCollationError> {
13560 self.record_buffer
13561 .extend_from_slice(&encoded_stitched_coord_record(record));
13562 self.records += 1;
13563 if self.record_buffer.len() >= STITCH_COORD_RECORD_WRITE_BUFFER {
13564 self.flush_record_buffer()?;
13565 }
13566 Ok(())
13567 }
13568
13569 fn flush_record_buffer(&mut self) -> Result<(), SerialCollationError> {
13570 if self.record_buffer.is_empty() {
13571 return Ok(());
13572 }
13573 self.out
13574 .write_all(&self.record_buffer)
13575 .map_err(|source| SerialCollationError::Io {
13576 path: self.path.clone(),
13577 source,
13578 })?;
13579 self.record_buffer.clear();
13580 Ok(())
13581 }
13582
13583 fn finish(mut self) -> Result<StitchedCoordBucketEntry, SerialCollationError> {
13584 self.flush_record_buffer()?;
13585 self.out
13586 .flush()
13587 .map_err(|source| SerialCollationError::Io {
13588 path: self.path.clone(),
13589 source,
13590 })?;
13591 self.out
13592 .seek(SeekFrom::Start(16))
13593 .map_err(|source| SerialCollationError::Io {
13594 path: self.path.clone(),
13595 source,
13596 })?;
13597 self.out
13598 .write_all(&self.records.to_le_bytes())
13599 .map_err(|source| SerialCollationError::Io {
13600 path: self.path.clone(),
13601 source,
13602 })?;
13603 self.out
13604 .flush()
13605 .map_err(|source| SerialCollationError::Io {
13606 path: self.path.clone(),
13607 source,
13608 })?;
13609 Ok(StitchedCoordBucketEntry {
13610 bucket_id: self.bucket_id,
13611 records: self.records,
13612 path: self.path,
13613 })
13614 }
13615}
13616
13617fn stitched_coord_bucket_count(threads: usize) -> usize {
13618 (threads.max(1) * 1024)
13619 .next_power_of_two()
13620 .clamp(1024, 16_384)
13621}
13622
13623fn stitch_endpoint_bucket_count(threads: usize) -> usize {
13624 let _ = threads;
13625 MAX_OPEN_STITCH_ENDPOINT_WRITERS
13626}
13627
13628fn materialized_stitched_coord_bucket_count(threads: usize) -> usize {
13629 let target = (threads.max(1) * 16).next_power_of_two().clamp(32, 64);
13630 let max_open_buckets = MAX_OPEN_MATERIALIZED_STITCH_WRITERS
13631 .checked_div(threads.max(1))
13632 .unwrap_or(1)
13633 .max(1);
13634 target.min(floor_power_of_two(max_open_buckets)).max(1)
13635}
13636
13637fn floor_power_of_two(value: usize) -> usize {
13638 if value == 0 {
13639 0
13640 } else {
13641 1usize << (usize::BITS - 1 - value.leading_zeros())
13642 }
13643}
13644
13645fn final_unitig_bucket_count(threads: usize) -> usize {
13646 (threads.max(1) * 8)
13647 .next_power_of_two()
13648 .clamp(128, MAX_OPEN_STITCH_ENDPOINT_WRITERS)
13649}
13650
13651fn stitched_coord_bucket(path_id: u64, bucket_mask: usize) -> usize {
13652 (xxh3_64(&path_id.to_ne_bytes()) as usize) & bucket_mask
13653}
13654
13655fn encoded_stitched_coord_record(
13656 record: StitchedCoordRecord,
13657) -> [u8; STITCH_PATH_INFO_RECORD_LEN as usize] {
13658 let mut bytes = [0u8; STITCH_PATH_INFO_RECORD_LEN as usize];
13659 bytes[..8].copy_from_slice(&record.path_id.to_le_bytes());
13660 bytes[8..16].copy_from_slice(&record.rank.to_le_bytes());
13661 bytes[16..20].copy_from_slice(&record.unitig_index.to_le_bytes());
13662 let mut flags = 0u8;
13663 if record.reverse {
13664 flags |= STITCH_COORD_REVERSE_FLAG;
13665 }
13666 if record.is_cycle {
13667 flags |= STITCH_COORD_CYCLE_FLAG;
13668 }
13669 bytes[20] = flags;
13670 bytes
13671}
13672
13673fn encoded_materialized_stitched_coord_record(
13674 record: MaterializedStitchedCoordRecord,
13675) -> [u8; STITCH_COORD_RECORD_LEN as usize] {
13676 let mut bytes = [0u8; STITCH_COORD_RECORD_LEN as usize];
13677 bytes[..8].copy_from_slice(&record.path_id.to_le_bytes());
13678 debug_assert!(u32::try_from(record.label_offset).is_ok());
13679 bytes[8..12].copy_from_slice(&(record.label_offset as u32).to_le_bytes());
13680 bytes[12..16].copy_from_slice(&record.color_index.to_le_bytes());
13681 bytes[16..18].copy_from_slice(
13682 &u16::try_from(record.rank)
13683 .expect("materialized rank fits C++ weight_t")
13684 .to_le_bytes(),
13685 );
13686 bytes[18..20].copy_from_slice(
13687 &u16::try_from(record.label_len)
13688 .expect("materialized label length fits C++ uni_len_t")
13689 .to_le_bytes(),
13690 );
13691 bytes[20..22].copy_from_slice(
13692 &u16::try_from(record.color_count)
13693 .expect("materialized color count fits C++ uni_len_t")
13694 .to_le_bytes(),
13695 );
13696 let mut flags = 0u16;
13697 if record.reverse {
13698 flags |= LoadedMaterializedStitchedCoordRecord::REVERSE_FLAG;
13699 }
13700 if record.is_cycle {
13701 flags |= LoadedMaterializedStitchedCoordRecord::CYCLE_FLAG;
13702 }
13703 bytes[22..24].copy_from_slice(&flags.to_le_bytes());
13704 bytes
13705}
13706
13707fn reduce_materialized_stitched_coord_bucket_files_to_final<const K: usize>(
13708 manifest: &[MaterializedStitchedCoordBucketEntry],
13709 retained: &[Vec<PendingMaterializedBucket>],
13710 threads: usize,
13711 final_buckets: &mut FinalUnitigBucketWriters,
13712) -> Result<u64, SerialCollationError> {
13713 if manifest.is_empty() && retained.iter().all(Vec::is_empty) {
13714 return Ok(0);
13715 }
13716
13717 let bucket_count = manifest
13718 .iter()
13719 .map(|entry| entry.bucket_id + 1)
13720 .max()
13721 .unwrap_or(0)
13722 .max(retained.len());
13723 let mut files_by_bucket = vec![Vec::new(); bucket_count];
13724 for entry in manifest {
13725 files_by_bucket[entry.bucket_id].push(entry.clone());
13726 }
13727 let groups = files_by_bucket
13728 .into_iter()
13729 .enumerate()
13730 .filter(|(bucket_id, files)| {
13731 !files.is_empty()
13732 || retained
13733 .get(*bucket_id)
13734 .is_some_and(|tails| !tails.is_empty())
13735 })
13736 .collect::<Vec<_>>();
13737
13738 let workers = threads.max(1).min(groups.len());
13739 if final_buckets.direct_output.is_some() && workers > 1 {
13740 return reduce_materialized_groups_to_direct_fasta::<K>(
13741 &groups,
13742 retained,
13743 workers,
13744 final_buckets,
13745 );
13746 }
13747 let mut emitted = 0u64;
13748 if workers == 1 {
13749 for (bucket_id, files) in &groups {
13750 let unitigs = reduce_materialized_stitched_coord_bucket_file_group_with_tails::<K>(
13751 files,
13752 retained.get(*bucket_id).map_or(&[], Vec::as_slice),
13753 )?;
13754 for unitig in unitigs {
13755 if unitig.colors.is_empty() {
13756 final_buckets.write_label(&unitig.label)?;
13757 } else {
13758 final_buckets.write_colored_label(&unitig.label, &unitig.colors)?;
13759 }
13760 emitted += 1;
13761 }
13762 }
13763 return Ok(emitted);
13764 }
13765
13766 let next_group = AtomicUsize::new(0);
13767 let (tx, rx) =
13768 mpsc::sync_channel::<Result<Vec<FinalUnitigRecord>, SerialCollationError>>(workers * 2);
13769 std::thread::scope(|scope| {
13770 let mut handles = Vec::new();
13771 for _ in 0..workers {
13772 let tx = tx.clone();
13773 let next_group = &next_group;
13774 let groups = &groups;
13775 handles.push(
13776 scope.spawn(move || {
13777 loop {
13778 let group_idx = next_group.fetch_add(1, Ordering::Relaxed);
13779 let Some(group) = groups.get(group_idx) else {
13780 break;
13781 };
13782 let result =
13783 reduce_materialized_stitched_coord_bucket_file_group_with_tails::<K>(
13784 &group.1,
13785 retained.get(group.0).map_or(&[], Vec::as_slice),
13786 );
13787 let should_stop = result.is_err();
13788 if tx.send(result).is_err() || should_stop {
13789 break;
13790 }
13791 }
13792 }),
13793 );
13794 }
13795 drop(tx);
13796
13797 let mut first_error = None;
13798 for result in rx {
13799 match result {
13800 Ok(unitigs) if first_error.is_none() => {
13801 for unitig in unitigs {
13802 let result = if unitig.colors.is_empty() {
13803 final_buckets.write_label(&unitig.label)
13804 } else {
13805 final_buckets.write_colored_label(&unitig.label, &unitig.colors)
13806 };
13807 if let Err(err) = result {
13808 first_error = Some(err);
13809 break;
13810 }
13811 emitted += 1;
13812 }
13813 }
13814 Ok(_) => {}
13815 Err(err) if first_error.is_none() => first_error = Some(err),
13816 Err(_) => {}
13817 }
13818 }
13819
13820 for handle in handles {
13821 handle
13822 .join()
13823 .map_err(|_| SerialCollationError::WorkerPanic)?;
13824 }
13825
13826 if let Some(err) = first_error {
13827 Err(err)
13828 } else {
13829 Ok(())
13830 }
13831 })?;
13832
13833 Ok(emitted)
13834}
13835
13836struct EncodedFinalBatch {
13837 bytes: Vec<u8>,
13838 records: u64,
13839 bases: u64,
13840}
13841
13842struct DirectFinalBatchBuilder {
13843 bytes: Vec<u8>,
13844 records: u64,
13845 bases: u64,
13846}
13847
13848impl DirectFinalBatchBuilder {
13849 fn new(_first_record: u64) -> Self {
13850 Self {
13851 bytes: Vec::with_capacity(512 * 1024),
13852 records: 0,
13853 bases: 0,
13854 }
13855 }
13856
13857 fn push(&mut self, label: &[u8], colors: &[UnitigColor]) {
13858 self.bytes.extend_from_slice(b">0");
13859 for color in colors {
13860 self.bytes.push(b' ');
13861 append_decimal_u64(&mut self.bytes, color.raw());
13862 }
13863 self.bytes.push(b'\n');
13864 self.bytes.extend_from_slice(label);
13865 self.bytes.push(b'\n');
13866 self.records += 1;
13867 self.bases += label.len() as u64;
13868 }
13869
13870 fn finish(self) -> EncodedFinalBatch {
13871 EncodedFinalBatch {
13872 bytes: self.bytes,
13873 records: self.records,
13874 bases: self.bases,
13875 }
13876 }
13877}
13878
13879fn reduce_materialized_groups_to_direct_fasta<const K: usize>(
13880 groups: &[(usize, Vec<MaterializedStitchedCoordBucketEntry>)],
13881 retained: &[Vec<PendingMaterializedBucket>],
13882 workers: usize,
13883 final_buckets: &mut FinalUnitigBucketWriters,
13884) -> Result<u64, SerialCollationError> {
13885 const DIRECT_FINAL_BATCH_RECORDS: usize = 16 * 1024;
13886 let next_group = AtomicUsize::new(0);
13887 let next_record = AtomicU64::new(final_buckets.direct_record_id_highwater + 1);
13888 let (output, output_path, initial_offset) = final_buckets.prepare_parallel_direct_output()?;
13889 let next_offset = AtomicU64::new(initial_offset);
13890 let emitted = AtomicU64::new(0);
13891 let emitted_bases = AtomicU64::new(0);
13892 let load_ns = AtomicU64::new(0);
13893 let sort_ns = AtomicU64::new(0);
13894 let assemble_ns = AtomicU64::new(0);
13895 let send_ns = AtomicU64::new(0);
13896 let write_started = Instant::now();
13897 let worker_timings = std::thread::scope(|scope| {
13898 let mut handles = Vec::with_capacity(workers);
13899 for _ in 0..workers {
13900 let next_group = &next_group;
13901 let next_record = &next_record;
13902 let next_offset = &next_offset;
13903 let emitted = &emitted;
13904 let emitted_bases = &emitted_bases;
13905 let output = &output;
13906 let output_path = &output_path;
13907 let load_ns = &load_ns;
13908 let sort_ns = &sort_ns;
13909 let assemble_ns = &assemble_ns;
13910 let send_ns = &send_ns;
13911 handles.push(scope.spawn(move || {
13912 let worker_started = Instant::now();
13913 let mut groups_processed = 0u64;
13914 let mut max_group_ns = 0u64;
13915 let mut max_load_ns = 0u64;
13916 let mut max_sort_ns = 0u64;
13917 let mut max_assemble_ns = 0u64;
13918 let mut max_send_ns = 0u64;
13919 let mut shard = MaterializedStitchedCoordBucket::default();
13920 loop {
13921 let group_idx = next_group.fetch_add(1, Ordering::Relaxed);
13922 let Some(group) = groups.get(group_idx) else {
13923 break;
13924 };
13925 let group_started = Instant::now();
13926 let started = Instant::now();
13927 load_materialized_stitched_coord_bucket_file_group_with_tails_into(
13928 &group.1,
13929 retained.get(group.0).map_or(&[], Vec::as_slice),
13930 &mut shard,
13931 )?;
13932 let elapsed_ns = started.elapsed().as_nanos() as u64;
13933 load_ns.fetch_add(elapsed_ns, Ordering::Relaxed);
13934 max_load_ns = max_load_ns.max(elapsed_ns);
13935 let started = Instant::now();
13936 shard
13937 .records
13938 .sort_unstable_by_key(|record| (record.path_id, record.rank));
13939 let elapsed_ns = started.elapsed().as_nanos() as u64;
13940 sort_ns.fetch_add(elapsed_ns, Ordering::Relaxed);
13941 max_sort_ns = max_sort_ns.max(elapsed_ns);
13942 let mut batch = None::<DirectFinalBatchBuilder>;
13943 let mut write_error = None;
13944 let started = Instant::now();
13945 reduce_sorted_materialized_stitched_coord_bucket_with::<K, _>(
13946 &mut shard.records,
13947 &shard.labels,
13948 &shard.colors,
13949 |label, colors| {
13950 if write_error.is_some() {
13951 return;
13952 }
13953 let builder = batch.get_or_insert_with(|| {
13954 let first_record = next_record.fetch_add(
13955 DIRECT_FINAL_BATCH_RECORDS as u64,
13956 Ordering::Relaxed,
13957 );
13958 DirectFinalBatchBuilder::new(first_record)
13959 });
13960 builder.push(label, colors);
13961 if builder.records as usize >= DIRECT_FINAL_BATCH_RECORDS {
13962 let write_started = Instant::now();
13963 let encoded = batch.take().expect("batch exists").finish();
13964 let result = write_encoded_final_batch_at(
13965 output,
13966 output_path,
13967 next_offset,
13968 emitted,
13969 emitted_bases,
13970 encoded,
13971 );
13972 let elapsed_ns = write_started.elapsed().as_nanos() as u64;
13973 send_ns.fetch_add(elapsed_ns, Ordering::Relaxed);
13974 max_send_ns = max_send_ns.max(elapsed_ns);
13975 if let Err(error) = result {
13976 write_error = Some(error);
13977 }
13978 }
13979 },
13980 );
13981 let elapsed_ns = started.elapsed().as_nanos() as u64;
13982 assemble_ns.fetch_add(elapsed_ns, Ordering::Relaxed);
13983 max_assemble_ns = max_assemble_ns.max(elapsed_ns);
13984 if let Some(error) = write_error {
13985 return Err(error);
13986 }
13987 if let Some(batch) = batch {
13988 let started = Instant::now();
13989 write_encoded_final_batch_at(
13990 output,
13991 output_path,
13992 next_offset,
13993 emitted,
13994 emitted_bases,
13995 batch.finish(),
13996 )?;
13997 let elapsed_ns = started.elapsed().as_nanos() as u64;
13998 send_ns.fetch_add(elapsed_ns, Ordering::Relaxed);
13999 max_send_ns = max_send_ns.max(elapsed_ns);
14000 }
14001 groups_processed += 1;
14002 max_group_ns = max_group_ns.max(group_started.elapsed().as_nanos() as u64);
14003 }
14004 Ok::<_, SerialCollationError>(ReducerWorkerTiming {
14005 elapsed_ns: worker_started.elapsed().as_nanos() as u64,
14006 groups: groups_processed,
14007 max_group_ns,
14008 max_load_ns,
14009 max_sort_ns,
14010 max_assemble_ns,
14011 max_send_ns,
14012 })
14013 }));
14014 }
14015 let mut first_error = None;
14016 let mut timings = Vec::with_capacity(handles.len());
14017 for handle in handles {
14018 let result = handle
14019 .join()
14020 .map_err(|_| SerialCollationError::WorkerPanic)?;
14021 match result {
14022 Ok(timing) => timings.push(timing),
14023 Err(err) if first_error.is_none() => first_error = Some(err),
14024 Err(_) => {}
14025 }
14026 }
14027 if let Some(err) = first_error {
14028 Err(err)
14029 } else {
14030 Ok(timings)
14031 }
14032 })?;
14033 final_buckets.direct_record_id_highwater =
14034 next_record.load(Ordering::Relaxed).saturating_sub(1);
14035 let emitted = emitted.load(Ordering::Relaxed);
14036 final_buckets.direct_records += emitted;
14037 final_buckets.direct_bases += emitted_bases.load(Ordering::Relaxed);
14038 eprintln!(
14039 "cuttlefish: path-info reduce worker detail: load {:.3}s, sort {:.3}s, assemble/encode {:.3}s, blocked send {:.3}s; reducer wall/write-drain {:.3}s",
14040 load_ns.load(Ordering::Relaxed) as f64 / 1e9,
14041 sort_ns.load(Ordering::Relaxed) as f64 / 1e9,
14042 assemble_ns.load(Ordering::Relaxed) as f64 / 1e9,
14043 send_ns.load(Ordering::Relaxed) as f64 / 1e9,
14044 write_started.elapsed().as_secs_f64(),
14045 );
14046 if let Some(slowest) = worker_timings.iter().max_by_key(|timing| timing.elapsed_ns) {
14047 eprintln!(
14048 "cuttlefish: path-info reduce critical worker: elapsed {:.3}s, groups {}; largest group {:.3}s (load {:.3}s, sort {:.3}s, assemble/write {:.3}s, write {:.3}s)",
14049 slowest.elapsed_ns as f64 / 1e9,
14050 slowest.groups,
14051 slowest.max_group_ns as f64 / 1e9,
14052 slowest.max_load_ns as f64 / 1e9,
14053 slowest.max_sort_ns as f64 / 1e9,
14054 slowest.max_assemble_ns as f64 / 1e9,
14055 slowest.max_send_ns as f64 / 1e9,
14056 );
14057 }
14058 Ok(emitted)
14059}
14060
14061struct ReducerWorkerTiming {
14062 elapsed_ns: u64,
14063 groups: u64,
14064 max_group_ns: u64,
14065 max_load_ns: u64,
14066 max_sort_ns: u64,
14067 max_assemble_ns: u64,
14068 max_send_ns: u64,
14069}
14070
14071fn write_encoded_final_batch_at(
14072 output: &File,
14073 output_path: &Path,
14074 next_offset: &AtomicU64,
14075 emitted: &AtomicU64,
14076 emitted_bases: &AtomicU64,
14077 batch: EncodedFinalBatch,
14078) -> Result<(), SerialCollationError> {
14079 let offset = next_offset.fetch_add(batch.bytes.len() as u64, Ordering::Relaxed);
14080 output
14081 .write_all_at(&batch.bytes, offset)
14082 .map_err(|source| SerialCollationError::Io {
14083 path: output_path.to_path_buf(),
14084 source,
14085 })?;
14086 emitted.fetch_add(batch.records, Ordering::Relaxed);
14087 emitted_bases.fetch_add(batch.bases, Ordering::Relaxed);
14088 Ok(())
14089}
14090
14091fn encode_final_unitig_batch(
14092 unitigs: Vec<FinalUnitigRecord>,
14093 _first_record: u64,
14094) -> EncodedFinalBatch {
14095 let records = unitigs.len() as u64;
14096 let bases = unitigs.iter().map(|unitig| unitig.label.len() as u64).sum();
14097 let mut bytes = Vec::with_capacity(bases as usize + unitigs.len() * 32);
14098 for unitig in unitigs {
14099 bytes.extend_from_slice(b">0");
14100 for color in unitig.colors {
14101 bytes.push(b' ');
14102 append_decimal_u64(&mut bytes, color.raw());
14103 }
14104 bytes.push(b'\n');
14105 bytes.extend_from_slice(&unitig.label);
14106 bytes.push(b'\n');
14107 }
14108 EncodedFinalBatch {
14109 bytes,
14110 records,
14111 bases,
14112 }
14113}
14114
14115#[allow(dead_code)]
14117fn reduce_materialized_stitched_coord_bucket_file_group<const K: usize>(
14118 group: &[MaterializedStitchedCoordBucketEntry],
14119) -> Result<Vec<FinalUnitigRecord>, SerialCollationError> {
14120 let mut shard = load_materialized_stitched_coord_bucket_file_group(group)?;
14121 Ok(reduce_materialized_stitched_coord_bucket::<K>(
14122 &mut shard.records,
14123 &shard.labels,
14124 &shard.colors,
14125 ))
14126}
14127
14128fn reduce_materialized_stitched_coord_bucket_file_group_with_tails<const K: usize>(
14129 group: &[MaterializedStitchedCoordBucketEntry],
14130 tails: &[PendingMaterializedBucket],
14131) -> Result<Vec<FinalUnitigRecord>, SerialCollationError> {
14132 let mut shard = load_materialized_stitched_coord_bucket_file_group_with_tails(group, tails)?;
14133 Ok(reduce_materialized_stitched_coord_bucket::<K>(
14134 &mut shard.records,
14135 &shard.labels,
14136 &shard.colors,
14137 ))
14138}
14139
14140fn load_materialized_stitched_coord_bucket_file_group_with_tails(
14141 group: &[MaterializedStitchedCoordBucketEntry],
14142 tails: &[PendingMaterializedBucket],
14143) -> Result<MaterializedStitchedCoordBucket, SerialCollationError> {
14144 let mut shard = if group.is_empty() {
14145 MaterializedStitchedCoordBucket {
14146 records: Vec::new(),
14147 labels: Vec::new(),
14148 colors: Vec::new(),
14149 }
14150 } else {
14151 load_materialized_stitched_coord_bucket_file_group(group)?
14152 };
14153 for tail in tails {
14154 append_pending_materialized_bucket(&mut shard, tail)?;
14155 }
14156 Ok(shard)
14157}
14158
14159fn load_materialized_stitched_coord_bucket_file_group_with_tails_into(
14160 group: &[MaterializedStitchedCoordBucketEntry],
14161 tails: &[PendingMaterializedBucket],
14162 shard: &mut MaterializedStitchedCoordBucket,
14163) -> Result<(), SerialCollationError> {
14164 shard.records.clear();
14165 shard.labels.clear();
14166 shard.colors.clear();
14167 let total_records = group.iter().map(|entry| entry.records).sum::<u64>()
14168 + tails
14169 .iter()
14170 .map(|tail| tail.records.len() as u64)
14171 .sum::<u64>();
14172 let total_label_bytes = group.iter().map(|entry| entry.label_bytes).sum::<u64>()
14173 + tails
14174 .iter()
14175 .map(|tail| tail.labels.len() as u64)
14176 .sum::<u64>();
14177 let total_colors = group.iter().map(|entry| entry.color_runs).sum::<u64>()
14178 + tails
14179 .iter()
14180 .map(|tail| tail.colors.len() as u64)
14181 .sum::<u64>();
14182 shard.records.reserve(total_records as usize);
14183 shard.labels.reserve(total_label_bytes as usize);
14184 shard.colors.reserve(total_colors as usize);
14185 for entry in group {
14186 append_materialized_stitched_coord_bucket_file(entry, shard)?;
14187 remove_serial_file(&entry.coord_path)?;
14188 remove_serial_file(&entry.label_path)?;
14189 if let Some(path) = entry.color_path.as_ref() {
14190 remove_serial_file(path)?;
14191 }
14192 }
14193 for tail in tails {
14194 append_pending_materialized_bucket(shard, tail)?;
14195 }
14196 Ok(())
14197}
14198
14199fn append_pending_materialized_bucket(
14200 shard: &mut MaterializedStitchedCoordBucket,
14201 tail: &PendingMaterializedBucket,
14202) -> Result<(), SerialCollationError> {
14203 let label_base = u32::try_from(shard.labels.len()).map_err(|_| {
14204 SerialCollationError::MalformedCoordBucket(PathBuf::from("retained-materialized-bucket"))
14205 })?;
14206 let color_base = u32::try_from(shard.colors.len()).map_err(|_| {
14207 SerialCollationError::MalformedCoordBucket(PathBuf::from("retained-materialized-bucket"))
14208 })?;
14209 shard.records.reserve(tail.records.len());
14210 shard.labels.extend_from_slice(&tail.labels);
14211 for pending in &tail.records {
14212 let color_start = if pending.color_start == u32::MAX {
14213 u32::MAX
14214 } else {
14215 color_base.checked_add(pending.color_start).ok_or_else(|| {
14216 SerialCollationError::MalformedCoordBucket(PathBuf::from(
14217 "retained-materialized-bucket",
14218 ))
14219 })?
14220 };
14221 let label_offset = label_base
14222 .checked_add(pending.label_offset)
14223 .ok_or_else(|| {
14224 SerialCollationError::MalformedCoordBucket(PathBuf::from(
14225 "retained-materialized-bucket",
14226 ))
14227 })?;
14228 shard
14229 .records
14230 .push(LoadedMaterializedStitchedCoordRecord::new(
14231 pending.path_id,
14232 u64::from(pending.rank),
14233 label_offset,
14234 u32::from(pending.label_len),
14235 pending.reverse(),
14236 pending.is_cycle(),
14237 color_start,
14238 pending.color_count(),
14239 ));
14240 }
14241 shard.colors.extend_from_slice(&tail.colors);
14242 Ok(())
14243}
14244
14245fn load_materialized_stitched_coord_bucket_file_group(
14246 group: &[MaterializedStitchedCoordBucketEntry],
14247) -> Result<MaterializedStitchedCoordBucket, SerialCollationError> {
14248 if let [entry] = group {
14249 let shard = read_materialized_stitched_coord_bucket_file(entry)?;
14250 remove_serial_file(&entry.coord_path)?;
14251 remove_serial_file(&entry.label_path)?;
14252 if let Some(path) = entry.color_path.as_ref() {
14253 remove_serial_file(path)?;
14254 }
14255 return Ok(shard);
14256 }
14257 let total_records = group.iter().map(|entry| entry.records).sum::<u64>();
14258 let total_label_bytes = group.iter().map(|entry| entry.label_bytes).sum::<u64>();
14259 let mut records = Vec::with_capacity(total_records as usize);
14260 let mut labels = Vec::with_capacity(total_label_bytes as usize);
14261 let mut colors = Vec::new();
14262 for entry in group {
14263 let label_offset = u32::try_from(labels.len())
14264 .map_err(|_| SerialCollationError::MalformedCoordBucket(entry.label_path.clone()))?;
14265 let mut shard = read_materialized_stitched_coord_bucket_file(entry)?;
14266 remove_serial_file(&entry.coord_path)?;
14267 remove_serial_file(&entry.label_path)?;
14268 if let Some(path) = entry.color_path.as_ref() {
14269 remove_serial_file(path)?;
14270 }
14271 let color_offset = u32::try_from(colors.len())
14272 .map_err(|_| SerialCollationError::MalformedCoordBucket(entry.coord_path.clone()))?;
14273 for record in &mut shard.records {
14274 record.label_offset =
14275 record
14276 .label_offset
14277 .checked_add(label_offset)
14278 .ok_or_else(|| {
14279 SerialCollationError::MalformedCoordBucket(entry.label_path.clone())
14280 })?;
14281 if record.color_start != u32::MAX {
14282 record.color_start =
14283 record
14284 .color_start
14285 .checked_add(color_offset)
14286 .ok_or_else(|| {
14287 SerialCollationError::MalformedCoordBucket(entry.coord_path.clone())
14288 })?;
14289 }
14290 }
14291 records.extend(shard.records);
14292 labels.extend(shard.labels);
14293 colors.extend(shard.colors);
14294 }
14295 Ok(MaterializedStitchedCoordBucket {
14296 records,
14297 labels,
14298 colors,
14299 })
14300}
14301
14302#[derive(Default)]
14303struct MaterializedStitchedCoordBucket {
14304 records: Vec<LoadedMaterializedStitchedCoordRecord>,
14305 labels: Vec<u8>,
14306 colors: Vec<UnitigColor>,
14307}
14308
14309struct FinalUnitigRecord {
14310 label: Vec<u8>,
14311 colors: Vec<UnitigColor>,
14312}
14313
14314fn read_materialized_stitched_coord_bucket_file(
14315 entry: &MaterializedStitchedCoordBucketEntry,
14316) -> Result<MaterializedStitchedCoordBucket, SerialCollationError> {
14317 let mut bucket = MaterializedStitchedCoordBucket::default();
14318 append_materialized_stitched_coord_bucket_file(entry, &mut bucket)?;
14319 Ok(bucket)
14320}
14321
14322fn append_materialized_stitched_coord_bucket_file(
14323 entry: &MaterializedStitchedCoordBucketEntry,
14324 bucket: &mut MaterializedStitchedCoordBucket,
14325) -> Result<(), SerialCollationError> {
14326 let mut file = File::open(&entry.coord_path).map_err(|source| SerialCollationError::Io {
14327 path: entry.coord_path.clone(),
14328 source,
14329 })?;
14330 let actual_len = file
14331 .metadata()
14332 .map_err(|source| SerialCollationError::Io {
14333 path: entry.coord_path.clone(),
14334 source,
14335 })?
14336 .len();
14337 let mut header = [0u8; STITCH_COORD_HEADER_LEN as usize];
14338 file.read_exact(&mut header)
14339 .map_err(|source| SerialCollationError::Io {
14340 path: entry.coord_path.clone(),
14341 source,
14342 })?;
14343 if &header[..8] != MATERIALIZED_STITCH_COORD_MAGIC {
14344 return Err(SerialCollationError::MalformedCoordBucket(
14345 entry.coord_path.clone(),
14346 ));
14347 }
14348 let bucket_id = u64::from_le_bytes(header[8..16].try_into().expect("bucket ID")) as usize;
14349 let records = u64::from_le_bytes(header[16..24].try_into().expect("record count"));
14350 let label_bytes = u64::from_le_bytes(header[24..32].try_into().expect("label bytes"));
14351 if bucket_id != entry.bucket_id || records != entry.records || label_bytes != entry.label_bytes
14352 {
14353 return Err(SerialCollationError::MalformedCoordBucket(
14354 entry.coord_path.clone(),
14355 ));
14356 }
14357 let expected_len = STITCH_COORD_HEADER_LEN
14358 .checked_add(
14359 records
14360 .checked_mul(STITCH_COORD_RECORD_LEN)
14361 .ok_or_else(|| {
14362 SerialCollationError::MalformedCoordBucket(entry.coord_path.clone())
14363 })?,
14364 )
14365 .ok_or_else(|| SerialCollationError::MalformedCoordBucket(entry.coord_path.clone()))?;
14366 if actual_len != expected_len {
14367 return Err(SerialCollationError::MalformedCoordBucket(
14368 entry.coord_path.clone(),
14369 ));
14370 }
14371 let record_base = bucket.records.len();
14372 let label_base = u32::try_from(bucket.labels.len())
14373 .map_err(|_| SerialCollationError::MalformedCoordBucket(entry.label_path.clone()))?;
14374 bucket.records.reserve(records as usize);
14375 let coord_bytes = unsafe {
14377 std::slice::from_raw_parts_mut(
14378 bucket
14379 .records
14380 .spare_capacity_mut()
14381 .as_mut_ptr()
14382 .cast::<u8>(),
14383 records as usize * std::mem::size_of::<LoadedMaterializedStitchedCoordRecord>(),
14384 )
14385 };
14386 file.read_exact(coord_bytes)
14387 .map_err(|source| SerialCollationError::Io {
14388 path: entry.coord_path.clone(),
14389 source,
14390 })?;
14391 unsafe { bucket.records.set_len(record_base + records as usize) };
14393
14394 let mut label_file =
14395 File::open(&entry.label_path).map_err(|source| SerialCollationError::Io {
14396 path: entry.label_path.clone(),
14397 source,
14398 })?;
14399 let actual_label_len = label_file
14400 .metadata()
14401 .map_err(|source| SerialCollationError::Io {
14402 path: entry.label_path.clone(),
14403 source,
14404 })?
14405 .len();
14406 if actual_label_len != entry.label_bytes {
14407 return Err(SerialCollationError::MalformedCoordBucket(
14408 entry.label_path.clone(),
14409 ));
14410 }
14411 let label_record_base = bucket.labels.len();
14412 bucket.labels.reserve(entry.label_bytes as usize);
14413 let uninitialized_labels = unsafe {
14414 std::slice::from_raw_parts_mut(
14415 bucket.labels.spare_capacity_mut().as_mut_ptr().cast::<u8>(),
14416 entry.label_bytes as usize,
14417 )
14418 };
14419 label_file
14420 .read_exact(uninitialized_labels)
14421 .map_err(|source| SerialCollationError::Io {
14422 path: entry.label_path.clone(),
14423 source,
14424 })?;
14425 unsafe {
14427 bucket
14428 .labels
14429 .set_len(label_record_base + entry.label_bytes as usize)
14430 };
14431
14432 let color_base = u32::try_from(bucket.colors.len())
14433 .map_err(|_| SerialCollationError::MalformedCoordBucket(entry.coord_path.clone()))?;
14434 bucket.colors.reserve(entry.color_runs as usize);
14435 if let Some(color_path) = &entry.color_path {
14436 let mut color_file = File::open(color_path).map_err(|source| SerialCollationError::Io {
14437 path: color_path.clone(),
14438 source,
14439 })?;
14440 let expected_color_bytes = entry.color_runs * std::mem::size_of::<UnitigColor>() as u64;
14441 if color_file
14442 .metadata()
14443 .map_err(|source| SerialCollationError::Io {
14444 path: color_path.clone(),
14445 source,
14446 })?
14447 .len()
14448 != expected_color_bytes
14449 {
14450 return Err(SerialCollationError::MalformedCoordBucket(
14451 color_path.clone(),
14452 ));
14453 }
14454 let color_record_base = bucket.colors.len();
14455 let uninitialized_colors = unsafe {
14456 std::slice::from_raw_parts_mut(
14457 bucket.colors.spare_capacity_mut().as_mut_ptr().cast::<u8>(),
14458 expected_color_bytes as usize,
14459 )
14460 };
14461 color_file
14462 .read_exact(uninitialized_colors)
14463 .map_err(|source| SerialCollationError::Io {
14464 path: color_path.clone(),
14465 source,
14466 })?;
14467 unsafe {
14470 bucket
14471 .colors
14472 .set_len(color_record_base + entry.color_runs as usize)
14473 };
14474 }
14475
14476 for record in &mut bucket.records[record_base..] {
14477 record.label_offset = record
14478 .label_offset
14479 .checked_add(label_base)
14480 .ok_or_else(|| SerialCollationError::MalformedCoordBucket(entry.label_path.clone()))?;
14481 if record.color_start != u32::MAX {
14482 let color_end = u64::from(record.color_start) + u64::from(record.color_count());
14483 if color_end > entry.color_runs {
14484 return Err(SerialCollationError::MalformedCoordBucket(
14485 entry.coord_path.clone(),
14486 ));
14487 }
14488 record.color_start = color_base.checked_add(record.color_start).ok_or_else(|| {
14489 SerialCollationError::MalformedCoordBucket(entry.coord_path.clone())
14490 })?;
14491 }
14492 }
14493
14494 Ok(())
14495}
14496
14497#[allow(dead_code)]
14499fn decoded_materialized_stitched_coord_record(
14500 bytes: &[u8],
14501 path: &Path,
14502) -> Result<MaterializedStitchedCoordRecord, SerialCollationError> {
14503 if bytes.len() != STITCH_COORD_RECORD_LEN as usize {
14504 return Err(SerialCollationError::MalformedCoordBucket(
14505 path.to_path_buf(),
14506 ));
14507 }
14508 let path_id = u64::from_le_bytes(bytes[..8].try_into().expect("u64 path_id field"));
14509 let label_offset = u32::from_le_bytes(bytes[8..12].try_into().expect("u32 label offset field"));
14510 let color_index = u32::from_le_bytes(bytes[12..16].try_into().expect("u32 color index field"));
14511 let rank = u16::from_le_bytes(bytes[16..18].try_into().expect("u16 rank field"));
14512 let label_len = u16::from_le_bytes(bytes[18..20].try_into().expect("u16 label length field"));
14513 let color_count = u16::from_le_bytes(bytes[20..22].try_into().expect("u16 color count field"));
14514 let flags = u16::from_le_bytes(bytes[22..24].try_into().expect("u16 flags field"));
14515 Ok(MaterializedStitchedCoordRecord {
14516 path_id,
14517 rank: u64::from(rank),
14518 label_offset: u64::from(label_offset),
14519 label_len: u32::from(label_len),
14520 reverse: flags & LoadedMaterializedStitchedCoordRecord::REVERSE_FLAG != 0,
14521 is_cycle: flags & LoadedMaterializedStitchedCoordRecord::CYCLE_FLAG != 0,
14522 color_index,
14523 color_count: u32::from(color_count),
14524 })
14525}
14526
14527fn reduce_materialized_stitched_coord_bucket<const K: usize>(
14528 records: &mut [LoadedMaterializedStitchedCoordRecord],
14529 labels: &[u8],
14530 color_runs: &[UnitigColor],
14531) -> Vec<FinalUnitigRecord> {
14532 let mut unitigs = Vec::new();
14533 reduce_materialized_stitched_coord_bucket_with::<K, _>(
14534 records,
14535 labels,
14536 color_runs,
14537 |label, colors| {
14538 unitigs.push(FinalUnitigRecord {
14539 label: label.to_vec(),
14540 colors: colors.to_vec(),
14541 });
14542 },
14543 );
14544 unitigs
14545}
14546
14547fn reduce_materialized_stitched_coord_bucket_with<const K: usize, F>(
14548 records: &mut [LoadedMaterializedStitchedCoordRecord],
14549 labels: &[u8],
14550 color_runs: &[UnitigColor],
14551 emit: F,
14552) where
14553 F: FnMut(&[u8], &[UnitigColor]),
14554{
14555 if records.is_empty() {
14556 return;
14557 }
14558
14559 if !materialized_stitched_coord_records_are_ordered(records) {
14560 records.sort_by_key(|record| (record.path_id, record.rank));
14561 }
14562 reduce_sorted_materialized_stitched_coord_bucket_with::<K, _>(
14563 records, labels, color_runs, emit,
14564 );
14565}
14566
14567fn reduce_sorted_materialized_stitched_coord_bucket_with<const K: usize, F>(
14568 records: &mut [LoadedMaterializedStitchedCoordRecord],
14569 labels: &[u8],
14570 color_runs: &[UnitigColor],
14571 mut emit: F,
14572) where
14573 F: FnMut(&[u8], &[UnitigColor]),
14574{
14575 if records.is_empty() {
14576 return;
14577 }
14578 let mut label = Vec::new();
14579 let mut colors = Vec::new();
14580 let mut start = 0;
14581 while start < records.len() {
14582 let path_id = records[start].path_id;
14583 let is_cycle = records[start].is_cycle();
14584 let mut end = start + 1;
14585 while end < records.len() && records[end].path_id == path_id {
14586 end += 1;
14587 }
14588
14589 label.clear();
14590 colors.clear();
14591 if end - start == 2 && !is_cycle && records[start].rank == 0 && records[start + 1].rank == 0
14592 {
14593 for (idx, record) in records[start..end].iter().enumerate() {
14594 let label_start = record.label_offset as usize;
14595 let label_end = label_start + record.label_len as usize;
14596 let unitig_label = &labels[label_start..label_end];
14597 let reverse = if idx == 0 {
14598 record.reverse()
14599 } else {
14600 !record.reverse()
14601 };
14602 append_materialized_colors::<K>(
14603 &mut colors,
14604 label.len(),
14605 record,
14606 color_runs,
14607 reverse,
14608 );
14609 append_or_init_oriented_fast::<K>(&mut label, unitig_label, reverse);
14610 }
14611 } else {
14612 for record in &records[start..end] {
14613 let label_start = record.label_offset as usize;
14614 let label_end = label_start + record.label_len as usize;
14615 let unitig_label = &labels[label_start..label_end];
14616 append_materialized_colors::<K>(
14617 &mut colors,
14618 label.len(),
14619 record,
14620 color_runs,
14621 record.reverse(),
14622 );
14623 append_or_init_oriented_fast::<K>(&mut label, unitig_label, record.reverse());
14624 }
14625 }
14626
14627 if label.len() >= K {
14628 let colored = !colors.is_empty();
14629 if is_cycle && colored {
14630 label.pop();
14631 let vertex_count = label.len().saturating_sub(K - 1) as u32;
14632 while colors
14633 .last()
14634 .is_some_and(|run| run.offset() >= vertex_count)
14635 {
14636 colors.pop();
14637 }
14638 } else if is_cycle {
14639 label = normalize_stitched_cycle::<K>(&label);
14640 } else {
14641 let reverse = reverse_complement_is_less(&label);
14642 if reverse && colored {
14643 reverse_color_runs_in_place(&mut colors, (label.len() - K + 1) as u32);
14644 }
14645 if reverse {
14646 reverse_complement_label_in_place(&mut label);
14647 }
14648 }
14649 emit(&label, &colors);
14650 }
14651 start = end;
14652 }
14653}
14654
14655fn append_materialized_colors<const K: usize>(
14656 output: &mut Vec<UnitigColor>,
14657 output_label_len: usize,
14658 record: &LoadedMaterializedStitchedCoordRecord,
14659 color_runs: &[UnitigColor],
14660 reverse: bool,
14661) {
14662 if record.color_start == u32::MAX {
14663 return;
14664 }
14665 let start = record.color_start as usize;
14666 let Some(runs) = color_runs.get(start..start + record.color_count() as usize) else {
14667 return;
14668 };
14669 let unitig_vertex_count = record.label_len as usize - K + 1;
14670 if output.is_empty() {
14671 output.reserve(runs.len());
14672 if reverse {
14673 for index in (0..runs.len()).rev() {
14674 let end = runs
14675 .get(index + 1)
14676 .map_or(unitig_vertex_count as u32, |next| next.offset());
14677 output.push(UnitigColor::new(
14678 unitig_vertex_count as u32 - end,
14679 crate::state::ColorCoordinate::from_u40(runs[index].coordinate()),
14680 ));
14681 }
14682 } else {
14683 output.extend_from_slice(runs);
14684 }
14685 } else {
14686 let output_vertex_count = output_label_len - K + 1;
14687 append_color_runs(
14688 output,
14689 output_vertex_count as u32,
14690 runs,
14691 unitig_vertex_count as u32,
14692 reverse,
14693 );
14694 }
14695}
14696
14697fn materialized_stitched_coord_records_are_ordered(
14698 records: &[LoadedMaterializedStitchedCoordRecord],
14699) -> bool {
14700 records
14701 .windows(2)
14702 .all(|pair| (pair[0].path_id, pair[0].rank) <= (pair[1].path_id, pair[1].rank))
14703}
14704
14705fn reduce_stitched_coord_buckets<const K: usize>(
14706 inputs: &DiscontinuityInputs<K>,
14707 records: Vec<StitchedCoordRecord>,
14708 threads: usize,
14709) -> Vec<Vec<u8>> {
14710 if records.is_empty() {
14711 return Vec::new();
14712 }
14713
14714 let bucket_count = (threads.max(1) * 1024)
14715 .next_power_of_two()
14716 .clamp(1024, 16_384);
14717 let bucket_mask = bucket_count - 1;
14718 let mut buckets = (0..bucket_count).map(|_| Vec::new()).collect::<Vec<_>>();
14719 for record in records {
14720 let bucket = (hash_u64(record.path_id, 0) as usize) & bucket_mask;
14721 buckets[bucket].push(record);
14722 }
14723
14724 let workers = threads.max(1).min(bucket_count);
14725 let mut reduced_buckets = Vec::new();
14726 if workers == 1 {
14727 for mut bucket in buckets {
14728 reduced_buckets.push(reduce_stitched_coord_bucket::<K>(
14729 inputs,
14730 &mut bucket,
14731 false,
14732 ));
14733 }
14734 } else {
14735 let chunk_size = bucket_count.div_ceil(workers);
14736 reduced_buckets = std::thread::scope(|scope| {
14737 let mut handles = Vec::new();
14738 for chunk in buckets.chunks_mut(chunk_size) {
14739 handles.push(scope.spawn(move || {
14740 let mut chunk_unitigs = Vec::new();
14741 for bucket in chunk {
14742 chunk_unitigs
14743 .extend(reduce_stitched_coord_bucket::<K>(inputs, bucket, false));
14744 }
14745 chunk_unitigs
14746 }));
14747 }
14748
14749 let mut reduced = Vec::new();
14750 for handle in handles {
14751 reduced.push(
14752 handle
14753 .join()
14754 .expect("stitched coordinate reducer worker panicked"),
14755 );
14756 }
14757 reduced
14758 });
14759 }
14760
14761 let total = reduced_buckets.iter().map(Vec::len).sum();
14762 let mut unitigs = Vec::with_capacity(total);
14763 for mut bucket in reduced_buckets {
14764 unitigs.append(&mut bucket);
14765 }
14766 unitigs
14767}
14768
14769fn reduce_stitched_coord_bucket_files<const K: usize>(
14770 inputs: &DiscontinuityInputs<K>,
14771 manifest: &[StitchedCoordBucketEntry],
14772 threads: usize,
14773) -> Result<Vec<Vec<u8>>, SerialCollationError> {
14774 if manifest.is_empty() {
14775 return Ok(Vec::new());
14776 }
14777
14778 let mut groups: Vec<Vec<StitchedCoordBucketEntry>> = Vec::new();
14779 for entry in manifest {
14780 if groups
14781 .last()
14782 .and_then(|group| group.first())
14783 .is_some_and(|first| first.bucket_id == entry.bucket_id)
14784 {
14785 groups
14786 .last_mut()
14787 .expect("checked that a final group exists")
14788 .push(entry.clone());
14789 } else {
14790 groups.push(vec![entry.clone()]);
14791 }
14792 }
14793
14794 let workers = threads.max(1).min(groups.len());
14795 let reduced_buckets = if workers == 1 {
14796 let mut reduced = Vec::with_capacity(groups.len());
14797 for group in &groups {
14798 reduced.push(reduce_stitched_coord_bucket_file_group::<K>(inputs, group)?);
14799 }
14800 reduced
14801 } else {
14802 let next_group = AtomicUsize::new(0);
14803 std::thread::scope(|scope| {
14804 let mut handles = Vec::new();
14805 for _ in 0..workers {
14806 let next_group = &next_group;
14807 let groups = &groups;
14808 handles.push(scope.spawn(move || {
14809 let mut local = Vec::new();
14810 loop {
14811 let group_idx = next_group.fetch_add(1, Ordering::Relaxed);
14812 let Some(group) = groups.get(group_idx) else {
14813 break;
14814 };
14815 local.extend(reduce_stitched_coord_bucket_file_group::<K>(inputs, group)?);
14816 }
14817 Ok::<_, SerialCollationError>(local)
14818 }));
14819 }
14820
14821 let mut reduced = Vec::new();
14822 for handle in handles {
14823 reduced.push(
14824 handle
14825 .join()
14826 .map_err(|_| SerialCollationError::WorkerPanic)??,
14827 );
14828 }
14829 Ok::<_, SerialCollationError>(reduced)
14830 })?
14831 };
14832
14833 let total = reduced_buckets.iter().map(Vec::len).sum();
14834 let mut unitigs = Vec::with_capacity(total);
14835 for mut bucket in reduced_buckets {
14836 unitigs.append(&mut bucket);
14837 }
14838 Ok(unitigs)
14839}
14840
14841fn reduce_stitched_coord_bucket_file_group<const K: usize>(
14842 inputs: &DiscontinuityInputs<K>,
14843 group: &[StitchedCoordBucketEntry],
14844) -> Result<Vec<Vec<u8>>, SerialCollationError> {
14845 let total_records = group.iter().map(|entry| entry.records).sum::<u64>();
14846 let mut records = Vec::with_capacity(total_records as usize);
14847 for entry in group {
14848 records.extend(read_stitched_coord_bucket_file(entry)?);
14849 }
14850 Ok(reduce_stitched_coord_bucket::<K>(
14851 inputs,
14852 &mut records,
14853 true,
14854 ))
14855}
14856
14857fn read_stitched_coord_bucket_file(
14858 entry: &StitchedCoordBucketEntry,
14859) -> Result<Vec<StitchedCoordRecord>, SerialCollationError> {
14860 let file = File::open(&entry.path).map_err(|source| SerialCollationError::Io {
14861 path: entry.path.clone(),
14862 source,
14863 })?;
14864 let mut input = BufReader::with_capacity(1024 * 1024, file);
14865 let mut magic = [0u8; 8];
14866 read_exact_coord(&mut input, &entry.path, &mut magic)?;
14867 if &magic != STITCH_COORD_MAGIC {
14868 return Err(SerialCollationError::MalformedCoordBucket(
14869 entry.path.clone(),
14870 ));
14871 }
14872
14873 let bucket_id = read_u64_coord(&mut input, &entry.path)? as usize;
14874 let records = read_u64_coord(&mut input, &entry.path)?;
14875 let _reserved = read_u64_coord(&mut input, &entry.path)?;
14876 if bucket_id != entry.bucket_id || records != entry.records {
14877 return Err(SerialCollationError::MalformedCoordBucket(
14878 entry.path.clone(),
14879 ));
14880 }
14881
14882 let expected_len = STITCH_COORD_HEADER_LEN
14883 .checked_add(
14884 records
14885 .checked_mul(STITCH_PATH_INFO_RECORD_LEN)
14886 .ok_or_else(|| SerialCollationError::MalformedCoordBucket(entry.path.clone()))?,
14887 )
14888 .ok_or_else(|| SerialCollationError::MalformedCoordBucket(entry.path.clone()))?;
14889 let actual_len = input
14890 .get_ref()
14891 .metadata()
14892 .map_err(|source| SerialCollationError::Io {
14893 path: entry.path.clone(),
14894 source,
14895 })?
14896 .len();
14897 if actual_len != expected_len {
14898 return Err(SerialCollationError::MalformedCoordBucket(
14899 entry.path.clone(),
14900 ));
14901 }
14902
14903 let mut coord_bytes = vec![0u8; (records * STITCH_PATH_INFO_RECORD_LEN) as usize];
14904 input
14905 .read_exact(&mut coord_bytes)
14906 .map_err(|source| SerialCollationError::Io {
14907 path: entry.path.clone(),
14908 source,
14909 })?;
14910 let mut out = Vec::with_capacity(records as usize);
14911 for chunk in coord_bytes.chunks_exact(STITCH_PATH_INFO_RECORD_LEN as usize) {
14912 out.push(decoded_stitched_coord_record(chunk, &entry.path)?);
14913 }
14914 Ok(out)
14915}
14916
14917fn read_stitched_coord_bucket_group(
14918 group: &[StitchedCoordBucketEntry],
14919) -> Result<Vec<StitchedCoordRecord>, SerialCollationError> {
14920 if let [entry] = group {
14921 return read_stitched_coord_bucket_file(entry);
14922 }
14923 let total_records = group.iter().map(|entry| entry.records).sum::<u64>();
14924 let mut records = Vec::with_capacity(total_records as usize);
14925 for entry in group {
14926 records.extend(read_stitched_coord_bucket_file(entry)?);
14927 }
14928 Ok(records)
14929}
14930
14931fn read_stitched_coord_bucket_group_dense(
14932 group: &[StitchedCoordBucketEntry],
14933 unitigs: usize,
14934 unitig_path: &Path,
14935 dense: &mut Vec<DenseLocalPathInfo>,
14936) -> Result<(), SerialCollationError> {
14937 dense.clear();
14938 dense.resize(unitigs, DenseLocalPathInfo::EMPTY);
14939 for entry in group {
14940 let file = File::open(&entry.path).map_err(|source| SerialCollationError::Io {
14941 path: entry.path.clone(),
14942 source,
14943 })?;
14944 let mut input = BufReader::with_capacity(1024 * 1024, file);
14945 let mut magic = [0u8; 8];
14946 read_exact_coord(&mut input, &entry.path, &mut magic)?;
14947 let bucket_id = read_u64_coord(&mut input, &entry.path)? as usize;
14948 let records = read_u64_coord(&mut input, &entry.path)?;
14949 let _reserved = read_u64_coord(&mut input, &entry.path)?;
14950 let expected_len = STITCH_COORD_HEADER_LEN
14951 .checked_add(
14952 records
14953 .checked_mul(STITCH_PATH_INFO_RECORD_LEN)
14954 .ok_or_else(|| {
14955 SerialCollationError::MalformedCoordBucket(entry.path.clone())
14956 })?,
14957 )
14958 .ok_or_else(|| SerialCollationError::MalformedCoordBucket(entry.path.clone()))?;
14959 let actual_len = input
14960 .get_ref()
14961 .metadata()
14962 .map_err(|source| SerialCollationError::Io {
14963 path: entry.path.clone(),
14964 source,
14965 })?
14966 .len();
14967 if &magic != STITCH_COORD_MAGIC
14968 || bucket_id != entry.bucket_id
14969 || records != entry.records
14970 || actual_len != expected_len
14971 {
14972 return Err(SerialCollationError::MalformedCoordBucket(
14973 entry.path.clone(),
14974 ));
14975 }
14976
14977 let mut bytes = vec![0u8; (records * STITCH_PATH_INFO_RECORD_LEN) as usize];
14978 read_exact_coord(&mut input, &entry.path, &mut bytes)?;
14979 for record in bytes.chunks_exact(STITCH_PATH_INFO_RECORD_LEN as usize) {
14980 let path_id = u64::from_le_bytes(record[..8].try_into().expect("path ID"));
14981 let rank = u64::from_le_bytes(record[8..16].try_into().expect("path rank"));
14982 let unitig_index =
14983 u32::from_le_bytes(record[16..20].try_into().expect("local unitig index")) as usize;
14984 let flags = record[20];
14985 let Some(slot) = dense.get_mut(unitig_index) else {
14986 return Err(SerialCollationError::MalformedCoordBucket(
14987 unitig_path.to_path_buf(),
14988 ));
14989 };
14990 if slot.rank_and_flags != u64::MAX || record[21..].iter().any(|&byte| byte != 0) {
14991 return Err(SerialCollationError::MalformedCoordBucket(
14992 entry.path.clone(),
14993 ));
14994 }
14995 *slot = DenseLocalPathInfo {
14996 path_id,
14997 rank_and_flags: (rank << 2)
14998 | u64::from((flags & STITCH_COORD_REVERSE_FLAG) != 0)
14999 | (u64::from((flags & STITCH_COORD_CYCLE_FLAG) != 0) << 1),
15000 };
15001 }
15002 }
15003 Ok(())
15004}
15005
15006fn decoded_stitched_coord_record(
15007 bytes: &[u8],
15008 path: &Path,
15009) -> Result<StitchedCoordRecord, SerialCollationError> {
15010 if bytes.len() != STITCH_PATH_INFO_RECORD_LEN as usize {
15011 return Err(SerialCollationError::MalformedCoordBucket(
15012 path.to_path_buf(),
15013 ));
15014 }
15015 let path_id = u64::from_le_bytes(bytes[..8].try_into().expect("u64 path_id field"));
15016 let rank = u64::from_le_bytes(bytes[8..16].try_into().expect("u64 rank field"));
15017 let unitig_index = u32::from_le_bytes(bytes[16..20].try_into().expect("u32 unitig field"));
15018 let flags = bytes[20];
15019 if bytes[21..].iter().any(|&byte| byte != 0) {
15020 return Err(SerialCollationError::MalformedCoordBucket(
15021 path.to_path_buf(),
15022 ));
15023 }
15024 Ok(StitchedCoordRecord {
15025 path_id,
15026 rank,
15027 unitig_index,
15028 reverse: (flags & STITCH_COORD_REVERSE_FLAG) != 0,
15029 is_cycle: (flags & STITCH_COORD_CYCLE_FLAG) != 0,
15030 })
15031}
15032
15033fn read_u64_coord(input: &mut BufReader<File>, path: &Path) -> Result<u64, SerialCollationError> {
15034 let mut bytes = [0u8; 8];
15035 read_exact_coord(input, path, &mut bytes)?;
15036 Ok(u64::from_le_bytes(bytes))
15037}
15038
15039fn read_exact_coord(
15040 input: &mut BufReader<File>,
15041 path: &Path,
15042 bytes: &mut [u8],
15043) -> Result<(), SerialCollationError> {
15044 input
15045 .read_exact(bytes)
15046 .map_err(|source| SerialCollationError::Io {
15047 path: path.to_path_buf(),
15048 source,
15049 })
15050}
15051
15052fn reduce_stitched_coord_bucket<const K: usize>(
15053 inputs: &DiscontinuityInputs<K>,
15054 records: &mut [StitchedCoordRecord],
15055 sort_records: bool,
15056) -> Vec<Vec<u8>> {
15057 if records.is_empty() {
15058 return Vec::new();
15059 }
15060
15061 if sort_records {
15062 records.sort_unstable_by_key(|record| (record.path_id, record.rank));
15063 }
15064 debug_assert!(
15065 records
15066 .windows(2)
15067 .all(|pair| (pair[0].path_id, pair[0].rank) <= (pair[1].path_id, pair[1].rank))
15068 );
15069
15070 let mut unitigs = Vec::new();
15071 let mut start = 0;
15072 while start < records.len() {
15073 let path_id = records[start].path_id;
15074 let is_cycle = records[start].is_cycle;
15075 let mut end = start + 1;
15076 while end < records.len() && records[end].path_id == path_id {
15077 end += 1;
15078 }
15079
15080 let mut label = Vec::new();
15081 for record in &records[start..end] {
15082 let unitig = &inputs.unitigs[record.unitig_index as usize];
15083 let unitig_label = unitig.label(inputs);
15084 let mut reverse = record.reverse;
15085 if !label.is_empty()
15086 && !labels_overlap_oriented_fast::<K>(&label, unitig_label, reverse)
15087 {
15088 let alternate = oriented_label(unitig_label, !reverse);
15089 if labels_overlap::<K>(&label, &alternate) {
15090 reverse = !reverse;
15091 }
15092 }
15093 append_or_init_oriented_fast::<K>(&mut label, unitig_label, reverse);
15094 }
15095
15096 if label.len() >= K {
15097 let label = if is_cycle {
15098 normalize_stitched_cycle::<K>(&label)
15099 } else {
15100 canonical_label(label)
15101 };
15102 unitigs.push(label);
15103 }
15104
15105 start = end;
15106 }
15107
15108 unitigs
15109}
15110
15111fn debug_endpoint<const K: usize>(endpoint: DiscontinuityEndpoint<K>) -> (String, Side) {
15112 (endpoint.vertex.to_ascii_string(), endpoint.side)
15113}
15114
15115fn endpoints_by_label_end<const K: usize>(
15116 unitig: &DiscontinuityUnitig<K>,
15117) -> (
15118 Option<DiscontinuityEndpoint<K>>,
15119 Option<DiscontinuityEndpoint<K>>,
15120) {
15121 (unitig.left_exit(), unitig.right_exit())
15122}
15123
15124fn push_stitch_edge(
15125 adjacency: &mut [StitchAdjacencyList],
15126 first: usize,
15127 second: usize,
15128 unitig_index: Option<usize>,
15129) {
15130 adjacency[first].push(StitchAdjacency {
15131 to: second,
15132 unitig_index,
15133 });
15134 adjacency[second].push(StitchAdjacency {
15135 to: first,
15136 unitig_index,
15137 });
15138}
15139
15140fn walk_stitched_path<const K: usize>(
15141 inputs: &DiscontinuityInputs<K>,
15142 adjacency: &[StitchAdjacencyList],
15143 visited_segments: &mut [bool],
15144 start: usize,
15145) -> Option<StitchedPath> {
15146 let mut label = Vec::new();
15147 let mut prev = None;
15148 let mut current = start;
15149 let mut emitted_segment = false;
15150 let mut is_cycle = false;
15151
15152 loop {
15153 let next = adjacency[current]
15154 .iter()
15155 .copied()
15156 .find(|edge| {
15157 Some(edge.to) != prev
15158 && edge
15159 .unitig_index
15160 .is_none_or(|unitig_index| !visited_segments[unitig_index])
15161 })
15162 .or_else(|| {
15163 adjacency[current].iter().copied().find(|edge| {
15164 edge.unitig_index
15165 .is_some_and(|unitig_index| !visited_segments[unitig_index])
15166 })
15167 });
15168
15169 let Some(edge) = next else {
15170 break;
15171 };
15172
15173 if let Some(unitig_index) = edge.unitig_index {
15174 if visited_segments[unitig_index] {
15175 break;
15176 }
15177 visited_segments[unitig_index] = true;
15178 let unitig = &inputs.unitigs[unitig_index];
15179 let unitig_label = unitig.label(inputs);
15180 let reverse = reverse_for_stitch_node(current);
15181 let mut append_reverse = reverse;
15182 if !label.is_empty()
15183 && !labels_overlap_oriented_fast::<K>(&label, unitig_label, reverse)
15184 {
15185 let alternate = oriented_label(unitig_label, !reverse);
15186 if labels_overlap::<K>(&label, &alternate) {
15187 append_reverse = !reverse;
15188 }
15189 }
15190 append_or_init_oriented_fast::<K>(&mut label, unitig_label, append_reverse);
15191 emitted_segment = true;
15192 }
15193
15194 prev = Some(current);
15195 current = edge.to;
15196 if emitted_segment && current == start {
15197 is_cycle = true;
15198 break;
15199 }
15200 }
15201
15202 emitted_segment.then_some(StitchedPath { label, is_cycle })
15203}
15204
15205fn walk_simple_stitched_component_coords(
15206 half_ends: &[HalfEnd],
15207 join_neighbor: &[u32],
15208 start: usize,
15209 path_id: u64,
15210 records: &mut Vec<StitchedCoordRecord>,
15211) {
15212 let record_start = records.len();
15213 let mut current = start;
15214 let mut rank = 0u64;
15215 let mut is_cycle = false;
15216
15217 loop {
15218 let unitig_index = half_ends[current].unitig_index;
15219 let reverse = reverse_for_stitch_node(current);
15220 records.push(StitchedCoordRecord {
15221 path_id,
15222 rank,
15223 unitig_index,
15224 reverse,
15225 is_cycle: false,
15226 });
15227 rank += 1;
15228
15229 let other = current ^ 1;
15230 let next = join_neighbor[other];
15231 if next == STITCH_NO_NODE {
15232 break;
15233 }
15234 if stitch_node_index(next) == start {
15235 is_cycle = true;
15236 break;
15237 }
15238 current = stitch_node_index(next);
15239 }
15240
15241 if rank == 0 {
15242 records.truncate(record_start);
15243 return;
15244 }
15245
15246 if is_cycle {
15247 for record in &mut records[record_start..] {
15248 record.is_cycle = true;
15249 }
15250 }
15251}
15252
15253fn walk_simple_stitched_component_coords_from_adjacency(
15254 adjacency: &[StitchAdjacencyList],
15255 start: usize,
15256 path_id: u64,
15257 records: &mut Vec<StitchedCoordRecord>,
15258) {
15259 let record_start = records.len();
15260 let mut prev = None;
15261 let mut current = start;
15262 let mut rank = 0u64;
15263 let mut emitted_segment = false;
15264 let mut is_cycle = false;
15265
15266 loop {
15267 let next = adjacency[current]
15268 .iter()
15269 .copied()
15270 .find(|edge| Some(edge.to) != prev);
15271
15272 let Some(edge) = next else {
15273 break;
15274 };
15275
15276 if let Some(unitig_index) = edge.unitig_index {
15277 let reverse = reverse_for_stitch_node(current);
15278 records.push(StitchedCoordRecord {
15279 path_id,
15280 rank,
15281 unitig_index: u32::try_from(unitig_index).expect("unitig index exceeds u32"),
15282 reverse,
15283 is_cycle: false,
15284 });
15285 rank += 1;
15286 emitted_segment = true;
15287 }
15288
15289 prev = Some(current);
15290 current = edge.to;
15291 if emitted_segment && current == start {
15292 is_cycle = true;
15293 break;
15294 }
15295 }
15296
15297 if !emitted_segment {
15298 records.truncate(record_start);
15299 return;
15300 }
15301
15302 if is_cycle {
15303 for record in &mut records[record_start..] {
15304 record.is_cycle = true;
15305 }
15306 }
15307}
15308
15309fn labels_overlap<const K: usize>(left: &[u8], right: &[u8]) -> bool {
15310 left.len() >= K && right.len() >= K && left[left.len() - K..] == right[..K]
15311}
15312
15313fn labels_overlap_oriented_fast<const K: usize>(left: &[u8], right: &[u8], reverse: bool) -> bool {
15314 if left.len() < K || right.len() < K {
15315 return false;
15316 }
15317
15318 let suffix = &left[left.len() - K..];
15319 if reverse {
15320 suffix
15321 .iter()
15322 .enumerate()
15323 .all(|(i, &base)| base == complement_ascii(right[right.len() - 1 - i]))
15324 } else {
15325 suffix == &right[..K]
15326 }
15327}
15328
15329fn append_or_init_oriented_fast<const K: usize>(label: &mut Vec<u8>, next: &[u8], reverse: bool) {
15330 let start = if label.is_empty() { 0 } else { K };
15331 if !reverse {
15332 label.extend_from_slice(&next[start..]);
15333 return;
15334 }
15335
15336 label.reserve(next.len().saturating_sub(start));
15337 for &base in next[..next.len() - start].iter().rev() {
15338 label.push(complement_ascii(base));
15339 }
15340}
15341
15342#[inline]
15343fn reverse_complement_label_in_place(label: &mut [u8]) {
15344 let mut left = 0;
15345 let mut right = label.len();
15346 while left < right {
15347 right -= 1;
15348 if left == right {
15349 label[left] = complement_ascii(label[left]);
15350 break;
15351 }
15352 let left_base = label[left];
15353 label[left] = complement_ascii(label[right]);
15354 label[right] = complement_ascii(left_base);
15355 left += 1;
15356 }
15357}
15358
15359fn normalize_stitched_cycle<const K: usize>(label: &[u8]) -> Vec<u8> {
15360 let mut graph = StitchCycleGraph::<K>::new();
15361 if graph.add_label(label).is_err() {
15362 return canonical_cycle_label(label.to_vec());
15363 }
15364
15365 let mut labels = graph.contract();
15366 labels.sort_by_key(|candidate| std::cmp::Reverse(candidate.len()));
15367 labels
15368 .into_iter()
15369 .next()
15370 .map(canonical_label)
15371 .unwrap_or_else(|| canonical_cycle_label(label.to_vec()))
15372}
15373
15374#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
15375struct StitchCycleEdge<const K: usize> {
15376 from: Kmer<K>,
15377 to: Kmer<K>,
15378}
15379
15380#[derive(Debug, Clone, Copy, PartialEq, Eq)]
15381struct StitchCycleKmer<const K: usize> {
15382 observed: Kmer<K>,
15383}
15384
15385impl<const K: usize> StitchCycleKmer<K> {
15386 fn new(observed: Kmer<K>) -> Self {
15387 Self { observed }
15388 }
15389
15390 fn canonical(self) -> Kmer<K> {
15391 self.observed.canonical()
15392 }
15393
15394 fn in_canonical_form(self) -> bool {
15395 self.observed.is_canonical()
15396 }
15397
15398 fn entrance_side(self) -> Side {
15399 if self.in_canonical_form() {
15400 Side::Front
15401 } else {
15402 Side::Back
15403 }
15404 }
15405
15406 fn roll_forward(self, base: Base) -> Self {
15407 Self::new(self.observed.roll_forward(base))
15408 }
15409}
15410
15411#[derive(Debug, Clone, PartialEq, Eq)]
15412struct StitchCycleWalk<const K: usize> {
15413 label: Vec<u8>,
15414 anchor: Kmer<K>,
15415}
15416
15417impl<const K: usize> StitchCycleWalk<K> {
15418 fn init(v: StitchCycleKmer<K>) -> Self {
15419 Self {
15420 label: v.observed.to_ascii_string().into_bytes(),
15421 anchor: v.canonical(),
15422 }
15423 }
15424
15425 fn extend(&mut self, v: StitchCycleKmer<K>, base: Base) -> bool {
15426 if v.canonical() == self.anchor {
15427 return false;
15428 }
15429
15430 self.label.push(base.to_ascii());
15431 true
15432 }
15433}
15434
15435struct StitchCycleGraph<const K: usize> {
15436 vertices: FastHashMap<Kmer<K>, VertexState>,
15437 unique_edges: BTreeSet<StitchCycleEdge<K>>,
15438}
15439
15440impl<const K: usize> StitchCycleGraph<K> {
15441 fn new() -> Self {
15442 Self {
15443 vertices: FastHashMap::default(),
15444 unique_edges: BTreeSet::new(),
15445 }
15446 }
15447
15448 fn add_label(&mut self, label: &[u8]) -> Result<(), ()> {
15449 if label.len() < K {
15450 return Ok(());
15451 }
15452
15453 let last_vertex_offset = label.len() - K;
15454 let mut prev = None;
15455 for offset in 0..=last_vertex_offset {
15456 let directed = StitchCycleKmer::new(
15457 Kmer::<K>::from_ascii(&label[offset..offset + K]).map_err(|_| ())?,
15458 );
15459 let canonical = directed.canonical();
15460 let pred_base = if offset == 0 {
15461 Base::E
15462 } else {
15463 Base::from_ascii(label[offset - 1])
15464 };
15465 let succ_base = if offset == last_vertex_offset {
15466 Base::E
15467 } else {
15468 Base::from_ascii(label[offset + K])
15469 };
15470 let mut front = if directed.in_canonical_form() {
15471 pred_base
15472 } else {
15473 succ_base.complement()
15474 };
15475 let mut back = if directed.in_canonical_form() {
15476 succ_base
15477 } else {
15478 pred_base.complement()
15479 };
15480
15481 if offset > 0 && Some(canonical) == prev {
15482 if directed.in_canonical_form() {
15483 front = Base::E;
15484 } else {
15485 back = Base::E;
15486 }
15487 }
15488
15489 self.vertices
15490 .entry(canonical)
15491 .or_default()
15492 .update_edges(front, back);
15493
15494 if let Some(from) = prev {
15495 self.unique_edges.insert(StitchCycleEdge {
15496 from,
15497 to: canonical,
15498 });
15499 }
15500 prev = Some(canonical);
15501 }
15502
15503 Ok(())
15504 }
15505
15506 fn contract(&mut self) -> Vec<Vec<u8>> {
15507 let mut unitigs = Vec::new();
15508 let mut vertices = self.vertices.keys().copied().collect::<Vec<_>>();
15509 vertices.sort_unstable();
15510
15511 for v_hat in vertices {
15512 let Some(state) = self.vertices.get(&v_hat).copied() else {
15513 continue;
15514 };
15515 if state.is_visited() || state.is_isolated(1) {
15516 continue;
15517 }
15518
15519 unitigs.push(self.extract_unitig(v_hat));
15520 }
15521
15522 unitigs
15523 }
15524
15525 fn extract_unitig(&mut self, v_hat: Kmer<K>) -> Vec<u8> {
15526 let (back_walk, back_is_cycle) = self.walk_unitig(v_hat, Side::Back);
15527 if back_is_cycle {
15528 return back_walk.label;
15529 }
15530
15531 let (front_walk, _) = self.walk_unitig(v_hat, Side::Front);
15532 let mut label = reverse_complement_label(&front_walk.label);
15533 label.extend_from_slice(&back_walk.label[K..]);
15534 label
15535 }
15536
15537 fn walk_unitig(&mut self, v_hat: Kmer<K>, start_side: Side) -> (StitchCycleWalk<K>, bool) {
15538 let icc_return_side = start_side.inverse();
15539 let mut v = if start_side == Side::Back {
15540 StitchCycleKmer::new(v_hat)
15541 } else {
15542 StitchCycleKmer::new(v_hat.reverse_complement())
15543 };
15544 let mut side = start_side;
15545 let mut walk = StitchCycleWalk::init(v);
15546
15547 loop {
15548 let canonical = v.canonical();
15549 let Some(state) = self.vertices.get(&canonical).copied() else {
15550 return (walk, false);
15551 };
15552 let Some(vertex_state) = self.vertices.get_mut(&canonical) else {
15553 return (walk, false);
15554 };
15555 vertex_state.mark_visited();
15556
15557 let mut edge = state.edge_at(side, 1);
15558 if edge == Base::N || edge == Base::E {
15559 return (walk, false);
15560 }
15561
15562 if side == Side::Front {
15563 edge = edge.complement();
15564 }
15565 v = v.roll_forward(edge);
15566
15567 let Some(next_state) = self.vertices.get(&v.canonical()).copied() else {
15568 return (walk, false);
15569 };
15570 side = v.entrance_side();
15571 if next_state.is_branching_side(side, 1) {
15572 return (walk, false);
15573 }
15574 if next_state.is_visited() {
15575 return (walk, v.canonical() == v_hat && side == icc_return_side);
15576 }
15577
15578 if !walk.extend(v, edge) {
15579 return (walk, false);
15580 }
15581 side = side.inverse();
15582 }
15583 }
15584}
15585
15586fn canonical_cycle_label(label: Vec<u8>) -> Vec<u8> {
15587 if label.is_empty() {
15588 return label;
15589 }
15590 let forward = minimal_rotation(&label);
15591 let reverse = minimal_rotation(&reverse_complement_label(&label));
15592 if reverse < forward { reverse } else { forward }
15593}
15594
15595fn minimal_rotation(label: &[u8]) -> Vec<u8> {
15596 let start = least_rotation_start(label);
15597 label[start..]
15598 .iter()
15599 .chain(label[..start].iter())
15600 .copied()
15601 .collect()
15602}
15603
15604fn least_rotation_start(s: &[u8]) -> usize {
15605 let n = s.len();
15606 if n <= 1 {
15607 return 0;
15608 }
15609
15610 let mut i = 0;
15611 let mut j = 1;
15612 let mut k = 0;
15613 while i < n && j < n && k < n {
15614 let a = s[(i + k) % n];
15615 let b = s[(j + k) % n];
15616 if a == b {
15617 k += 1;
15618 } else if a > b {
15619 i += k + 1;
15620 if i <= j {
15621 i = j + 1;
15622 }
15623 k = 0;
15624 } else {
15625 j += k + 1;
15626 if j <= i {
15627 j = i + 1;
15628 }
15629 k = 0;
15630 }
15631 }
15632
15633 i.min(j)
15634}
15635
15636#[derive(Debug, Clone, Copy, PartialEq, Eq)]
15637struct DiagonalOtherEnd<const K: usize> {
15638 vertex: Kmer<K>,
15639 side_at_vertex: Side,
15640 side_at_current: Side,
15641 weight: u64,
15642 unitig_index: usize,
15643 unitig_exit_side: Side,
15644 is_phi: bool,
15645}
15646
15647#[derive(Debug, Clone, Copy, PartialEq, Eq)]
15648struct PartitionOtherEnd<const K: usize> {
15649 endpoint: MatrixEndpoint<K>,
15650 side_at_current: Side,
15651 weight: u64,
15652 in_same_part: bool,
15653 processed: bool,
15654}
15655
15656struct PathInfoSlot<const K: usize> {
15657 tag: AtomicU64,
15658 key: UnsafeCell<MaybeUninit<Kmer<K>>>,
15659 value: UnsafeCell<MaybeUninit<PathInfo<K>>>,
15660}
15661
15662unsafe impl<const K: usize> Sync for PathInfoSlot<K> {}
15663
15664struct ConcurrentPathInfoTable<const K: usize> {
15665 slots: Vec<PathInfoSlot<K>>,
15666 mask: usize,
15667}
15668
15669struct CompactPathInfoSlot {
15670 key: AtomicU64,
15671 path_id: AtomicU64,
15672 rank_and_flags: AtomicU64,
15673 epoch: AtomicU8,
15674}
15675
15676const _: () = assert!(std::mem::size_of::<CompactPathInfoSlot>() == 32);
15677
15678#[derive(Clone, Copy)]
15679struct CompactExpansionPathInfo {
15680 path_id: u64,
15681 rank_and_flags: u64,
15682}
15683
15684impl CompactExpansionPathInfo {
15685 #[inline(always)]
15686 fn new(path_id: u64, rank: u64, exit_side: Side, is_cycle: bool) -> Self {
15687 let flags = u64::from(exit_side == Side::Back) | (u64::from(is_cycle) << 1);
15688 Self {
15689 path_id,
15690 rank_and_flags: (rank << 2) | flags,
15691 }
15692 }
15693
15694 #[inline(always)]
15695 fn rank(self) -> u64 {
15696 self.rank_and_flags >> 2
15697 }
15698
15699 #[inline(always)]
15700 fn exit_side(self) -> Side {
15701 if self.rank_and_flags & 1 == 0 {
15702 Side::Front
15703 } else {
15704 Side::Back
15705 }
15706 }
15707
15708 #[inline(always)]
15709 fn is_cycle(self) -> bool {
15710 self.rank_and_flags & 2 != 0
15711 }
15712}
15713
15714struct CompactPathInfoTable {
15715 slots: Vec<CompactPathInfoSlot>,
15716 mask: usize,
15717 key_mask: u64,
15718 generation: AtomicU8,
15719}
15720
15721impl CompactPathInfoTable {
15722 fn with_max_entries(max_entries: usize, k: usize) -> Self {
15723 debug_assert!(k <= 31);
15724 let capacity = max_entries
15725 .saturating_mul(2)
15729 .next_power_of_two()
15730 .max(8);
15731 let key_mask = (1u64 << (2 * k)) - 1;
15732 Self {
15733 slots: (0..capacity)
15734 .map(|_| CompactPathInfoSlot {
15735 key: AtomicU64::new(0),
15736 path_id: AtomicU64::new(0),
15737 rank_and_flags: AtomicU64::new(0),
15738 epoch: AtomicU8::new(0),
15739 })
15740 .collect(),
15741 mask: capacity - 1,
15742 key_mask,
15743 generation: AtomicU8::new(0),
15744 }
15745 }
15746
15747 fn clear(&self) {
15748 let current = self.generation.load(Ordering::Relaxed);
15749 let mut next = current.wrapping_add(1);
15750 if next == 0 || next == u8::MAX {
15751 self.slots
15752 .par_iter()
15753 .for_each(|slot| slot.epoch.store(0, Ordering::Relaxed));
15754 next = 1;
15755 }
15756 self.generation.store(next, Ordering::Relaxed);
15757 }
15758
15759 fn insert<const K: usize>(&self, vertex: Kmer<K>, value: PathInfo<K>) -> bool {
15760 let flags = u8::from(value.exit_side == Side::Back) | (u8::from(value.is_cycle) << 1);
15761 self.insert_raw(
15762 vertex.as_u128() as u64,
15763 value.path_id.as_u128() as u64,
15764 value.rank,
15765 flags,
15766 )
15767 }
15768
15769 #[inline(always)]
15770 fn insert_compact<const K: usize>(
15771 &self,
15772 vertex: Kmer<K>,
15773 value: CompactExpansionPathInfo,
15774 ) -> bool {
15775 self.insert_raw(
15776 vertex.as_u128() as u64,
15777 value.path_id,
15778 value.rank(),
15779 value.rank_and_flags as u8 & 3,
15780 )
15781 }
15782
15783 #[inline(always)]
15784 fn insert_raw(&self, key: u64, path_id: u64, rank: u64, flags: u8) -> bool {
15785 debug_assert!(rank <= (u64::MAX >> 2));
15786 debug_assert!(flags < 4);
15787 self.insert_packed(key, path_id, (rank << 2) | u64::from(flags))
15788 }
15789
15790 #[inline(always)]
15791 fn insert_packed(&self, key: u64, path_id: u64, rank_and_flags: u64) -> bool {
15792 const BUSY: u8 = u8::MAX;
15793 debug_assert_eq!(key & !self.key_mask, 0);
15794 let generation = self.generation.load(Ordering::Relaxed);
15795 let mut idx = wyhash_u64(key, 0) as usize & self.mask;
15796 loop {
15797 let slot = &self.slots[idx];
15798 let observed_epoch = slot.epoch.load(Ordering::Acquire);
15799 if observed_epoch == generation {
15800 if slot.key.load(Ordering::Relaxed) == key {
15801 return false;
15802 }
15803 idx = (idx + 1) & self.mask;
15804 continue;
15805 }
15806 if observed_epoch == BUSY {
15807 std::hint::spin_loop();
15808 continue;
15809 }
15810 if slot
15811 .epoch
15812 .compare_exchange(observed_epoch, BUSY, Ordering::AcqRel, Ordering::Acquire)
15813 .is_ok()
15814 {
15815 slot.key.store(key, Ordering::Relaxed);
15818 slot.path_id.store(path_id, Ordering::Relaxed);
15819 slot.rank_and_flags.store(rank_and_flags, Ordering::Relaxed);
15820 slot.epoch.store(generation, Ordering::Release);
15821 return true;
15822 }
15823 }
15824 }
15825
15826 #[inline(always)]
15827 fn get<const K: usize>(&self, vertex: Kmer<K>) -> Option<PathInfo<K>> {
15828 let value = self.get_compact(vertex)?;
15829 Some(PathInfo {
15830 path_id: Kmer::from_bits(value.path_id as u128),
15831 rank: value.rank(),
15832 exit_side: value.exit_side(),
15833 is_cycle: value.is_cycle(),
15834 })
15835 }
15836
15837 #[inline(always)]
15838 fn get_compact<const K: usize>(&self, vertex: Kmer<K>) -> Option<CompactExpansionPathInfo> {
15839 self.get_compact_raw(vertex.as_u128() as u64)
15840 }
15841
15842 #[inline(always)]
15843 fn get_compact_raw(&self, key: u64) -> Option<CompactExpansionPathInfo> {
15844 let generation = self.generation.load(Ordering::Relaxed);
15845 let mut idx = wyhash_u64(key, 0) as usize & self.mask;
15846 loop {
15847 let slot = &self.slots[idx];
15848 if slot.epoch.load(Ordering::Acquire) != generation {
15849 return None;
15850 }
15851 if slot.key.load(Ordering::Relaxed) == key {
15852 let path_id = slot.path_id.load(Ordering::Relaxed);
15853 let rank_and_flags = slot.rank_and_flags.load(Ordering::Relaxed);
15854 return Some(CompactExpansionPathInfo {
15855 path_id,
15856 rank_and_flags,
15857 });
15858 }
15859 idx = (idx + 1) & self.mask;
15860 }
15861 }
15862}
15863
15864enum ExpansionPathInfoTable<const K: usize> {
15865 Compact(CompactPathInfoTable),
15866 Wide(ConcurrentPathInfoTable<K>),
15867}
15868
15869impl<const K: usize> ExpansionPathInfoTable<K> {
15870 fn with_max_entries(max_entries: usize) -> Self {
15871 if K <= 31 {
15872 Self::Compact(CompactPathInfoTable::with_max_entries(max_entries, K))
15873 } else {
15874 Self::Wide(ConcurrentPathInfoTable::with_max_entries(max_entries))
15875 }
15876 }
15877
15878 fn clear(&self) {
15879 match self {
15880 Self::Compact(table) => table.clear(),
15881 Self::Wide(table) => table.clear(),
15882 }
15883 }
15884
15885 #[inline(always)]
15886 fn insert(&self, vertex: Kmer<K>, value: PathInfo<K>) -> bool {
15887 match self {
15888 Self::Compact(table) => table.insert(vertex, value),
15889 Self::Wide(table) => table.insert(vertex, value),
15890 }
15891 }
15892
15893 #[inline(always)]
15894 fn get(&self, vertex: Kmer<K>) -> Option<PathInfo<K>> {
15895 match self {
15896 Self::Compact(table) => table.get(vertex),
15897 Self::Wide(table) => table.get(vertex),
15898 }
15899 }
15900
15901 #[inline(always)]
15902 fn get_compact(&self, vertex: Kmer<K>) -> Option<CompactExpansionPathInfo> {
15903 match self {
15904 Self::Compact(table) => table.get_compact(vertex),
15905 Self::Wide(_) => unreachable!("compact path info is only used for k <= 31"),
15906 }
15907 }
15908
15909 #[inline(always)]
15910 fn insert_compact(&self, vertex: Kmer<K>, value: CompactExpansionPathInfo) -> bool {
15911 match self {
15912 Self::Compact(table) => table.insert_compact(vertex, value),
15913 Self::Wide(_) => unreachable!("compact path info is only used for k <= 31"),
15914 }
15915 }
15916
15917 #[inline(always)]
15918 fn insert_compact_record(&self, record: CompactVertexPathInfoRecord) -> bool {
15919 match self {
15920 Self::Compact(table) => {
15921 table.insert_packed(record.vertex, record.path_id, record.rank_and_flags)
15922 }
15923 Self::Wide(_) => unreachable!("compact path info is only used for k <= 31"),
15924 }
15925 }
15926
15927 #[allow(dead_code)]
15929 fn contains_key(&self, vertex: &Kmer<K>) -> bool {
15930 self.get(*vertex).is_some()
15931 }
15932
15933 fn capacity(&self) -> usize {
15934 match self {
15935 Self::Compact(table) => table.slots.len(),
15936 Self::Wide(table) => table.slots.len(),
15937 }
15938 }
15939
15940 fn slot_size(&self) -> usize {
15941 match self {
15942 Self::Compact(_) => std::mem::size_of::<CompactPathInfoSlot>(),
15943 Self::Wide(_) => std::mem::size_of::<PathInfoSlot<K>>(),
15944 }
15945 }
15946
15947 #[inline(always)]
15948 fn insert_encoded(&self, bytes: &[u8]) -> bool {
15949 match self {
15950 Self::Compact(table) => {
15951 if discontinuity_edge_kmer_bytes::<K>() != 8 {
15952 let record = decoded_vertex_path_info_record::<K>(bytes);
15953 return table.insert(record.vertex, record.info);
15954 }
15955 let mut key = [0u8; 8];
15956 key.copy_from_slice(&bytes[..8]);
15957 let mut path_id = [0u8; 8];
15958 path_id.copy_from_slice(&bytes[8..16]);
15959 let mut rank = [0u8; 8];
15960 rank.copy_from_slice(&bytes[16..24]);
15961 table.insert_packed(
15962 u64::from_le_bytes(key),
15963 u64::from_le_bytes(path_id),
15964 u64::from_le_bytes(rank),
15965 )
15966 }
15967 Self::Wide(table) => {
15968 let record = decoded_vertex_path_info_record::<K>(bytes);
15969 table.insert(record.vertex, record.info)
15970 }
15971 }
15972 }
15973}
15974
15975impl<const K: usize> ConcurrentPathInfoTable<K> {
15976 const EMPTY: u64 = 0;
15977 const BUSY: u64 = 1;
15978
15979 fn with_max_entries(max_entries: usize) -> Self {
15980 let capacity = max_entries.saturating_mul(2).next_power_of_two().max(8);
15981 Self {
15982 slots: (0..capacity)
15983 .map(|_| PathInfoSlot {
15984 tag: AtomicU64::new(Self::EMPTY),
15985 key: UnsafeCell::new(MaybeUninit::uninit()),
15986 value: UnsafeCell::new(MaybeUninit::uninit()),
15987 })
15988 .collect(),
15989 mask: capacity - 1,
15990 }
15991 }
15992
15993 #[inline]
15994 fn tag(vertex: Kmer<K>) -> u64 {
15995 vertex.hash64(0).wrapping_add(2).max(2)
15996 }
15997
15998 fn clear(&self) {
15999 self.slots
16000 .par_iter()
16001 .for_each(|slot| slot.tag.store(Self::EMPTY, Ordering::Relaxed));
16002 }
16003
16004 fn insert(&self, vertex: Kmer<K>, value: PathInfo<K>) -> bool {
16005 let tag = Self::tag(vertex);
16006 let mut idx = vertex.hash64(0) as usize & self.mask;
16007 loop {
16008 let slot = &self.slots[idx];
16009 let observed = slot.tag.load(Ordering::Acquire);
16010 if observed == Self::EMPTY
16011 && slot
16012 .tag
16013 .compare_exchange(
16014 Self::EMPTY,
16015 Self::BUSY,
16016 Ordering::Acquire,
16017 Ordering::Relaxed,
16018 )
16019 .is_ok()
16020 {
16021 unsafe {
16022 (*slot.key.get()).write(vertex);
16023 (*slot.value.get()).write(value);
16024 }
16025 slot.tag.store(tag, Ordering::Release);
16026 return true;
16027 }
16028 if observed == Self::BUSY {
16029 std::hint::spin_loop();
16030 continue;
16031 }
16032 if observed == tag && unsafe { (*slot.key.get()).assume_init() == vertex } {
16033 return false;
16034 }
16035 idx = (idx + 1) & self.mask;
16036 }
16037 }
16038
16039 fn get(&self, vertex: Kmer<K>) -> Option<PathInfo<K>> {
16040 let tag = Self::tag(vertex);
16041 let mut idx = vertex.hash64(0) as usize & self.mask;
16042 loop {
16043 let slot = &self.slots[idx];
16044 let observed = slot.tag.load(Ordering::Acquire);
16045 if observed == Self::EMPTY {
16046 return None;
16047 }
16048 if observed == Self::BUSY {
16049 std::hint::spin_loop();
16050 continue;
16051 }
16052 if observed == tag && unsafe { (*slot.key.get()).assume_init() == vertex } {
16053 return Some(unsafe { (*slot.value.get()).assume_init() });
16054 }
16055 idx = (idx + 1) & self.mask;
16056 }
16057 }
16058
16059 #[inline]
16060 #[allow(dead_code)]
16062 fn contains_key(&self, vertex: &Kmer<K>) -> bool {
16063 self.get(*vertex).is_some()
16064 }
16065}
16066
16067#[allow(dead_code)]
16069struct OwnedPartitionTables<const K: usize> {
16070 maps: Vec<FastHashMap<Kmer<K>, PartitionOtherEnd<K>>>,
16071 mask: usize,
16072}
16073
16074impl<const K: usize> OwnedPartitionTables<K> {
16076 #[allow(dead_code)]
16077 fn new(threads: usize) -> Self {
16078 let count = threads.max(1).next_power_of_two();
16079 Self {
16080 maps: (0..count).map(|_| FastHashMap::default()).collect(),
16081 mask: count - 1,
16082 }
16083 }
16084
16085 #[inline]
16086 #[allow(dead_code)]
16087 fn owner(&self, vertex: Kmer<K>) -> usize {
16088 partition_column_vertex_shard(vertex, self.mask)
16089 }
16090
16091 fn clear(&mut self) {
16092 for map in &mut self.maps {
16093 map.clear();
16094 }
16095 }
16096
16097 fn contains_key(&self, vertex: &Kmer<K>) -> bool {
16098 self.maps[self.owner(*vertex)].contains_key(vertex)
16099 }
16100
16101 fn get(&self, vertex: &Kmer<K>) -> Option<&PartitionOtherEnd<K>> {
16102 self.maps[self.owner(*vertex)].get(vertex)
16103 }
16104
16105 fn get_mut(&mut self, vertex: &Kmer<K>) -> Option<&mut PartitionOtherEnd<K>> {
16106 let owner = self.owner(*vertex);
16107 self.maps[owner].get_mut(vertex)
16108 }
16109
16110 fn insert(&mut self, vertex: Kmer<K>, end: PartitionOtherEnd<K>) {
16111 let owner = self.owner(vertex);
16112 self.maps[owner].insert(vertex, end);
16113 }
16114}
16115
16116#[derive(Clone, Copy)]
16117struct CompactPartitionOtherEnd {
16118 endpoint: u64,
16119 weight: u16,
16120 flags: u8,
16121}
16122
16123impl CompactPartitionOtherEnd {
16124 const PHI: u8 = 1 << 4;
16125
16126 fn pack<const K: usize>(end: PartitionOtherEnd<K>) -> Self {
16127 let (endpoint, endpoint_side, phi) = match end.endpoint {
16128 MatrixEndpoint::Phi => (0, Side::Front, true),
16129 MatrixEndpoint::Vertex(endpoint) => {
16130 (endpoint.vertex.as_u128() as u64, endpoint.side, false)
16131 }
16132 };
16133 let mut flags = u8::from(endpoint_side == Side::Back)
16134 | (u8::from(end.side_at_current == Side::Back) << 1)
16135 | (u8::from(end.in_same_part) << 2)
16136 | (u8::from(end.processed) << 3);
16137 if phi {
16138 flags |= Self::PHI;
16139 }
16140 Self {
16141 endpoint,
16142 weight: u16::try_from(end.weight).expect("partition edge weight fits u16"),
16143 flags,
16144 }
16145 }
16146
16147 fn unpack<const K: usize>(self) -> PartitionOtherEnd<K> {
16148 let side = |bit| {
16149 if self.flags & bit == 0 {
16150 Side::Front
16151 } else {
16152 Side::Back
16153 }
16154 };
16155 PartitionOtherEnd {
16156 endpoint: if self.flags & Self::PHI != 0 {
16157 MatrixEndpoint::Phi
16158 } else {
16159 MatrixEndpoint::Vertex(DiscontinuityEndpoint {
16160 vertex: Kmer::from_bits(self.endpoint as u128),
16161 side: side(1),
16162 })
16163 },
16164 side_at_current: side(1 << 1),
16165 weight: u64::from(self.weight),
16166 in_same_part: self.flags & (1 << 2) != 0,
16167 processed: self.flags & (1 << 3) != 0,
16168 }
16169 }
16170}
16171
16172struct AtomicPartitionSlot {
16173 key: AtomicU64,
16174 value: UnsafeCell<MaybeUninit<CompactPartitionOtherEnd>>,
16175}
16176
16177#[allow(dead_code)]
16179struct FlatPartitionTable<const K: usize> {
16180 keys: Vec<u64>,
16181 values: Vec<MaybeUninit<PartitionOtherEnd<K>>>,
16182 occupied: Vec<usize>,
16183 mask: usize,
16184}
16185
16186impl<const K: usize> FlatPartitionTable<K> {
16188 #[allow(dead_code)]
16189 const EMPTY: u64 = u64::MAX;
16190
16191 #[allow(dead_code)]
16192 fn with_max_entries(max_entries: usize) -> Self {
16193 let capacity = max_entries
16194 .saturating_mul(4)
16195 .div_ceil(3)
16196 .next_power_of_two()
16197 .max(8);
16198 Self {
16199 keys: vec![Self::EMPTY; capacity],
16200 values: (0..capacity).map(|_| MaybeUninit::uninit()).collect(),
16201 occupied: Vec::with_capacity(max_entries),
16202 mask: capacity - 1,
16203 }
16204 }
16205
16206 #[allow(dead_code)]
16207 fn clear(&mut self) {
16208 for idx in self.occupied.drain(..) {
16209 self.keys[idx] = Self::EMPTY;
16210 }
16211 }
16212
16213 #[inline]
16214 #[allow(dead_code)]
16215 fn find_index(&self, vertex: Kmer<K>) -> Result<usize, usize> {
16216 let key = vertex.as_u128() as u64;
16217 let mut idx = vertex.hash64(0) as usize & self.mask;
16218 loop {
16219 let observed = self.keys[idx];
16220 if observed == key {
16221 return Ok(idx);
16222 }
16223 if observed == Self::EMPTY {
16224 return Err(idx);
16225 }
16226 idx = (idx + 1) & self.mask;
16227 }
16228 }
16229
16230 fn absorb(
16231 &mut self,
16232 vertex: Kmer<K>,
16233 incoming: PartitionOtherEnd<K>,
16234 partition: usize,
16235 edges: &mut Vec<DiscontinuityEdge<K>>,
16236 meta_vertices: &mut Vec<SerialMetaVertex<K>>,
16237 ) {
16238 match self.find_index(vertex) {
16239 Ok(idx) => {
16240 let existing = unsafe { self.values[idx].assume_init() };
16241 absorb_partition_collision_outputs(
16242 vertex,
16243 incoming,
16244 existing,
16245 partition,
16246 edges,
16247 meta_vertices,
16248 );
16249 let mut updated = existing;
16250 if updated.in_same_part {
16251 updated = incoming;
16252 }
16253 updated.processed = true;
16254 self.values[idx].write(updated);
16255 }
16256 Err(idx) => {
16257 self.keys[idx] = vertex.as_u128() as u64;
16258 self.values[idx].write(incoming);
16259 self.occupied.push(idx);
16260 }
16261 }
16262 }
16263
16264 fn get(&self, vertex: Kmer<K>) -> Option<PartitionOtherEnd<K>> {
16265 self.find_index(vertex)
16266 .ok()
16267 .map(|idx| unsafe { self.values[idx].assume_init() })
16268 }
16269
16270 fn insert(&mut self, vertex: Kmer<K>, value: PartitionOtherEnd<K>) {
16271 match self.find_index(vertex) {
16272 Ok(idx) => {
16273 self.values[idx].write(value);
16274 }
16275 Err(idx) => {
16276 self.keys[idx] = vertex.as_u128() as u64;
16277 self.values[idx].write(value);
16278 self.occupied.push(idx);
16279 }
16280 }
16281 }
16282
16283 fn mark_processed(&mut self, vertex: Kmer<K>) {
16284 let idx = self.find_index(vertex).expect("partition endpoint exists");
16285 let mut value = unsafe { self.values[idx].assume_init() };
16286 value.processed = true;
16287 self.values[idx].write(value);
16288 }
16289}
16290
16291unsafe impl Sync for AtomicPartitionSlot {}
16292
16293struct AtomicPartitionTable<const K: usize> {
16294 slots: Vec<AtomicPartitionSlot>,
16295 mask: usize,
16296}
16297
16298impl<const K: usize> AtomicPartitionTable<K> {
16299 const EMPTY: u64 = u64::MAX;
16300 const LOCKED: u64 = 1 << 63;
16301 const HASH_SEED: u64 = 0xAAAA_AAAA_5555_5555;
16302
16303 #[inline(always)]
16304 fn index(&self, vertex: Kmer<K>) -> usize {
16305 vertex.hash64(Self::HASH_SEED) as usize & self.mask
16306 }
16307
16308 #[allow(dead_code)]
16310 fn with_max_entries(max_entries: usize) -> Self {
16311 let capacity = max_entries
16312 .saturating_mul(4)
16313 .div_ceil(3)
16314 .next_power_of_two()
16315 .max(8);
16316 Self {
16317 slots: (0..capacity)
16318 .map(|_| AtomicPartitionSlot {
16319 key: AtomicU64::new(Self::EMPTY),
16320 value: UnsafeCell::new(MaybeUninit::uninit()),
16321 })
16322 .collect(),
16323 mask: capacity - 1,
16324 }
16325 }
16326
16327 fn clear(&self, threads: usize, pool: &ThreadPool) {
16328 let workers = threads.max(1).min(self.slots.len());
16329 let chunk = self.slots.len().div_ceil(workers);
16330 pool.install(|| {
16331 self.slots.par_chunks(chunk).for_each(|slots| {
16332 for slot in slots {
16333 slot.key.store(Self::EMPTY, Ordering::Relaxed);
16334 }
16335 })
16336 });
16337 }
16338
16339 #[allow(dead_code)]
16341 fn absorb(
16342 &self,
16343 vertex: Kmer<K>,
16344 incoming: PartitionOtherEnd<K>,
16345 partition: usize,
16346 edges: &mut Vec<DiscontinuityEdge<K>>,
16347 meta_vertices: &mut Vec<SerialMetaVertex<K>>,
16348 ) {
16349 let key = vertex.as_u128() as u64;
16350 debug_assert!(K <= 31 && key != Self::EMPTY);
16351 let mut idx = self.index(vertex);
16352 loop {
16353 let slot = &self.slots[idx];
16354 let observed = slot.key.load(Ordering::Acquire);
16355 if observed == key {
16356 if slot
16357 .key
16358 .compare_exchange_weak(
16359 key,
16360 key | Self::LOCKED,
16361 Ordering::Acquire,
16362 Ordering::Relaxed,
16363 )
16364 .is_err()
16365 {
16366 std::hint::spin_loop();
16367 continue;
16368 }
16369 let existing = unsafe { (*slot.value.get()).assume_init() }.unpack::<K>();
16370 absorb_partition_collision_outputs(
16371 vertex,
16372 incoming,
16373 existing,
16374 partition,
16375 edges,
16376 meta_vertices,
16377 );
16378 let mut updated = existing;
16379 if updated.in_same_part {
16380 updated = incoming;
16381 }
16382 updated.processed = true;
16383 unsafe { (*slot.value.get()).write(CompactPartitionOtherEnd::pack(updated)) };
16384 slot.key.store(key, Ordering::Release);
16385 return;
16386 }
16387 if observed == Self::EMPTY {
16388 if slot
16389 .key
16390 .compare_exchange_weak(
16391 Self::EMPTY,
16392 Self::LOCKED,
16393 Ordering::Acquire,
16394 Ordering::Relaxed,
16395 )
16396 .is_err()
16397 {
16398 std::hint::spin_loop();
16399 continue;
16400 }
16401 unsafe { (*slot.value.get()).write(CompactPartitionOtherEnd::pack(incoming)) };
16402 slot.key.store(key, Ordering::Release);
16403 return;
16404 }
16405 if observed & Self::LOCKED != 0 {
16406 std::hint::spin_loop();
16407 continue;
16408 }
16409 idx = (idx + 1) & self.mask;
16410 }
16411 }
16412
16413 fn absorb_prepared(
16414 &self,
16415 vertex: Kmer<K>,
16416 incoming: PartitionOtherEnd<K>,
16417 partition: usize,
16418 vertex_partitions: usize,
16419 edges: &mut Vec<PreparedBlockedEdge>,
16420 meta_vertices: &mut Vec<SerialMetaVertex<K>>,
16421 ) {
16422 let key = vertex.as_u128() as u64;
16423 debug_assert!(K <= 31 && key != Self::EMPTY);
16424 let mut idx = self.index(vertex);
16425 loop {
16426 let slot = &self.slots[idx];
16427 let observed = slot.key.load(Ordering::Acquire);
16428 if observed == key {
16429 if slot
16430 .key
16431 .compare_exchange_weak(
16432 key,
16433 key | Self::LOCKED,
16434 Ordering::Acquire,
16435 Ordering::Relaxed,
16436 )
16437 .is_err()
16438 {
16439 std::hint::spin_loop();
16440 continue;
16441 }
16442 let existing = unsafe { (*slot.value.get()).assume_init() }.unpack::<K>();
16443 if existing.endpoint.is_phi() && incoming.endpoint.is_phi() {
16444 meta_vertices.push(two_weight_meta_vertex(
16445 vertex,
16446 partition,
16447 incoming.side_at_current,
16448 incoming.weight,
16449 existing.weight,
16450 false,
16451 ));
16452 } else if !existing.in_same_part {
16453 edges.push(prepare_existing_blocked_edge(
16454 join_other_ends(
16455 incoming.endpoint,
16456 existing.endpoint,
16457 incoming.weight + existing.weight,
16458 ),
16459 vertex_partitions,
16460 ));
16461 }
16462 let mut updated = existing;
16463 if updated.in_same_part {
16464 updated = incoming;
16465 }
16466 updated.processed = true;
16467 unsafe { (*slot.value.get()).write(CompactPartitionOtherEnd::pack(updated)) };
16468 slot.key.store(key, Ordering::Release);
16469 return;
16470 }
16471 if observed == Self::EMPTY {
16472 if slot
16473 .key
16474 .compare_exchange_weak(
16475 Self::EMPTY,
16476 Self::LOCKED,
16477 Ordering::Acquire,
16478 Ordering::Relaxed,
16479 )
16480 .is_err()
16481 {
16482 std::hint::spin_loop();
16483 continue;
16484 }
16485 unsafe { (*slot.value.get()).write(CompactPartitionOtherEnd::pack(incoming)) };
16486 slot.key.store(key, Ordering::Release);
16487 return;
16488 }
16489 if observed & Self::LOCKED != 0 {
16490 std::hint::spin_loop();
16491 continue;
16492 }
16493 idx = (idx + 1) & self.mask;
16494 }
16495 }
16496
16497 #[allow(dead_code)]
16499 fn to_fast_map(&self) -> FastHashMap<Kmer<K>, PartitionOtherEnd<K>> {
16500 let occupied = self
16501 .slots
16502 .iter()
16503 .filter(|slot| slot.key.load(Ordering::Relaxed) != Self::EMPTY)
16504 .count();
16505 let mut map = FastHashMap::with_capacity_and_hasher(occupied, FastBuildHasher::default());
16506 for slot in &self.slots {
16507 let key = slot.key.load(Ordering::Relaxed);
16508 if key != Self::EMPTY {
16509 let value = unsafe { (*slot.value.get()).assume_init() }.unpack::<K>();
16510 map.insert(Kmer::from_bits(key as u128), value);
16511 }
16512 }
16513 map
16514 }
16515
16516 fn get(&self, vertex: Kmer<K>) -> Option<PartitionOtherEnd<K>> {
16517 let key = vertex.as_u128() as u64;
16518 let mut idx = self.index(vertex);
16519 loop {
16520 let slot = &self.slots[idx];
16521 let observed = slot.key.load(Ordering::Acquire);
16522 if observed == key {
16523 return Some(unsafe { (*slot.value.get()).assume_init() }.unpack::<K>());
16524 }
16525 if observed == Self::EMPTY {
16526 return None;
16527 }
16528 idx = (idx + 1) & self.mask;
16529 }
16530 }
16531
16532 #[allow(dead_code)]
16534 fn insert_serial(&self, vertex: Kmer<K>, value: PartitionOtherEnd<K>) {
16535 let key = vertex.as_u128() as u64;
16536 let mut idx = self.index(vertex);
16537 loop {
16538 let slot = &self.slots[idx];
16539 let observed = slot.key.load(Ordering::Relaxed);
16540 if observed == key {
16541 unsafe { (*slot.value.get()).write(CompactPartitionOtherEnd::pack(value)) };
16542 return;
16543 }
16544 if observed == Self::EMPTY {
16545 unsafe { (*slot.value.get()).write(CompactPartitionOtherEnd::pack(value)) };
16546 slot.key.store(key, Ordering::Release);
16547 return;
16548 }
16549 idx = (idx + 1) & self.mask;
16550 }
16551 }
16552
16553 fn mark_processed(&self, vertex: Kmer<K>) {
16554 let key = vertex.as_u128() as u64;
16555 let mut idx = self.index(vertex);
16556 loop {
16557 let slot = &self.slots[idx];
16558 let observed = slot.key.load(Ordering::Relaxed);
16559 if observed == key {
16560 let mut value = unsafe { (*slot.value.get()).assume_init() }.unpack::<K>();
16561 value.processed = true;
16562 unsafe { (*slot.value.get()).write(CompactPartitionOtherEnd::pack(value)) };
16563 return;
16564 }
16565 assert_ne!(observed, Self::EMPTY, "partition endpoint disappeared");
16566 idx = (idx + 1) & self.mask;
16567 }
16568 }
16569}
16570
16571fn absorb_partition_other_end<const K: usize>(
16572 vertex: Kmer<K>,
16573 incoming: PartitionOtherEnd<K>,
16574 partition: usize,
16575 table: &mut FastHashMap<Kmer<K>, PartitionOtherEnd<K>>,
16576 edges: &mut Vec<DiscontinuityEdge<K>>,
16577 meta_vertices: &mut Vec<SerialMetaVertex<K>>,
16578) {
16579 match table.get_mut(&vertex) {
16580 None => {
16581 table.insert(vertex, incoming);
16582 }
16583 Some(existing) => {
16584 if existing.endpoint.is_phi() && incoming.endpoint.is_phi() {
16585 meta_vertices.push(two_weight_meta_vertex(
16586 vertex,
16587 partition,
16588 incoming.side_at_current,
16589 incoming.weight,
16590 existing.weight,
16591 false,
16592 ));
16593 } else if existing.in_same_part {
16594 *existing = incoming;
16595 } else {
16596 edges.push(join_other_ends(
16597 incoming.endpoint,
16598 existing.endpoint,
16599 incoming.weight + existing.weight,
16600 ));
16601 }
16602 existing.processed = true;
16603 }
16604 }
16605}
16606
16607fn absorb_partition_collision_outputs<const K: usize>(
16608 vertex: Kmer<K>,
16609 incoming: PartitionOtherEnd<K>,
16610 existing: PartitionOtherEnd<K>,
16611 partition: usize,
16612 edges: &mut Vec<DiscontinuityEdge<K>>,
16613 meta_vertices: &mut Vec<SerialMetaVertex<K>>,
16614) {
16615 if existing.endpoint.is_phi() && incoming.endpoint.is_phi() {
16616 meta_vertices.push(two_weight_meta_vertex(
16617 vertex,
16618 partition,
16619 incoming.side_at_current,
16620 incoming.weight,
16621 existing.weight,
16622 false,
16623 ));
16624 } else if !existing.in_same_part {
16625 edges.push(join_other_ends(
16626 incoming.endpoint,
16627 existing.endpoint,
16628 incoming.weight + existing.weight,
16629 ));
16630 }
16631}
16632
16633fn endpoint_in_partition<const K: usize>(
16634 matrix: &SerialEdgeMatrix<K>,
16635 edge: &DiscontinuityEdge<K>,
16636 partition: usize,
16637) -> Option<(MatrixEndpoint<K>, DiscontinuityEndpoint<K>)> {
16638 match (edge.first, edge.second) {
16639 (lower, MatrixEndpoint::Vertex(current))
16640 if matrix.partition(MatrixEndpoint::Vertex(current)) == partition =>
16641 {
16642 Some((lower, current))
16643 }
16644 (MatrixEndpoint::Vertex(current), lower)
16645 if matrix.partition(MatrixEndpoint::Vertex(current)) == partition =>
16646 {
16647 Some((lower, current))
16648 }
16649 _ => None,
16650 }
16651}
16652
16653fn join_other_ends<const K: usize>(
16654 first: MatrixEndpoint<K>,
16655 second: MatrixEndpoint<K>,
16656 weight: u64,
16657) -> DiscontinuityEdge<K> {
16658 join_other_ends_with_phantom(first, second, weight, None)
16659}
16660
16661fn join_other_ends_with_phantom<const K: usize>(
16662 first: MatrixEndpoint<K>,
16663 second: MatrixEndpoint<K>,
16664 weight: u64,
16665 phantom_unitig: Option<DiscontinuityEndpoint<K>>,
16666) -> DiscontinuityEdge<K> {
16667 DiscontinuityEdge {
16668 first,
16669 second,
16670 weight,
16671 unitig_bucket: 0,
16672 unitig_index: 0,
16673 unitig_exit_side: Side::Back,
16674 phantom_unitig,
16675 swapped: false,
16676 }
16677}
16678
16679fn two_weight_meta_vertex<const K: usize>(
16680 vertex: Kmer<K>,
16681 partition: usize,
16682 side: Side,
16683 front_weight: u64,
16684 back_weight: u64,
16685 is_cycle: bool,
16686) -> SerialMetaVertex<K> {
16687 SerialMetaVertex {
16688 vertex,
16689 partition,
16690 entry_side: Side::Back,
16691 weight: if side == Side::Front {
16692 front_weight
16693 } else {
16694 back_weight
16695 },
16696 is_cycle,
16697 }
16698}
16699
16700fn max_endpoint_partition<const K: usize>(
16701 matrix: &SerialEdgeMatrix<K>,
16702 first: MatrixEndpoint<K>,
16703 second: MatrixEndpoint<K>,
16704) -> usize {
16705 matrix.partition(first).max(matrix.partition(second))
16706}
16707
16708fn edge_matrix_partition<const K: usize>(
16709 vertex_partitions: usize,
16710 endpoint: MatrixEndpoint<K>,
16711) -> usize {
16712 match endpoint {
16713 MatrixEndpoint::Phi => 0,
16714 MatrixEndpoint::Vertex(endpoint) => {
16715 ((endpoint.vertex.hash64(0) as usize) & (vertex_partitions - 1)) + 1
16716 }
16717 }
16718}
16719
16720fn edge_matrix_row_col<const K: usize>(
16721 vertex_partitions: usize,
16722 first: MatrixEndpoint<K>,
16723 second: MatrixEndpoint<K>,
16724) -> (usize, usize) {
16725 let first_partition = edge_matrix_partition(vertex_partitions, first);
16726 let second_partition = edge_matrix_partition(vertex_partitions, second);
16727 if first_partition <= second_partition {
16728 (first_partition, second_partition)
16729 } else {
16730 (second_partition, first_partition)
16731 }
16732}
16733
16734fn endpoint_sort_key<const K: usize>(endpoint: MatrixEndpoint<K>) -> (u8, u128, u8) {
16735 match endpoint {
16736 MatrixEndpoint::Phi => (0, 0, 0),
16737 MatrixEndpoint::Vertex(endpoint) => (1, endpoint.vertex.as_u128(), endpoint.side as u8),
16738 }
16739}
16740
16741fn endpoint_id<const K: usize>(
16742 ids: &mut FastHashMap<EndpointKey<K>, usize>,
16743 adjacency: &mut Vec<Vec<(usize, u64, bool)>>,
16744 endpoint: MatrixEndpoint<K>,
16745) -> usize {
16746 let next = ids.len();
16747 *ids.entry(EndpointKey(endpoint)).or_insert_with(|| {
16748 adjacency.push(Vec::new());
16749 next
16750 })
16751}
16752
16753#[derive(Debug, Clone, Copy, PartialEq, Eq)]
16754pub enum SerialEdgeMatrixError {
16755 InvalidPartitionCount(usize),
16756}
16757
16758impl std::fmt::Display for SerialEdgeMatrixError {
16759 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
16760 match self {
16761 Self::InvalidPartitionCount(count) => write!(
16762 f,
16763 "vertex partition count must be a non-zero power of two, got {count}"
16764 ),
16765 }
16766 }
16767}
16768
16769impl std::error::Error for SerialEdgeMatrixError {}
16770
16771pub fn emit_uncolored_discontinuity_inputs<const K: usize>(
16772 bucket_dir: impl AsRef<Path>,
16773 cutoff: u32,
16774) -> Result<DiscontinuityInputs<K>, DiscontinuityInputError> {
16775 emit_uncolored_discontinuity_inputs_with_threads::<K>(bucket_dir, cutoff, 1)
16776}
16777
16778pub fn emit_uncolored_discontinuity_inputs_with_threads<const K: usize>(
16779 bucket_dir: impl AsRef<Path>,
16780 cutoff: u32,
16781 threads: usize,
16782) -> Result<DiscontinuityInputs<K>, DiscontinuityInputError> {
16783 emit_uncolored_discontinuity_inputs_with_threads_impl::<K>(bucket_dir, cutoff, threads, None)
16784}
16785
16786pub fn emit_uncolored_discontinuity_inputs_with_threads_in_dir<const K: usize>(
16787 bucket_dir: impl AsRef<Path>,
16788 cutoff: u32,
16789 threads: usize,
16790 label_path: impl AsRef<Path>,
16791) -> Result<DiscontinuityInputs<K>, DiscontinuityInputError> {
16792 emit_uncolored_discontinuity_inputs_with_threads_impl::<K>(
16793 bucket_dir,
16794 cutoff,
16795 threads,
16796 Some(label_path.as_ref()),
16797 )
16798}
16799
16800pub fn emit_uncolored_external_discontinuity_inputs_with_threads_in_dir<const K: usize>(
16808 bucket_dir: impl AsRef<Path>,
16809 cutoff: u32,
16810 threads: usize,
16811 label_path: impl AsRef<Path>,
16812 direct_output_path: Option<&Path>,
16813) -> Result<ExternalDiscontinuityInputs<K>, DiscontinuityInputError> {
16814 if cutoff == 0 {
16815 return Err(DiscontinuityInputError::InvalidCutoff);
16816 }
16817 if threads == 0 {
16818 return Err(DiscontinuityInputError::InvalidThreadCount);
16819 }
16820 let (store, entries) = BucketStore::open_dir(bucket_dir.as_ref())?;
16821 contract_local_subgraphs_into_external_inputs::<K>(
16822 &store,
16823 &entries,
16824 cutoff,
16825 threads,
16826 label_path.as_ref(),
16827 None,
16828 direct_output_path,
16829 None,
16830 0,
16831 )
16832}
16833
16834pub fn emit_colored_external_discontinuity_inputs_with_threads_in_dir<const K: usize>(
16839 bucket_dir: impl AsRef<Path>,
16840 cutoff: u32,
16841 threads: usize,
16842 label_path: impl AsRef<Path>,
16843 color_path: impl AsRef<Path>,
16844 color_repository_dir: impl AsRef<Path>,
16845 num_colors: u32,
16846) -> Result<ExternalDiscontinuityInputs<K>, DiscontinuityInputError> {
16847 if cutoff == 0 {
16848 return Err(DiscontinuityInputError::InvalidCutoff);
16849 }
16850 if threads == 0 {
16851 return Err(DiscontinuityInputError::InvalidThreadCount);
16852 }
16853 let (store, entries) = BucketStore::open_dir(bucket_dir.as_ref())?;
16854 contract_local_subgraphs_into_external_inputs::<K>(
16855 &store,
16856 &entries,
16857 cutoff,
16858 threads,
16859 label_path.as_ref(),
16860 Some(color_path.as_ref()),
16861 None,
16863 Some(color_repository_dir.as_ref()),
16864 num_colors,
16865 )
16866}
16867
16868fn emit_uncolored_discontinuity_inputs_with_threads_impl<const K: usize>(
16869 bucket_dir: impl AsRef<Path>,
16870 cutoff: u32,
16871 threads: usize,
16872 label_path: Option<&Path>,
16873) -> Result<DiscontinuityInputs<K>, DiscontinuityInputError> {
16874 if cutoff == 0 {
16875 return Err(DiscontinuityInputError::InvalidCutoff);
16876 }
16877 if threads == 0 {
16878 return Err(DiscontinuityInputError::InvalidThreadCount);
16879 }
16880
16881 let (store, entries) = BucketStore::open_dir(bucket_dir.as_ref())?;
16882 if let Some(label_path) = label_path {
16883 let external = contract_local_subgraphs_into_external_inputs::<K>(
16884 &store, &entries, cutoff, threads, label_path, None, None, None, 0,
16885 )?;
16886 return external_inputs_to_memory_inputs(external);
16887 }
16888
16889 let outputs = contract_local_subgraphs::<K>(&store, &entries, cutoff, threads)?;
16890 let mut inputs = DiscontinuityInputs::empty(DiscontinuityInputStats {
16891 input_buckets: entries.len(),
16892 ..DiscontinuityInputStats::default()
16893 });
16894
16895 for output in outputs {
16896 inputs.stats.weak_superkmers += output.weak_superkmers;
16897 let label_offset = inputs.labels.len() as u64;
16898 inputs.labels.extend_from_slice(&output.labels);
16899 for mut unitig in output.unitigs {
16900 inputs.stats.local_unitigs += 1;
16901 inputs.stats.discontinuity_exits +=
16902 u64::from(unitig.left_exit().is_some()) + u64::from(unitig.right_exit().is_some());
16903 inputs.stats.unitig_bases += unitig.label_len as u64;
16904 unitig.label_start += label_offset;
16905 inputs.unitigs.push(unitig);
16906 }
16907 }
16908
16909 Ok(inputs)
16910}
16911
16912fn external_inputs_to_memory_inputs<const K: usize>(
16913 external: ExternalDiscontinuityInputs<K>,
16914) -> Result<DiscontinuityInputs<K>, DiscontinuityInputError> {
16915 let unitig_file =
16916 File::open(&external.unitig_path).map_err(|source| DiscontinuityInputError::Io {
16917 path: external.unitig_path.clone(),
16918 source,
16919 })?;
16920 let mut unitig_input = BufReader::with_capacity(1024 * 1024, unitig_file);
16921 let mut unitigs = Vec::with_capacity(external.unitigs);
16922 for _ in 0..external.unitigs {
16923 unitigs.push(read_discontinuity_unitig_from_reader_for_input(
16924 &mut unitig_input,
16925 &external.unitig_path,
16926 )?);
16927 }
16928
16929 let mut label_file =
16930 File::open(&external.label_path).map_err(|source| DiscontinuityInputError::Io {
16931 path: external.label_path.clone(),
16932 source,
16933 })?;
16934 let label_len = label_file
16935 .metadata()
16936 .map_err(|source| DiscontinuityInputError::Io {
16937 path: external.label_path.clone(),
16938 source,
16939 })?
16940 .len() as usize;
16941 let mut labels = vec![0u8; label_len];
16942 label_file
16943 .read_exact(&mut labels)
16944 .map_err(|source| DiscontinuityInputError::Io {
16945 path: external.label_path.clone(),
16946 source,
16947 })?;
16948
16949 Ok(DiscontinuityInputs {
16950 unitigs,
16951 labels,
16952 stats: external.stats,
16953 })
16954}
16955
16956#[allow(clippy::too_many_arguments)]
16957fn contract_local_subgraphs_into_external_inputs<const K: usize>(
16958 store: &BucketStore,
16959 entries: &[BucketManifestEntry],
16960 cutoff: u32,
16961 threads: usize,
16962 label_path: &Path,
16963 color_path: Option<&Path>,
16964 direct_output_path: Option<&Path>,
16965 color_repository_dir: Option<&Path>,
16966 num_colors: u32,
16967) -> Result<ExternalDiscontinuityInputs<K>, DiscontinuityInputError> {
16968 let compact_unitigs = std::env::var_os("CF3_RS_ENDPOINT_STITCH").is_none();
16973 let mut inputs = DiscontinuityInputs::empty(DiscontinuityInputStats {
16974 input_buckets: entries.len(),
16975 ..DiscontinuityInputStats::default()
16976 });
16977
16978 if let Some(parent) = label_path.parent() {
16979 fs::create_dir_all(parent).map_err(|source| DiscontinuityInputError::Io {
16980 path: parent.to_path_buf(),
16981 source,
16982 })?;
16983 }
16984 let file = File::create(label_path).map_err(|source| DiscontinuityInputError::Io {
16985 path: label_path.to_path_buf(),
16986 source,
16987 })?;
16988 let mut labels = BufWriter::with_capacity(8 * 1024 * 1024, file);
16989 let unitig_path = unitig_table_path_for_labels(label_path);
16990 let unitig_file = File::create(&unitig_path).map_err(|source| DiscontinuityInputError::Io {
16991 path: unitig_path.clone(),
16992 source,
16993 })?;
16994 let mut unitigs = BufWriter::with_capacity(8 * 1024 * 1024, unitig_file);
16995 let trivial_path = direct_output_path
16999 .map(Path::to_path_buf)
17000 .unwrap_or_else(|| label_path.with_extension("trivial.fa"));
17001 let trivial_file =
17002 File::create(&trivial_path).map_err(|source| DiscontinuityInputError::Io {
17003 path: trivial_path.clone(),
17004 source,
17005 })?;
17006 let mut trivial_output = BufWriter::with_capacity(8 * 1024 * 1024, trivial_file);
17007 let mut trivial_unitigs = 0u64;
17008 let mut trivial_bases = 0u64;
17009 let matrix_dir = unitig_path.with_file_name(format!(
17010 "{}.edge-matrix",
17011 unitig_path
17012 .file_name()
17013 .and_then(|name| name.to_str())
17014 .unwrap_or("unitigs")
17015 ));
17016 let mut edge_matrix = BlockedEdgeMatrix::create(&matrix_dir, DEFAULT_VERTEX_PARTITIONS)
17017 .map_err(serial_collation_to_input_error)?;
17018 let mut label_offset = 0u64;
17019 let mut ranges = Vec::new();
17020 let expected_color_ceiling = std::env::var("CF3_RS_EXPECTED_COLORS")
17026 .ok()
17027 .and_then(|value| value.parse::<u64>().ok())
17028 .filter(|value| *value > 0)
17029 .unwrap_or(DEFAULT_EXPECTED_COLOR_CEILING);
17030 let expected_colors = entries
17031 .iter()
17032 .map(|entry| entry.records)
17033 .sum::<u64>()
17034 .div_ceil(16)
17035 .min(expected_color_ceiling) as usize;
17036 let color_repository = color_path
17037 .map(|path| {
17038 let dir = color_repository_dir
17041 .map(Path::to_path_buf)
17042 .unwrap_or_else(|| path.with_extension("color-repository"));
17043 ConcurrentColorRepository::create(
17044 dir,
17045 threads.min(256),
17046 expected_colors.max(entries.len() * 8),
17047 num_colors,
17048 )
17049 })
17050 .transpose()?;
17051 let mut groups = local_bucket_groups(entries)?;
17052 groups.sort_by_key(|group| std::cmp::Reverse((group.stored_bytes, group.graph_id)));
17053 let workers = threads.min(groups.len().max(1));
17054 let local_unitig_bucket_count = if compact_unitigs && workers > 1 {
17058 local_unitig_bucket_plan(open_file_limit(), current_open_file_count(), workers)
17059 } else {
17060 0
17061 };
17062 let local_unitig_bucket_dir = unitig_path.with_extension("local-unitig-buckets");
17063 let mut local_unitig_writers = if local_unitig_bucket_count == 0 {
17064 None
17065 } else {
17066 if local_unitig_bucket_dir.exists() {
17067 fs::remove_dir_all(&local_unitig_bucket_dir).map_err(|source| {
17068 DiscontinuityInputError::Io {
17069 path: local_unitig_bucket_dir.clone(),
17070 source,
17071 }
17072 })?;
17073 }
17074 fs::create_dir_all(&local_unitig_bucket_dir).map_err(|source| {
17075 DiscontinuityInputError::Io {
17076 path: local_unitig_bucket_dir.clone(),
17077 source,
17078 }
17079 })?;
17080 Some(
17081 (1..=local_unitig_bucket_count)
17082 .map(|bucket_id| {
17083 LocalUnitigBucketWriter::create(
17084 &local_unitig_bucket_dir,
17085 bucket_id as u16,
17086 color_path.is_some(),
17087 )
17088 .map(Mutex::new)
17089 })
17090 .collect::<Result<Vec<_>, _>>()?,
17091 )
17092 };
17093 let mut color_run_writer = if workers == 1 {
17094 color_path.map(ColorRunSidecarWriter::create).transpose()?
17095 } else {
17096 None
17097 };
17098 let concurrent_color_runs = if workers > 1 && local_unitig_bucket_count == 0 {
17099 color_path
17100 .map(ConcurrentColorRunSidecarWriter::create)
17101 .transpose()?
17102 } else {
17103 None
17104 };
17105 let started = Instant::now();
17106 if groups.len() >= 1024 {
17107 eprintln!(
17108 "cuttlefish: contracting {} local subgraph(s) from {} bucket file(s) with {} worker(s)",
17109 groups.len(),
17110 entries.len(),
17111 workers
17112 );
17113 }
17114
17115 let mut build_elapsed = Duration::default();
17116 let mut contract_elapsed = Duration::default();
17117 let mut sink_io_elapsed = Duration::default();
17118 let mut sink_color_elapsed = Duration::default();
17119 let mut sink_edge_elapsed = Duration::default();
17120
17121 if workers == 1 {
17122 let mut reusable_vertices = None;
17123 for (offset, group) in groups.iter().enumerate() {
17124 let output = contract_local_subgraph::<K>(
17125 store,
17126 group,
17127 cutoff,
17128 color_repository.as_ref(),
17129 &mut reusable_vertices,
17130 color_repository.is_none(),
17131 )?;
17132 SerialLocalOutput {
17133 trivial_output: &mut trivial_output,
17134 trivial_path: &trivial_path,
17135 labels: &mut labels,
17136 label_path,
17137 unitigs: &mut unitigs,
17138 unitig_path: &unitig_path,
17139 inputs: &mut inputs,
17140 edge_matrix: &mut edge_matrix,
17141 ranges: &mut ranges,
17142 color_run_writer: color_run_writer.as_mut(),
17143 trivial_unitigs: &mut trivial_unitigs,
17144 trivial_bases: &mut trivial_bases,
17145 label_offset: &mut label_offset,
17146 build_elapsed: &mut build_elapsed,
17147 contract_elapsed: &mut contract_elapsed,
17148 compact_unitigs,
17149 }
17150 .append(output)?;
17151 report_local_contraction_progress(offset + 1, groups.len(), started);
17152 }
17153 } else {
17154 let total_groups = groups.len();
17155 let next_group = AtomicUsize::new(0);
17156 let completed = Arc::new(AtomicUsize::new(0));
17157 let sink = ConcurrentLocalOutputSink::<K> {
17158 labels: labels
17159 .get_ref()
17160 .try_clone()
17161 .map_err(|source| DiscontinuityInputError::Io {
17162 path: label_path.to_path_buf(),
17163 source,
17164 })?,
17165 label_path,
17166 unitigs: unitigs.get_ref().try_clone().map_err(|source| {
17167 DiscontinuityInputError::Io {
17168 path: unitig_path.clone(),
17169 source,
17170 }
17171 })?,
17172 unitig_path: &unitig_path,
17173 next_label: AtomicU64::new(0),
17174 next_unitig: AtomicUsize::new(0),
17175 weak_superkmers: AtomicU64::new(0),
17176 discontinuity_exits: AtomicU64::new(0),
17177 unitig_bases: AtomicU64::new(0),
17178 build_nanos: AtomicU64::new(0),
17179 contract_nanos: AtomicU64::new(0),
17180 sink_io_nanos: AtomicU64::new(0),
17181 sink_color_nanos: AtomicU64::new(0),
17182 sink_edge_nanos: AtomicU64::new(0),
17183 ranges: Mutex::new(Vec::with_capacity(groups.len())),
17184 edge_writers: ConcurrentBlockedEdgeWriters::new(&edge_matrix),
17185 color_repository: color_repository.as_ref(),
17186 color_runs: concurrent_color_runs.as_ref(),
17187 local_unitig_writers: local_unitig_writers.as_deref(),
17188 compact_unitigs,
17189 trivial_output: trivial_output.get_ref().try_clone().map_err(|source| {
17190 DiscontinuityInputError::Io {
17191 path: trivial_path.clone(),
17192 source,
17193 }
17194 })?,
17195 trivial_path: &trivial_path,
17196 next_trivial_byte: AtomicU64::new(0),
17197 trivial_unitigs: AtomicU64::new(0),
17198 trivial_bases: AtomicU64::new(0),
17199 marker: PhantomData,
17200 };
17201 let worker_result = std::thread::scope(|scope| {
17202 let mut handles = Vec::new();
17203 for worker_id in 0..workers {
17204 let completed = Arc::clone(&completed);
17205 let next_group = &next_group;
17206 let groups = &groups;
17207 let sink = &sink;
17208 handles.push(scope.spawn(move || {
17209 let mut reusable_vertices = None;
17210 let buckets_per_worker = (local_unitig_bucket_count / workers.max(1)).max(1);
17213 let mut bucket_rotation = 0usize;
17214 loop {
17215 let bucket_id = if local_unitig_bucket_count == 0 {
17216 0
17217 } else {
17218 let owned = worker_id * buckets_per_worker
17219 + bucket_rotation % buckets_per_worker;
17220 bucket_rotation += 1;
17221 (owned % local_unitig_bucket_count + 1) as u16
17222 };
17223 let group_idx = next_group.fetch_add(1, Ordering::Relaxed);
17224 let Some(group) = groups.get(group_idx) else {
17225 break;
17226 };
17227 let output = contract_local_subgraph::<K>(
17228 store,
17229 group,
17230 cutoff,
17231 sink.color_repository,
17232 &mut reusable_vertices,
17233 sink.color_repository.is_none(),
17234 )
17235 .and_then(|output| sink.write(output, bucket_id));
17236 let done = completed.fetch_add(1, Ordering::Relaxed) + 1;
17237 report_local_contraction_progress(done, total_groups, started);
17238 let should_stop = output.is_err();
17239 if should_stop {
17240 return output;
17241 }
17242 }
17243 Ok(())
17244 }));
17245 }
17246 let mut first_error = None;
17247 for handle in handles {
17248 let result = handle
17249 .join()
17250 .map_err(|_| DiscontinuityInputError::WorkerPanic)?;
17251 if let Err(err) = result
17252 && first_error.is_none()
17253 {
17254 first_error = Some(err);
17255 }
17256 }
17257
17258 if let Some(err) = first_error {
17259 Err(err)
17260 } else {
17261 Ok(())
17262 }
17263 });
17264 worker_result?;
17265 let edge_finish_started = Instant::now();
17266 sink.edge_writers
17267 .finish_into(&mut edge_matrix)
17268 .map_err(serial_collation_to_input_error)?;
17269 let edge_finish_elapsed = edge_finish_started.elapsed();
17270 inputs.stats.weak_superkmers = sink.weak_superkmers.load(Ordering::Relaxed);
17271 inputs.stats.local_unitigs = sink.next_unitig.load(Ordering::Relaxed) as u64;
17272 trivial_unitigs = sink.trivial_unitigs.load(Ordering::Relaxed);
17273 trivial_bases = sink.trivial_bases.load(Ordering::Relaxed);
17274 inputs.stats.local_unitigs += trivial_unitigs;
17275 inputs.stats.discontinuity_exits = sink.discontinuity_exits.load(Ordering::Relaxed);
17276 inputs.stats.unitig_bases = sink.unitig_bases.load(Ordering::Relaxed);
17277 inputs.stats.unitig_bases += trivial_bases;
17278 build_elapsed = Duration::from_nanos(sink.build_nanos.load(Ordering::Relaxed));
17279 contract_elapsed = Duration::from_nanos(sink.contract_nanos.load(Ordering::Relaxed));
17280 sink_io_elapsed = Duration::from_nanos(sink.sink_io_nanos.load(Ordering::Relaxed));
17281 sink_color_elapsed = Duration::from_nanos(sink.sink_color_nanos.load(Ordering::Relaxed));
17282 sink_edge_elapsed = Duration::from_nanos(sink.sink_edge_nanos.load(Ordering::Relaxed));
17283 ranges = sink
17284 .ranges
17285 .into_inner()
17286 .map_err(|_| DiscontinuityInputError::WorkerPanic)?;
17287 ranges.sort_by_key(|range| range.start_unitig);
17288 eprintln!(
17289 "cuttlefish: local edge-writer finalization {:.3}s",
17290 edge_finish_elapsed.as_secs_f64()
17291 );
17292 }
17293
17294 eprintln!(
17295 "cuttlefish: local worker time: bucket read/build {:.3}s, unitig walk {:.3}s",
17296 build_elapsed.as_secs_f64(),
17297 contract_elapsed.as_secs_f64()
17298 );
17299 if workers > 1 {
17300 eprintln!(
17301 "cuttlefish: local sink worker time: label/unitig I/O {:.3}s, color resolve/write {:.3}s, edge/range emission {:.3}s",
17302 sink_io_elapsed.as_secs_f64(),
17303 sink_color_elapsed.as_secs_f64(),
17304 sink_edge_elapsed.as_secs_f64(),
17305 );
17306 }
17307
17308 let stream_finish_started = Instant::now();
17309 labels
17310 .flush()
17311 .map_err(|source| DiscontinuityInputError::Io {
17312 path: label_path.to_path_buf(),
17313 source,
17314 })?;
17315 unitigs
17316 .flush()
17317 .map_err(|source| DiscontinuityInputError::Io {
17318 path: unitig_path.clone(),
17319 source,
17320 })?;
17321 trivial_output
17322 .flush()
17323 .map_err(|source| DiscontinuityInputError::Io {
17324 path: trivial_path.clone(),
17325 source,
17326 })?;
17327 drop(labels);
17328 drop(unitigs);
17329 edge_matrix
17330 .flush_all_with_threads(threads)
17331 .map_err(serial_collation_to_input_error)?;
17332 let stream_finish_elapsed = stream_finish_started.elapsed();
17333 let color_runs_started = Instant::now();
17334 let color_runs = match (color_run_writer, concurrent_color_runs) {
17335 (Some(writer), None) => Some(writer.finish()?),
17336 (None, Some(writer)) => Some(writer.finish()?),
17337 (None, None) => None,
17338 (Some(_), Some(_)) => unreachable!("color writers are mutually exclusive"),
17339 };
17340 let color_runs_elapsed = color_runs_started.elapsed();
17341 let unitig_buckets_started = Instant::now();
17342 let local_unitig_buckets = local_unitig_writers
17343 .take()
17344 .map(|writers| finish_local_unitig_writers(writers, workers))
17345 .transpose()?;
17346 let unitig_buckets_elapsed = unitig_buckets_started.elapsed();
17347 let color_repository_started = Instant::now();
17348 let color_repository = color_repository
17349 .map(|repository| repository.finish())
17350 .transpose()?;
17351 let color_repository_elapsed = color_repository_started.elapsed();
17352 if workers > 1 {
17353 eprintln!(
17354 "cuttlefish: local finalization detail: streams/matrix {:.3}s, color runs {:.3}s, unitig buckets {:.3}s, color repository {:.3}s",
17355 stream_finish_elapsed.as_secs_f64(),
17356 color_runs_elapsed.as_secs_f64(),
17357 unitig_buckets_elapsed.as_secs_f64(),
17358 color_repository_elapsed.as_secs_f64(),
17359 );
17360 }
17361 let trivial_is_output = direct_output_path.is_some();
17362 let trivial_fasta = if trivial_unitigs == 0 {
17363 if !trivial_is_output {
17365 let _ = fs::remove_file(&trivial_path);
17366 }
17367 None
17368 } else {
17369 Some((trivial_path, trivial_unitigs, trivial_bases))
17370 };
17371 Ok(ExternalDiscontinuityInputs {
17372 unitig_path,
17373 label_path: label_path.to_path_buf(),
17374 unitigs: inputs.stats.local_unitigs.saturating_sub(trivial_unitigs) as usize,
17376 compact_unitigs,
17377 ranges,
17378 edge_matrix: Some(edge_matrix),
17379 color_runs,
17380 local_unitig_buckets,
17381 local_unitig_bucket_dir: (local_unitig_bucket_count != 0)
17382 .then(|| local_unitig_bucket_dir.clone()),
17383 trivial_fasta,
17384 trivial_is_output,
17385 color_repository,
17386 stats: inputs.stats,
17387 })
17388}
17389
17390fn unitig_table_path_for_labels(label_path: &Path) -> PathBuf {
17391 let Some(file_name) = label_path.file_name().and_then(|name| name.to_str()) else {
17392 return label_path.with_extension("unitigs");
17393 };
17394 let unitig_name = if let Some(prefix) = file_name.strip_suffix("labels") {
17395 format!("{prefix}unitigs")
17396 } else {
17397 format!("{file_name}.unitigs")
17398 };
17399 label_path.with_file_name(unitig_name)
17400}
17401
17402#[allow(clippy::too_many_arguments)]
17403struct SerialLocalOutput<'a, const K: usize> {
17411 trivial_output: &'a mut BufWriter<File>,
17412 trivial_path: &'a Path,
17413 labels: &'a mut BufWriter<File>,
17414 label_path: &'a Path,
17415 unitigs: &'a mut BufWriter<File>,
17416 unitig_path: &'a Path,
17417 inputs: &'a mut DiscontinuityInputs<K>,
17418 edge_matrix: &'a mut BlockedEdgeMatrix<K>,
17419 ranges: &'a mut Vec<ExternalLocalUnitigRange>,
17420 color_run_writer: Option<&'a mut ColorRunSidecarWriter>,
17421 trivial_unitigs: &'a mut u64,
17423 trivial_bases: &'a mut u64,
17424 label_offset: &'a mut u64,
17425 build_elapsed: &'a mut Duration,
17426 contract_elapsed: &'a mut Duration,
17427 compact_unitigs: bool,
17428}
17429
17430impl<const K: usize> SerialLocalOutput<'_, K> {
17431 fn append(&mut self, output: LocalContractionOutput<K>) -> Result<(), DiscontinuityInputError> {
17432 self.trivial_output
17433 .write_all(&output.trivial_fasta)
17434 .map_err(|source| DiscontinuityInputError::Io {
17435 path: self.trivial_path.to_path_buf(),
17436 source,
17437 })?;
17438 *self.trivial_unitigs += output.trivial_unitigs;
17439 *self.trivial_bases += output.trivial_bases;
17440 self.inputs.stats.local_unitigs += output.trivial_unitigs;
17441 self.inputs.stats.unitig_bases += output.trivial_bases;
17442 self.inputs.stats.weak_superkmers += output.weak_superkmers;
17443 let start_unitig = usize::try_from(self.inputs.stats.local_unitigs - *self.trivial_unitigs)
17451 .unwrap_or(usize::MAX);
17452 let output_unitigs = output.unitigs.len();
17453 let output_label_len = output.labels.len() as u64;
17454 let color_start = self
17455 .color_run_writer
17456 .as_deref()
17457 .map(ColorRunSidecarWriter::position)
17458 .unwrap_or(0);
17459 if output_unitigs != 0 {
17460 self.ranges.push(ExternalLocalUnitigRange {
17461 start_unitig,
17462 unitigs: output_unitigs,
17463 label_start: *self.label_offset,
17464 label_len: output_label_len,
17465 color_start,
17466 });
17467 }
17468 self.labels
17469 .write_all(&output.labels)
17470 .map_err(|source| DiscontinuityInputError::Io {
17471 path: self.label_path.to_path_buf(),
17472 source,
17473 })?;
17474 self.edge_matrix
17475 .add_prepared_edges(&output.matrix_edges, start_unitig)
17476 .map_err(serial_collation_to_input_error)?;
17477 let mut output_colors = output.color_runs.map(Vec::into_iter);
17478 for mut unitig in output.unitigs {
17479 self.inputs.stats.local_unitigs += 1;
17480 self.inputs.stats.discontinuity_exits +=
17481 u64::from(unitig.left_exit().is_some()) + u64::from(unitig.right_exit().is_some());
17482 self.inputs.stats.unitig_bases += unitig.label_len as u64;
17483 unitig.label_start += *self.label_offset;
17484 write_discontinuity_unitig_record(
17485 self.unitigs,
17486 self.unitig_path,
17487 &unitig,
17488 self.compact_unitigs,
17489 )?;
17490 if let Some(writer) = self.color_run_writer.as_deref_mut() {
17491 let runs = output_colors
17492 .as_mut()
17493 .and_then(Iterator::next)
17494 .ok_or(DiscontinuityInputError::MissingColorRuns)?;
17495 writer.write_unitig(&runs)?;
17496 }
17497 }
17498 if output_colors
17499 .as_mut()
17500 .is_some_and(|colors| colors.next().is_some())
17501 {
17502 return Err(DiscontinuityInputError::MissingColorRuns);
17503 }
17504 *self.label_offset += output.labels.len() as u64;
17505 *self.build_elapsed += output.build_elapsed;
17506 *self.contract_elapsed += output.contract_elapsed;
17507 Ok(())
17508 }
17509}
17510
17511fn write_discontinuity_unitig_record<const K: usize>(
17512 out: &mut BufWriter<File>,
17513 path: &Path,
17514 unitig: &DiscontinuityUnitig<K>,
17515 compact: bool,
17516) -> Result<(), DiscontinuityInputError> {
17517 if compact {
17518 let mut bytes = [0u8; 8];
17519 bytes[..4].copy_from_slice(&unitig.label_len.to_le_bytes());
17520 bytes[4] = unitig.flags;
17521 return out
17522 .write_all(&bytes)
17523 .map_err(|source| DiscontinuityInputError::Io {
17524 path: path.to_path_buf(),
17525 source,
17526 });
17527 }
17528 let mut stored = MaybeUninit::<DiscontinuityUnitig<K>>::zeroed();
17529 unsafe {
17530 let ptr = stored.as_mut_ptr();
17531 (*ptr).label_start = unitig.label_start;
17532 (*ptr).left_vertex = unitig.left_vertex;
17533 (*ptr).right_vertex = unitig.right_vertex;
17534 (*ptr).label_len = unitig.label_len;
17535 (*ptr).flags = unitig.flags;
17536 let stored = stored.assume_init();
17537 let bytes = std::slice::from_raw_parts(
17538 (&stored as *const DiscontinuityUnitig<K>).cast::<u8>(),
17539 std::mem::size_of::<DiscontinuityUnitig<K>>(),
17540 );
17541 out.write_all(bytes)
17542 .map_err(|source| DiscontinuityInputError::Io {
17543 path: path.to_path_buf(),
17544 source,
17545 })?;
17546 }
17547 Ok(())
17548}
17549
17550struct LocalContractionOutput<const K: usize> {
17551 index: usize,
17552 weak_superkmers: u64,
17553 build_elapsed: Duration,
17554 contract_elapsed: Duration,
17555 unitigs: Vec<DiscontinuityUnitig<K>>,
17556 labels: Vec<u8>,
17557 matrix_edges: Vec<PreparedBlockedEdge>,
17558 color_runs: Option<Vec<Vec<UnitigColor>>>,
17559 trivial_fasta: Vec<u8>,
17560 trivial_unitigs: u64,
17561 trivial_bases: u64,
17562}
17563
17564struct ConcurrentLocalOutputSink<'a, const K: usize> {
17565 labels: File,
17566 label_path: &'a Path,
17567 unitigs: File,
17568 unitig_path: &'a Path,
17569 next_label: AtomicU64,
17570 next_unitig: AtomicUsize,
17571 weak_superkmers: AtomicU64,
17572 discontinuity_exits: AtomicU64,
17573 unitig_bases: AtomicU64,
17574 build_nanos: AtomicU64,
17575 contract_nanos: AtomicU64,
17576 sink_io_nanos: AtomicU64,
17577 sink_color_nanos: AtomicU64,
17578 sink_edge_nanos: AtomicU64,
17579 ranges: Mutex<Vec<ExternalLocalUnitigRange>>,
17580 edge_writers: ConcurrentBlockedEdgeWriters,
17581 color_repository: Option<&'a ConcurrentColorRepository>,
17582 color_runs: Option<&'a ConcurrentColorRunSidecarWriter>,
17583 local_unitig_writers: Option<&'a [Mutex<LocalUnitigBucketWriter>]>,
17584 compact_unitigs: bool,
17585 trivial_output: File,
17586 trivial_path: &'a Path,
17587 next_trivial_byte: AtomicU64,
17588 trivial_unitigs: AtomicU64,
17589 trivial_bases: AtomicU64,
17590 marker: PhantomData<[(); K]>,
17591}
17592
17593impl<'a, const K: usize> ConcurrentLocalOutputSink<'a, K> {
17594 fn write(
17595 &self,
17596 mut output: LocalContractionOutput<K>,
17597 unitig_bucket: u16,
17598 ) -> Result<(), DiscontinuityInputError> {
17599 let io_started = Instant::now();
17600 if !output.trivial_fasta.is_empty() {
17601 let offset = self
17602 .next_trivial_byte
17603 .fetch_add(output.trivial_fasta.len() as u64, Ordering::Relaxed);
17604 self.trivial_output
17605 .write_all_at(&output.trivial_fasta, offset)
17606 .map_err(|source| DiscontinuityInputError::Io {
17607 path: self.trivial_path.to_path_buf(),
17608 source,
17609 })?;
17610 self.trivial_unitigs
17611 .fetch_add(output.trivial_unitigs, Ordering::Relaxed);
17612 self.trivial_bases
17613 .fetch_add(output.trivial_bases, Ordering::Relaxed);
17614 }
17615 let output_unitigs = output.unitigs.len();
17616 let unitig_base = self
17617 .next_unitig
17618 .fetch_add(output_unitigs, Ordering::Relaxed);
17619 let label_base = self
17620 .next_label
17621 .fetch_add(output.labels.len() as u64, Ordering::Relaxed);
17622
17623 let bucket_local = unitig_bucket != 0;
17624 if !bucket_local {
17625 self.labels
17626 .write_all_at(&output.labels, label_base)
17627 .map_err(|source| DiscontinuityInputError::Io {
17628 path: self.label_path.to_path_buf(),
17629 source,
17630 })?;
17631 }
17632
17633 let mut encoded_unitigs = Vec::with_capacity(
17634 output_unitigs * external_unitig_record_len::<K>(self.compact_unitigs),
17635 );
17636 let mut exits = 0u64;
17637 let mut bases = 0u64;
17638 for unitig in &mut output.unitigs {
17639 exits +=
17640 u64::from(unitig.left_exit().is_some()) + u64::from(unitig.right_exit().is_some());
17641 bases += u64::from(unitig.label_len);
17642 if !bucket_local {
17643 unitig.label_start += label_base;
17644 append_encoded_discontinuity_unitig_record(
17645 &mut encoded_unitigs,
17646 unitig,
17647 self.compact_unitigs,
17648 );
17649 }
17650 }
17651 let unitig_offset = unitig_base
17652 .checked_mul(external_unitig_record_len::<K>(self.compact_unitigs))
17653 .ok_or_else(|| DiscontinuityInputError::Io {
17654 path: self.unitig_path.to_path_buf(),
17655 source: std::io::Error::from(std::io::ErrorKind::FileTooLarge),
17656 })? as u64;
17657 if !bucket_local {
17658 self.unitigs
17659 .write_all_at(&encoded_unitigs, unitig_offset)
17660 .map_err(|source| DiscontinuityInputError::Io {
17661 path: self.unitig_path.to_path_buf(),
17662 source,
17663 })?;
17664 }
17665 self.sink_io_nanos.fetch_add(
17666 io_started.elapsed().as_nanos().min(u128::from(u64::MAX)) as u64,
17667 Ordering::Relaxed,
17668 );
17669
17670 let color_started = Instant::now();
17671 let resolved_colors = if self.color_repository.is_some() {
17672 let runs = output
17673 .color_runs
17674 .take()
17675 .ok_or(DiscontinuityInputError::MissingColorRuns)?;
17676 if runs.len() != output_unitigs {
17677 return Err(DiscontinuityInputError::MissingColorRuns);
17678 }
17679 Some(runs)
17680 } else {
17681 None
17682 };
17683 let color_start = if let Some(writer) = self.color_runs {
17684 writer.write_unitigs(
17685 resolved_colors
17686 .as_deref()
17687 .ok_or(DiscontinuityInputError::MissingColorRuns)?,
17688 )?
17689 } else {
17690 0
17691 };
17692 let local_unitig_base = if unitig_bucket != 0 {
17693 let writers = self
17694 .local_unitig_writers
17695 .ok_or(DiscontinuityInputError::WorkerPanic)?;
17696 let writer = writers
17697 .get(unitig_bucket as usize - 1)
17698 .ok_or(DiscontinuityInputError::WorkerPanic)?;
17699 writer
17700 .lock()
17701 .map_err(|_| DiscontinuityInputError::WorkerPanic)?
17702 .write::<K>(&output.labels, &output.unitigs, resolved_colors.as_deref())?
17703 } else {
17704 unitig_base
17705 };
17706 self.sink_color_nanos.fetch_add(
17707 color_started.elapsed().as_nanos().min(u128::from(u64::MAX)) as u64,
17708 Ordering::Relaxed,
17709 );
17710
17711 let edge_started = Instant::now();
17712 self.edge_writers
17713 .add_prepared_edges::<K>(&mut output.matrix_edges, local_unitig_base, unitig_bucket)
17714 .map_err(serial_collation_to_input_error)?;
17715 if output_unitigs != 0 && !bucket_local {
17716 self.ranges
17717 .lock()
17718 .map_err(|_| DiscontinuityInputError::WorkerPanic)?
17719 .push(ExternalLocalUnitigRange {
17720 start_unitig: unitig_base,
17721 unitigs: output_unitigs,
17722 label_start: label_base,
17723 label_len: output.labels.len() as u64,
17724 color_start,
17725 });
17726 }
17727 self.sink_edge_nanos.fetch_add(
17728 edge_started.elapsed().as_nanos().min(u128::from(u64::MAX)) as u64,
17729 Ordering::Relaxed,
17730 );
17731 self.weak_superkmers
17732 .fetch_add(output.weak_superkmers, Ordering::Relaxed);
17733 self.discontinuity_exits.fetch_add(exits, Ordering::Relaxed);
17734 self.unitig_bases.fetch_add(bases, Ordering::Relaxed);
17735 self.build_nanos.fetch_add(
17736 output.build_elapsed.as_nanos().min(u128::from(u64::MAX)) as u64,
17737 Ordering::Relaxed,
17738 );
17739 self.contract_nanos.fetch_add(
17740 output.contract_elapsed.as_nanos().min(u128::from(u64::MAX)) as u64,
17741 Ordering::Relaxed,
17742 );
17743 Ok(())
17744 }
17745}
17746
17747fn append_encoded_discontinuity_unitig_record<const K: usize>(
17748 output: &mut Vec<u8>,
17749 unitig: &DiscontinuityUnitig<K>,
17750 compact: bool,
17751) {
17752 if compact {
17753 output.extend_from_slice(&unitig.label_len.to_le_bytes());
17754 output.push(unitig.flags);
17755 output.extend_from_slice(&[0; 3]);
17756 return;
17757 }
17758 let mut stored = MaybeUninit::<DiscontinuityUnitig<K>>::zeroed();
17759 unsafe {
17760 let ptr = stored.as_mut_ptr();
17761 (*ptr).label_start = unitig.label_start;
17762 (*ptr).left_vertex = unitig.left_vertex;
17763 (*ptr).right_vertex = unitig.right_vertex;
17764 (*ptr).label_len = unitig.label_len;
17765 (*ptr).flags = unitig.flags;
17766 let stored = stored.assume_init();
17767 output.extend_from_slice(std::slice::from_raw_parts(
17768 (&stored as *const DiscontinuityUnitig<K>).cast::<u8>(),
17769 std::mem::size_of::<DiscontinuityUnitig<K>>(),
17770 ));
17771 }
17772}
17773
17774#[derive(Debug, Clone)]
17775struct LocalBucketGroup {
17776 index: usize,
17777 graph_id: usize,
17778 stored_bytes: u64,
17779 entries: Vec<BucketManifestEntry>,
17780}
17781
17782const VERTEX_MAP_REUSE_SLACK: usize = 8;
17785
17786const VERTEX_MAP_MEAN_DECAY: usize = 8;
17788
17789const VERTEX_MAP_ALWAYS_REUSE_CAPACITY: usize = 4096;
17792
17793struct ReusableVertexMap<const K: usize> {
17802 map: LocalVertexMap<K>,
17803 mean_vertices: usize,
17804}
17805
17806impl<const K: usize> ReusableVertexMap<K> {
17807 fn worth_carrying(&self) -> bool {
17810 self.map.capacity()
17811 <= self
17812 .mean_vertices
17813 .saturating_mul(vertex_map_reuse_slack())
17814 .max(VERTEX_MAP_ALWAYS_REUSE_CAPACITY)
17815 }
17816
17817 fn observe(mean: usize, vertices: usize) -> usize {
17819 if mean == 0 {
17820 return vertices;
17821 }
17822 (mean.saturating_mul(VERTEX_MAP_MEAN_DECAY - 1) + vertices) / VERTEX_MAP_MEAN_DECAY
17823 }
17824}
17825
17826fn vertex_map_reuse_slack() -> usize {
17827 static SLACK: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
17828 *SLACK.get_or_init(|| {
17829 std::env::var("CF3_RS_VERTEX_REUSE_SLACK")
17830 .ok()
17831 .and_then(|value| value.parse::<usize>().ok())
17832 .unwrap_or(VERTEX_MAP_REUSE_SLACK)
17833 })
17834}
17835
17836fn contract_local_subgraphs<const K: usize>(
17837 store: &BucketStore,
17838 entries: &[BucketManifestEntry],
17839 cutoff: u32,
17840 threads: usize,
17841) -> Result<Vec<LocalContractionOutput<K>>, DiscontinuityInputError> {
17842 let mut groups = local_bucket_groups(entries)?;
17843 groups.sort_by_key(|group| std::cmp::Reverse((group.stored_bytes, group.graph_id)));
17844 let workers = threads.min(groups.len().max(1));
17845 let started = Instant::now();
17846 if groups.len() >= 1024 {
17847 eprintln!(
17848 "cuttlefish: contracting {} local subgraph(s) from {} bucket file(s) with {} worker(s)",
17849 groups.len(),
17850 entries.len(),
17851 workers
17852 );
17853 }
17854 if workers == 1 {
17855 let mut outputs = Vec::with_capacity(groups.len());
17856 let mut reusable_vertices = None;
17857 for (offset, group) in groups.iter().enumerate() {
17858 outputs.push(contract_local_subgraph::<K>(
17859 store,
17860 group,
17861 cutoff,
17862 None,
17863 &mut reusable_vertices,
17864 false,
17865 )?);
17866 report_local_contraction_progress(offset + 1, groups.len(), started);
17867 }
17868 return Ok(outputs);
17869 }
17870
17871 let total_groups = groups.len();
17872 let next_group = AtomicUsize::new(0);
17873 let completed = Arc::new(AtomicUsize::new(0));
17874 let mut outputs = std::thread::scope(|scope| {
17875 let mut handles = Vec::new();
17876 for _ in 0..workers {
17877 let completed = Arc::clone(&completed);
17878 let next_group = &next_group;
17879 let groups = &groups;
17880 handles.push(scope.spawn(move || {
17881 let mut chunk_outputs = Vec::new();
17882 let mut reusable_vertices = None;
17883 loop {
17884 let group_idx = next_group.fetch_add(1, Ordering::Relaxed);
17885 let Some(group) = groups.get(group_idx) else {
17886 break;
17887 };
17888 chunk_outputs.push(contract_local_subgraph::<K>(
17889 store,
17890 group,
17891 cutoff,
17892 None,
17893 &mut reusable_vertices,
17894 false,
17895 )?);
17896 let done = completed.fetch_add(1, Ordering::Relaxed) + 1;
17897 report_local_contraction_progress(done, total_groups, started);
17898 }
17899 Ok::<_, DiscontinuityInputError>(chunk_outputs)
17900 }));
17901 }
17902
17903 let mut outputs = Vec::with_capacity(entries.len());
17904 for handle in handles {
17905 outputs.extend(
17906 handle
17907 .join()
17908 .map_err(|_| DiscontinuityInputError::WorkerPanic)??,
17909 );
17910 }
17911 Ok::<_, DiscontinuityInputError>(outputs)
17912 })?;
17913
17914 outputs.sort_by_key(|output| output.index);
17915 let build_elapsed = outputs
17916 .iter()
17917 .map(|output| output.build_elapsed)
17918 .sum::<Duration>();
17919 let contract_elapsed = outputs
17920 .iter()
17921 .map(|output| output.contract_elapsed)
17922 .sum::<Duration>();
17923 eprintln!(
17924 "cuttlefish: local worker time: bucket read/build {:.3}s, unitig walk {:.3}s",
17925 build_elapsed.as_secs_f64(),
17926 contract_elapsed.as_secs_f64()
17927 );
17928 Ok(outputs)
17929}
17930
17931fn local_bucket_groups(
17932 entries: &[BucketManifestEntry],
17933) -> Result<Vec<LocalBucketGroup>, DiscontinuityInputError> {
17934 let mut by_graph = BTreeMap::<usize, Vec<BucketManifestEntry>>::new();
17935 for entry in entries {
17936 by_graph
17937 .entry(entry.graph_id)
17938 .or_default()
17939 .push(entry.clone());
17940 }
17941
17942 by_graph
17943 .into_iter()
17944 .enumerate()
17945 .map(|(index, (graph_id, entries))| {
17946 let stored_bytes = entries.iter().try_fold(0u64, |bytes, entry| {
17950 entry
17951 .stored_bytes()
17952 .map(|len| bytes.saturating_add(len))
17953 .map_err(DiscontinuityInputError::from)
17954 })?;
17955 Ok(LocalBucketGroup {
17956 index,
17957 graph_id,
17958 stored_bytes,
17959 entries,
17960 })
17961 })
17962 .collect()
17963}
17964
17965fn report_local_contraction_progress(done: usize, total: usize, started: Instant) {
17966 if total < 1024 {
17967 return;
17968 }
17969 if done == total || done % 1024 == 0 {
17970 eprintln!(
17971 "cuttlefish: contracted {done}/{total} local subgraph bucket(s) in {:.1}s",
17972 started.elapsed().as_secs_f64()
17973 );
17974 report_process_memory(&format!("local contraction progress {done}/{total}"));
17975 }
17976}
17977
17978fn report_discontinuity_contraction_progress(done: usize, total: usize, started: Instant) {
17979 if total < 16 {
17980 return;
17981 }
17982 if done == total || done % 16 == 0 {
17983 eprintln!(
17984 "cuttlefish: contracted {done}/{total} discontinuity partition(s) in {:.1}s",
17985 started.elapsed().as_secs_f64()
17986 );
17987 }
17988}
17989
17990fn contract_local_subgraph<const K: usize>(
17991 store: &BucketStore,
17992 group: &LocalBucketGroup,
17993 cutoff: u32,
17994 color_repository: Option<&ConcurrentColorRepository>,
17995 reusable_vertices: &mut Option<ReusableVertexMap<K>>,
17996 emit_trivial_fasta: bool,
17997) -> Result<LocalContractionOutput<K>, DiscontinuityInputError> {
17998 let build_start = Instant::now();
17999 let carried = reusable_vertices.take();
18000 let mean_vertices = carried.as_ref().map_or(0, |held| held.mean_vertices);
18003 let mut subgraph = LocalSubgraph::<K>::from_manifest_entries_reusing(
18004 store,
18005 &group.entries,
18006 cutoff,
18007 carried
18008 .filter(ReusableVertexMap::worth_carrying)
18009 .map(|held| held.map),
18010 )?;
18011 let build_elapsed = build_start.elapsed();
18012 let weak_superkmers = subgraph.stats.weak_superkmers;
18013 let graph_id = subgraph.graph_id;
18014 debug_assert_eq!(graph_id, group.graph_id);
18015 let mut inputs = DiscontinuityInputs::empty(DiscontinuityInputStats::default());
18016
18017 let contract_start = Instant::now();
18018 let mut color_runs = None;
18019 let mut trivial_fasta = Vec::new();
18020 let mut trivial_unitigs = 0u64;
18021 let mut trivial_bases = 0u64;
18022 let mut matrix_edges = Vec::new();
18023 let mut emit_unitig = |unitig: LocalUnitig<K>| {
18024 let left_exit = unitig
18025 .left_exit
18026 .map(|(vertex, side)| DiscontinuityEndpoint { vertex, side });
18027 let right_exit = unitig
18028 .right_exit
18029 .map(|(vertex, side)| DiscontinuityEndpoint { vertex, side });
18030
18031 if emit_trivial_fasta && left_exit.is_none() && right_exit.is_none() {
18032 let label = canonical_label(unitig.label);
18033 trivial_fasta.extend_from_slice(b">0\n");
18034 trivial_fasta.extend_from_slice(&label);
18035 trivial_fasta.push(b'\n');
18036 trivial_unitigs += 1;
18037 trivial_bases += label.len() as u64;
18038 return;
18039 }
18040
18041 let unitig_index = inputs.unitigs.len();
18042 inputs.push_unitig(OwnedDiscontinuityUnitig {
18043 graph_id,
18044 label: unitig.label,
18045 left_exit,
18046 right_exit,
18047 is_cycle: unitig.is_cycle,
18048 });
18049 if let Some(edge) = prepare_unitig_blocked_edge(
18050 &inputs.unitigs[unitig_index],
18051 unitig_index,
18052 DEFAULT_VERTEX_PARTITIONS,
18053 ) {
18054 matrix_edges.push(edge);
18055 }
18056 };
18057 if subgraph.colored {
18058 let repository = color_repository.ok_or(DiscontinuityInputError::MissingColorRuns)?;
18059 let runs = subgraph.contract_colored_resolved_with(
18060 store,
18061 &group.entries,
18062 repository,
18063 group.index % repository.worker_count(),
18064 &mut emit_unitig,
18065 )?;
18066 color_runs = Some(runs);
18067 } else {
18068 subgraph.contract_compact_with(emit_unitig)?;
18069 }
18070 let contract_elapsed = contract_start.elapsed();
18071
18072 let output = LocalContractionOutput {
18073 index: group.index,
18074 weak_superkmers,
18075 build_elapsed,
18076 contract_elapsed,
18077 unitigs: inputs.unitigs,
18078 labels: inputs.labels,
18079 matrix_edges,
18080 color_runs,
18081 trivial_fasta,
18082 trivial_unitigs,
18083 trivial_bases,
18084 };
18085 let map = subgraph.into_vertex_map();
18086 *reusable_vertices = Some(ReusableVertexMap {
18087 mean_vertices: ReusableVertexMap::<K>::observe(mean_vertices, map.len()),
18088 map,
18089 });
18090 if !keep_intermediates() {
18091 for entry in &group.entries {
18096 let BucketLocation::File(path) = &entry.location else {
18097 if let BucketLocation::Container {
18103 container,
18104 segments,
18105 ..
18106 } = &entry.location
18107 && let Some(containers) = store.containers()
18108 {
18109 containers.release_segments(*container, segments);
18110 }
18111 continue;
18112 };
18113 match fs::remove_file(path) {
18114 Ok(()) => {}
18115 Err(source) if source.kind() == std::io::ErrorKind::NotFound => {}
18116 Err(source) => {
18117 return Err(DiscontinuityInputError::Io {
18118 path: path.to_path_buf(),
18119 source,
18120 });
18121 }
18122 }
18123 }
18124 }
18125 Ok(output)
18126}
18127
18128#[derive(Debug)]
18129pub enum DiscontinuityInputError {
18130 Bucket(BucketError),
18131 LocalSubgraph(LocalSubgraphError),
18132 Color(ColorError),
18133 Io {
18134 path: PathBuf,
18135 source: std::io::Error,
18136 },
18137 InvalidCutoff,
18138 InvalidThreadCount,
18139 WorkerPanic,
18140 MissingColorRuns,
18141}
18142
18143impl From<BucketError> for DiscontinuityInputError {
18144 fn from(value: BucketError) -> Self {
18145 Self::Bucket(value)
18146 }
18147}
18148
18149impl From<LocalSubgraphError> for DiscontinuityInputError {
18150 fn from(value: LocalSubgraphError) -> Self {
18151 Self::LocalSubgraph(value)
18152 }
18153}
18154
18155impl From<ColorError> for DiscontinuityInputError {
18156 fn from(value: ColorError) -> Self {
18157 Self::Color(value)
18158 }
18159}
18160
18161impl std::fmt::Display for DiscontinuityInputError {
18162 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
18163 match self {
18164 Self::Bucket(err) => write!(f, "{err}"),
18165 Self::LocalSubgraph(err) => write!(f, "{err}"),
18166 Self::Color(err) => write!(f, "{err}"),
18167 Self::Io { path, source } => write!(f, "{}: {source}", path.display()),
18168 Self::InvalidCutoff => write!(f, "discontinuity input cutoff must be at least 1"),
18169 Self::InvalidThreadCount => {
18170 write!(f, "discontinuity input thread count must be at least 1")
18171 }
18172 Self::WorkerPanic => write!(f, "local subgraph worker thread panicked"),
18173 Self::MissingColorRuns => write!(f, "colored local unitig is missing color runs"),
18174 }
18175 }
18176}
18177
18178impl std::error::Error for DiscontinuityInputError {}
18179
18180#[cfg(test)]
18181mod materialized_record_tests {
18182 use super::*;
18183
18184 #[test]
18185 fn compact_discontinuity_edge_uses_cpp_widths() {
18186 let edge = DiscontinuityEdge::<31> {
18187 first: MatrixEndpoint::Vertex(DiscontinuityEndpoint {
18188 vertex: Kmer::from_bits(0x1234),
18189 side: Side::Back,
18190 }),
18191 second: MatrixEndpoint::Vertex(DiscontinuityEndpoint {
18192 vertex: Kmer::from_bits(0x5678),
18193 side: Side::Front,
18194 }),
18195 weight: 65_534,
18196 unitig_bucket: 1_023,
18197 unitig_index: u32::MAX as usize - 1,
18198 unitig_exit_side: Side::Back,
18199 phantom_unitig: None,
18200 swapped: true,
18201 };
18202
18203 assert_eq!(discontinuity_edge_record_len::<31>(), 25);
18204 assert_eq!(blocked_edge_unitig_offset::<31>(), 18);
18205 let mut encoded = encode_discontinuity_edge(&edge);
18206 assert_eq!(decode_discontinuity_edge::<31>(&encoded[..25]), edge);
18207
18208 set_encoded_edge_unitig_bucket::<31>(&mut encoded, 511);
18209 assert_eq!(
18210 decode_discontinuity_edge::<31>(&encoded[..25]).unitig_bucket,
18211 511
18212 );
18213 }
18214
18215 #[test]
18224 fn merging_concurrent_appends_transfers_extents_and_counts() {
18225 let directory = std::env::temp_dir().join(format!(
18226 "cf3-edge-container-merge-{}-{:?}",
18227 std::process::id(),
18228 std::thread::current().id()
18229 ));
18230 let _ = std::fs::remove_dir_all(&directory);
18231
18232 const VERTEX_PARTITIONS: usize = 4;
18233 let record_len = discontinuity_edge_record_len::<31>();
18234 let mut matrix = BlockedEdgeMatrix::<31>::create(&directory, VERTEX_PARTITIONS).unwrap();
18235 let (row, col) = (1usize, 2usize);
18236 let block = matrix.block_index(row, col);
18237
18238 let edge = |id: usize| DiscontinuityEdge::<31> {
18241 first: MatrixEndpoint::Vertex(DiscontinuityEndpoint {
18242 vertex: Kmer::from_bits(id as u128 * 7 + 1),
18243 side: Side::Back,
18244 }),
18245 second: MatrixEndpoint::Vertex(DiscontinuityEndpoint {
18246 vertex: Kmer::from_bits(id as u128 * 11 + 2),
18247 side: Side::Front,
18248 }),
18249 weight: 1,
18250 unitig_bucket: 3,
18251 unitig_index: id,
18252 unitig_exit_side: Side::Front,
18253 phantom_unitig: None,
18254 swapped: false,
18255 };
18256 let prepared = |id: usize| PreparedBlockedEdge {
18257 block,
18258 bytes: encode_discontinuity_edge(&edge(id)),
18259 phi: false,
18260 diagonal: false,
18261 };
18262
18263 let per_writer = (BLOCKED_EDGE_WRITE_BUFFER_BYTES / record_len) * 3;
18267 let mut expected = Vec::new();
18268
18269 for round in 0..3usize {
18271 let appenders = ConcurrentBlockedEdgeWriters::new(&matrix);
18272 let base = round * per_writer * 2;
18273
18274 let direct = (0..per_writer)
18275 .map(|i| prepared(base + i))
18276 .collect::<Vec<_>>();
18277 matrix.add_prepared_edges_absolute(&direct).unwrap();
18278 expected.extend((0..per_writer).map(|i| edge(base + i)));
18279
18280 for i in 0..per_writer {
18281 appenders.add(&prepared(base + per_writer + i)).unwrap();
18282 expected.push(edge(base + per_writer + i));
18283 }
18284
18285 matrix.flush_block(row, col).unwrap();
18289 let added = appenders.merge_block_into(&mut matrix, row, col).unwrap();
18290 assert_eq!(added, per_writer);
18291 }
18292
18293 let flushed = matrix.blocks[block]
18294 .extents
18295 .iter()
18296 .map(|extent| extent.len as usize)
18297 .sum::<usize>();
18298 assert_eq!(matrix.blocks[block].edges, expected.len());
18299 assert_eq!(
18300 flushed + matrix.blocks[block].buffer.len(),
18301 matrix.blocks[block].edges * record_len,
18302 "extent bytes and edge count must agree",
18303 );
18304
18305 let mut read_back = matrix.read_flushed_block(row, col).unwrap();
18306 assert_eq!(read_back.len(), expected.len());
18307 read_back.sort_unstable_by_key(|edge| edge.unitig_index);
18308 expected.sort_unstable_by_key(|edge| edge.unitig_index);
18309 assert_eq!(read_back, expected);
18310
18311 assert!(
18313 matrix
18314 .blocks
18315 .iter()
18316 .enumerate()
18317 .all(|(index, block_state)| index == block || block_state.edges == 0)
18318 );
18319
18320 std::fs::remove_dir_all(&directory).unwrap();
18321 }
18322
18323 #[test]
18324 fn shared_materialized_batch_matches_cpp_buffer_thresholds() {
18325 let mut bucket = PendingMaterializedBucket::default();
18326 bucket.records.resize_with(
18327 SharedMaterializedBatch::FLUSH_BYTES / STITCH_COORD_RECORD_LEN as usize,
18328 || LoadedMaterializedStitchedCoordRecord::new(0, 0, 0, 31, false, false, u32::MAX, 0),
18329 );
18330 assert!(!SharedMaterializedBatch::uncolored_bucket_ready(&bucket));
18331
18332 bucket
18333 .labels
18334 .resize(SharedMaterializedBatch::FLUSH_BYTES, 0);
18335 assert!(SharedMaterializedBatch::uncolored_bucket_ready(&bucket));
18336 assert!(!SharedMaterializedBatch::colored_bucket_ready(&bucket));
18337
18338 bucket.colors.resize(
18339 SharedMaterializedBatch::FLUSH_BYTES / std::mem::size_of::<UnitigColor>(),
18340 UnitigColor::new(0, crate::state::ColorCoordinate::from_u40(0)),
18341 );
18342 assert!(SharedMaterializedBatch::colored_bucket_ready(&bucket));
18343 }
18344
18345 #[test]
18346 fn materialized_coordinate_plan_respects_descriptor_budget() {
18347 const SMALL_GRAPH: u64 = 4_000_000_000;
18350 for &(limit, open, threads, colored) in &[
18351 (1024usize, 64usize, 128usize, false),
18352 (1024, 64, 128, true),
18353 (1024, 64, 256, false),
18354 (1024, 64, 256, true),
18355 (65_536, 64, 256, true),
18356 ] {
18357 let (buckets, workers, open_writers) =
18358 materialized_coordinate_plan(limit, open, threads, colored, SMALL_GRAPH);
18359 let writer_files = if colored { 3 } else { 2 };
18360 assert!(buckets >= 1 && buckets.is_power_of_two());
18361 assert_eq!(open_writers, buckets);
18362 assert!(workers >= 1 && workers <= threads);
18363 assert!(
18365 open + open_writers * writer_files + workers * 2 <= limit,
18366 "plan exceeds descriptor limit {limit}: {open_writers} writers, {workers} workers"
18367 );
18368 }
18369
18370 assert_eq!(
18372 materialized_coordinate_plan(1_048_576, 64, 64, true, SMALL_GRAPH),
18373 (
18374 DEFAULT_MAX_UNITIG_COORD_BUCKETS,
18375 64,
18376 DEFAULT_MAX_UNITIG_COORD_BUCKETS
18377 )
18378 );
18379 assert_eq!(
18381 materialized_coordinate_plan(1_048_576, 64, 256, true, SMALL_GRAPH),
18382 (
18383 HIGH_THREAD_MAX_UNITIG_COORD_BUCKETS,
18384 256,
18385 HIGH_THREAD_MAX_UNITIG_COORD_BUCKETS
18386 )
18387 );
18388 let large_graph =
18390 MAX_NARROW_COORD_BUCKET_BASES * HIGH_THREAD_MAX_UNITIG_COORD_BUCKETS as u64 * 2;
18391 assert_eq!(
18392 materialized_coordinate_plan(1_048_576, 64, 256, true, large_graph),
18393 (
18394 DEFAULT_MAX_UNITIG_COORD_BUCKETS,
18395 256,
18396 DEFAULT_MAX_UNITIG_COORD_BUCKETS
18397 )
18398 );
18399
18400 let (buckets, workers, open_writers) =
18403 materialized_coordinate_plan(96, 64, 32, true, SMALL_GRAPH);
18404 assert!(buckets >= 1 && buckets.is_power_of_two());
18405 assert_eq!(open_writers, buckets);
18406 assert!(workers >= 1);
18407 }
18408
18409 #[test]
18410 fn local_unitig_bucket_plan_accounts_for_concurrent_files() {
18411 assert_eq!(local_unitig_bucket_plan(1024, 264, 256), 124);
18414
18415 assert_eq!(
18417 local_unitig_bucket_plan(1_048_576, 264, 256),
18418 MAX_LOCAL_UNITIG_BUCKETS
18419 );
18420 assert_eq!(
18421 local_unitig_bucket_plan(65_536, 264, 256),
18422 MAX_LOCAL_UNITIG_BUCKETS
18423 );
18424 assert_eq!(
18426 local_unitig_bucket_plan(1_048_576, 264, 64),
18427 64 * LOCAL_UNITIG_BUCKETS_PER_WORKER
18428 );
18429
18430 for &(limit, open, workers) in &[
18431 (1024usize, 16usize, 128usize),
18432 (4096, 64, 64),
18433 (262_144, 64, 256),
18434 ] {
18435 let buckets = local_unitig_bucket_plan(limit, open, workers);
18436 if buckets >= workers {
18439 assert_eq!(
18440 buckets % workers,
18441 0,
18442 "{buckets} not a multiple of {workers}"
18443 );
18444 }
18445 assert!(buckets * 2 + workers * 2 + open <= limit);
18447 }
18448 }
18449
18450 #[test]
18451 fn stitched_path_info_record_uses_compact_external_layout() {
18452 let record = StitchedCoordRecord {
18453 path_id: 17,
18454 rank: 23,
18455 unitig_index: u32::MAX - 1,
18456 reverse: true,
18457 is_cycle: true,
18458 };
18459 let encoded = encoded_stitched_coord_record(record);
18460 assert_eq!(encoded.len(), 24);
18461 assert_eq!(
18462 decoded_stitched_coord_record(&encoded, Path::new("test.scb")).unwrap(),
18463 record
18464 );
18465 }
18466
18467 #[test]
18468 fn in_place_reverse_complement_matches_allocating_version() {
18469 for label in [b"".as_slice(), b"A", b"AC", b"ACG", b"AACGTT"] {
18470 let mut in_place = label.to_vec();
18471 reverse_complement_label_in_place(&mut in_place);
18472 assert_eq!(in_place, reverse_complement_label(label));
18473 }
18474 }
18475
18476 #[test]
18477 fn compact_path_info_table_survives_collisions_and_generation_wraps() {
18478 let table = CompactPathInfoTable::with_max_entries(6, 31);
18479 assert_eq!(std::mem::size_of::<CompactPathInfoSlot>(), 32);
18480 assert_eq!(table.slots.len(), 16);
18481
18482 let mut first_by_bucket = vec![None; table.slots.len()];
18483 let (first, second) = (0u64..)
18484 .find_map(|key| {
18485 let bucket = wyhash_u64(key, 0) as usize & table.mask;
18486 if let Some(first) = first_by_bucket[bucket] {
18487 Some((first, key))
18488 } else {
18489 first_by_bucket[bucket] = Some(key);
18490 None
18491 }
18492 })
18493 .unwrap();
18494
18495 for generation in 0..300u64 {
18496 table.clear();
18497 assert!(table.get::<31>(Kmer::from_bits(first as u128)).is_none());
18498 let first_info = PathInfo {
18499 path_id: Kmer::from_bits((100 + generation) as u128),
18500 rank: 10 + generation,
18501 exit_side: Side::Front,
18502 is_cycle: false,
18503 };
18504 let second_info = PathInfo {
18505 path_id: Kmer::from_bits((200 + generation) as u128),
18506 rank: 20 + generation,
18507 exit_side: Side::Back,
18508 is_cycle: true,
18509 };
18510 assert!(table.insert::<31>(Kmer::from_bits(first as u128), first_info));
18511 assert!(table.insert::<31>(Kmer::from_bits(second as u128), second_info));
18512 assert!(!table.insert::<31>(Kmer::from_bits(first as u128), first_info));
18513 assert_eq!(
18514 table.get::<31>(Kmer::from_bits(first as u128)),
18515 Some(first_info)
18516 );
18517 assert_eq!(
18518 table.get::<31>(Kmer::from_bits(second as u128)),
18519 Some(second_info)
18520 );
18521 }
18522 }
18523
18524 #[test]
18525 fn materialized_color_index_uses_thirty_bits() {
18526 for color_index in [0, 0xff_ffff, 0x100_0000, 0x3fff_fffe, u32::MAX] {
18527 let record = MaterializedStitchedCoordRecord {
18528 path_id: 17,
18529 rank: 23,
18530 label_offset: 41,
18531 label_len: 59,
18532 reverse: true,
18533 is_cycle: true,
18534 color_index,
18535 color_count: 61,
18536 };
18537 let encoded = encoded_materialized_stitched_coord_record(record);
18538 assert_eq!(encoded.len(), 24);
18539 let decoded = decoded_materialized_stitched_coord_record(
18540 &encoded,
18541 Path::new("materialized-record-test"),
18542 )
18543 .unwrap();
18544 assert_eq!(decoded, record);
18545 }
18546 }
18547
18548 #[test]
18549 fn appending_materialized_shards_rebases_labels_and_colors() {
18550 let directory = std::env::temp_dir().join(format!(
18551 "cf3-materialized-append-{}-{:?}",
18552 std::process::id(),
18553 std::thread::current().id()
18554 ));
18555 std::fs::create_dir_all(&directory).unwrap();
18556
18557 let make_entry = |worker_id, path_id, label: &[u8], color_offset| {
18558 let mut writer =
18559 MaterializedStitchedCoordShardWriter::create(&directory, worker_id, 7).unwrap();
18560 writer
18561 .write_colored_record(
18562 &StitchedCoordRecord {
18563 path_id,
18564 rank: worker_id as u64,
18565 unitig_index: 0,
18566 reverse: false,
18567 is_cycle: false,
18568 },
18569 label,
18570 &[UnitigColor::new(
18571 color_offset,
18572 crate::state::ColorCoordinate::from_u40(path_id),
18573 )],
18574 )
18575 .unwrap();
18576 writer.finish().unwrap()
18577 };
18578
18579 let first = make_entry(0, 11, b"ACGT", 3);
18580 let second = make_entry(1, 29, b"TTAA", 5);
18581 assert_eq!(
18582 std::fs::metadata(first.color_path.as_ref().unwrap())
18583 .unwrap()
18584 .len(),
18585 std::mem::size_of::<UnitigColor>() as u64
18586 );
18587 assert_eq!(
18588 std::fs::metadata(second.color_path.as_ref().unwrap())
18589 .unwrap()
18590 .len(),
18591 std::mem::size_of::<UnitigColor>() as u64
18592 );
18593 let mut bucket = MaterializedStitchedCoordBucket::default();
18594 append_materialized_stitched_coord_bucket_file(&first, &mut bucket).unwrap();
18595 append_materialized_stitched_coord_bucket_file(&second, &mut bucket).unwrap();
18596
18597 assert_eq!(bucket.records.len(), 2);
18598 assert_eq!(bucket.labels, b"ACGTTTAA");
18599 assert_eq!(bucket.records[0].label_offset, 0);
18600 assert_eq!(bucket.records[1].label_offset, 4);
18601 assert_eq!(bucket.records[0].color_start, 0);
18602 assert_eq!(bucket.records[1].color_start, 1);
18603 assert_eq!(bucket.records[0].color_count(), 1);
18604 assert_eq!(bucket.records[1].color_count(), 1);
18605 assert_eq!(
18606 bucket.colors[0].raw(),
18607 UnitigColor::new(3, crate::state::ColorCoordinate::from_u40(11),).raw()
18608 );
18609 assert_eq!(
18610 bucket.colors[1].raw(),
18611 UnitigColor::new(5, crate::state::ColorCoordinate::from_u40(29),).raw()
18612 );
18613
18614 std::fs::remove_dir_all(directory).unwrap();
18615 }
18616
18617 #[test]
18618 fn paged_external_range_index_matches_binary_search() {
18619 let mut ranges = Vec::new();
18620 let mut start = 0;
18621 for (index, unitigs) in [1, 17, 65_535, 2, 131_073, 4096, 70_000]
18622 .into_iter()
18623 .enumerate()
18624 {
18625 ranges.push(ExternalLocalUnitigRange {
18626 start_unitig: start,
18627 unitigs,
18628 label_start: index as u64,
18629 label_len: 0,
18630 color_start: 0,
18631 });
18632 start += unitigs;
18633 }
18634 let index = ExternalRangeIndex::new(&ranges);
18635 for unitig in 0..start {
18636 assert_eq!(
18637 index.find(&ranges, unitig),
18638 external_range_id_for_unitig(&ranges, unitig)
18639 );
18640 }
18641 assert_eq!(index.find(&ranges, start), None);
18642 }
18643}