Skip to main content

cuttlefish_rs/
buckets.rs

1//! External weak-super-k-mer bucket storage.
2//!
3//! Writers use the same 128-subgraph atlas hierarchy as Cuttlefish 3. Worker
4//! buffers are drained in source/worker order, and open files are bounded so a
5//! build does not require one descriptor per subgraph. The on-disk format is a
6//! private, versioned Rust format and may be compressed with LZ4 blocks.
7
8use crate::discontinuity::{current_open_file_count, open_file_limit};
9use crate::dna::{Base, ascii_base_bits, valid_ascii_base_bits};
10use crate::params::BuildParams;
11use crate::partition::WeakSuperKmer;
12use std::collections::BTreeMap;
13use std::fs::{self, File, OpenOptions};
14use std::io::{BufRead, BufReader, Read, Seek, SeekFrom, Write};
15use std::os::unix::fs::FileExt;
16use std::path::{Path, PathBuf};
17use std::sync::{
18    Arc, Mutex,
19    atomic::{AtomicU64, AtomicUsize, Ordering},
20};
21use std::time::{Duration, Instant};
22
23const MAGIC: &[u8; 8] = b"CF3WSK1\0";
24const RECORD_COUNT_OFFSET: u64 = 34;
25const HEADER_LEN: u64 = 42;
26const COMPRESSED_BLOCK_HEADER_LEN: usize = 12;
27const MAX_SOURCE_ID: u32 = 0x1f_ffff;
28const MAX_OPEN_BUCKET_WRITERS: usize = 512;
29/// Largest record the staging buffers emit: attribute plus four label words.
30const MAX_RECORD_BYTES: usize = 4 + 4 * 8;
31const MAX_PENDING_BUCKET_BYTES: usize = 1024 * 1024;
32// Keep a colored source window coalesced in worker-local atlas buffers. C++
33// retains roughly this amount per active worker set; the larger cap avoids
34// repeatedly scanning every graph bucket merely to move tiny fragments into
35// the shared atlas.
36const MAX_TOTAL_PENDING_BYTES: usize = 128 * 1024 * 1024;
37const ATLAS_GRAPH_COUNT: usize = 128;
38const SUBGRAPH_CHUNK_BYTES: usize = 64 * 1024;
39
40/// Bytes reserved per segment when buckets share container files.
41///
42/// The segment decouples the write unit from the read and reclaim units. A
43/// flush stays at `SUBGRAPH_CHUNK_BYTES` because that is a memory budget --
44/// 16,384 buckets times 64 KiB of staging -- while reads become segment-sized
45/// and reclaim becomes block-aligned, which is what lets a consumed bucket be
46/// punched out in full.
47///
48/// Sized from measurement on 149,998 Salmonella assemblies: 237.8 GB across
49/// 16,385 buckets, mean 14.51 MB, p1 8.80 MB and p99 22.77 MB, so the
50/// distribution is tight and the only real cost is the partial final segment
51/// each bucket leaves. At 256 KiB that is 2.15 GB, 0.90% of the directory,
52/// against 7.3 MB of chain metadata. Larger segments trade waste for fewer
53/// reads, and reads were measured not to matter on local storage.
54const DEFAULT_BUCKET_SEGMENT_BYTES: u64 = 256 * 1024;
55
56/// Descriptors left for everything downstream of the bucket containers: the
57/// edge-matrix containers, local-unitig buckets, stitch writers and the
58/// coordinate-bucket fanout, all of which plan against the same budget.
59const RESERVED_NON_BUCKET_DESCRIPTORS: usize = 384;
60
61fn bucket_segment_bytes() -> u64 {
62    static BYTES: std::sync::OnceLock<u64> = std::sync::OnceLock::new();
63    *BYTES.get_or_init(|| {
64        std::env::var("CF3_RS_BUCKET_SEGMENT_BYTES")
65            .ok()
66            .and_then(|value| value.parse::<u64>().ok())
67            .filter(|bytes| *bytes >= MAX_RECORD_BYTES as u64 && bytes % 4096 == 0)
68            .unwrap_or(DEFAULT_BUCKET_SEGMENT_BYTES)
69    })
70}
71
72/// The physical files backing the weak-super-k-mer buckets.
73///
74/// One container per atlas rather than one file per bucket, taking a 16,384
75/// bucket build from 16,385 files to 129. Two things make this cheap rather
76/// than intricate. Every bucket already writes under its atlas's mutex
77/// (`SharedBucketSink::append_bucket`), and a container holds exactly one
78/// atlas's buckets, so a container is only ever written by the thread holding
79/// that lock and needs no lock of its own. And a flush no longer opens
80/// anything: `BucketFile::open_existing` cost an `openat`, seven unbuffered
81/// reads to re-read a 42-byte header, a revalidation, an `lseek` and a `close`
82/// on *every* 64 KiB flush, which is about eleven syscalls times the 14.4
83/// million flushes a full-corpus build performs. A container flush is one
84/// `pwrite`.
85///
86/// The measured prize is smaller than that count suggests -- partitioning's
87/// whole system time is 436.7 s of CPU across 64 threads, so the ceiling is
88/// a couple of seconds of wall -- and larger somewhere unexpected. XFS
89/// speculatively preallocates on extending writes, and with 16,385 repeatedly
90/// reopened files it held 332.7 GB for 237.8 GB of data. Writing 128 files
91/// instead returns that 94.9 GB.
92#[derive(Debug)]
93pub struct BucketContainers {
94    files: Vec<BucketContainerFile>,
95    segment_bytes: u64,
96    /// Latched so an unsupported filesystem is reported once, not per bucket.
97    punch_unsupported: std::sync::atomic::AtomicBool,
98}
99
100#[derive(Debug)]
101struct BucketContainerFile {
102    path: PathBuf,
103    file: File,
104    /// Next unreserved byte offset. Atomic for shape rather than contention:
105    /// one atlas owns one container.
106    cursor: AtomicU64,
107}
108
109impl BucketContainers {
110    /// Containers a build may hold open, given the descriptor budget.
111    ///
112    /// One per atlas is the natural choice and what a normal limit allows, but
113    /// it must not be a floor: the per-file layout this replaced held no
114    /// descriptors between flushes, so a tight `ulimit -n` merely narrowed the
115    /// fanout planners rather than failing the build. Sharing a container
116    /// between atlases costs nothing -- the segment cursor is atomic, so
117    /// concurrent reservations from different atlas locks are already safe --
118    /// and keeps that property.
119    fn plan_container_count(atlas_count: usize) -> usize {
120        let budget = open_file_limit()
121            .saturating_sub(current_open_file_count())
122            // Local contraction opens the edge-matrix containers and the
123            // local-unitig buckets on top of these, and the fanout planners
124            // want room of their own.
125            .saturating_sub(RESERVED_NON_BUCKET_DESCRIPTORS)
126            / 2;
127        atlas_count.min(budget.max(1))
128    }
129
130    fn create(bucket_dir: &Path, container_count: usize) -> Result<Self, BucketError> {
131        let mut files = Vec::with_capacity(container_count);
132        for index in 0..container_count {
133            let path = bucket_dir.join(format!("{index:05}.wskc"));
134            let file = OpenOptions::new()
135                .create(true)
136                .truncate(true)
137                .read(true)
138                .write(true)
139                .open(&path)
140                .map_err(|source| BucketError::Io {
141                    path: path.clone(),
142                    source,
143                })?;
144            files.push(BucketContainerFile {
145                path,
146                file,
147                cursor: AtomicU64::new(0),
148            });
149        }
150        Ok(Self {
151            files,
152            segment_bytes: bucket_segment_bytes(),
153            punch_unsupported: std::sync::atomic::AtomicBool::new(false),
154        })
155    }
156
157    /// Opens the containers a finished manifest names, for reading.
158    fn open(
159        bucket_dir: &Path,
160        container_count: usize,
161        segment_bytes: u64,
162    ) -> Result<Self, BucketError> {
163        let mut files = Vec::with_capacity(container_count);
164        for index in 0..container_count {
165            let path = bucket_dir.join(format!("{index:05}.wskc"));
166            let file = OpenOptions::new()
167                .read(true)
168                .write(true)
169                .open(&path)
170                .map_err(|source| BucketError::Io {
171                    path: path.clone(),
172                    source,
173                })?;
174            files.push(BucketContainerFile {
175                path,
176                file,
177                cursor: AtomicU64::new(0),
178            });
179        }
180        Ok(Self {
181            files,
182            segment_bytes,
183            punch_unsupported: std::sync::atomic::AtomicBool::new(false),
184        })
185    }
186
187    #[inline]
188    pub fn segment_bytes(&self) -> u64 {
189        self.segment_bytes
190    }
191
192    #[inline]
193    pub fn len(&self) -> usize {
194        self.files.len()
195    }
196
197    #[inline]
198    pub fn is_empty(&self) -> bool {
199        self.files.is_empty()
200    }
201
202    #[inline]
203    fn reserve_segment(&self, container: usize) -> u64 {
204        self.files[container]
205            .cursor
206            .fetch_add(self.segment_bytes, Ordering::Relaxed)
207    }
208
209    fn write_at(&self, container: usize, offset: u64, bytes: &[u8]) -> Result<(), BucketError> {
210        let file = &self.files[container];
211        file.file
212            .write_all_at(bytes, offset)
213            .map_err(|source| BucketError::Io {
214                path: file.path.clone(),
215                source,
216            })
217    }
218
219    fn read_at(&self, container: usize, offset: u64, buf: &mut [u8]) -> Result<(), BucketError> {
220        let file = &self.files[container];
221        file.file
222            .read_exact_at(buf, offset)
223            .map_err(|source| BucketError::Io {
224                path: file.path.clone(),
225                source,
226            })
227    }
228
229    /// Releases a consumed bucket's segments without disturbing its neighbours.
230    ///
231    /// This is not optional, which measurement rather than reasoning settled.
232    /// The per-file layout unlinked each bucket as local contraction consumed
233    /// it, and containers cannot: a container is only droppable once all 128
234    /// of its buckets are done. Containers do start about 94 GB below the
235    /// per-file layout, because they do not accumulate XFS speculative
236    /// preallocation, and the expectation was that this covered it. It does
237    /// not -- the work-directory peak moves out of the end of partitioning and
238    /// into local contraction, where the containers still hold everything
239    /// while local-unitig buckets, labels and the edge matrix accumulate on
240    /// top, and peak disk rose 24.5 GB.
241    ///
242    /// Punching restores the incremental release. Segments are reserved at a
243    /// 4 KiB multiple, so a punch frees whole filesystem blocks rather than
244    /// leaving partial ones behind -- the one thing raw extents, averaging
245    /// 16.1 KiB and unaligned, could not have done. Adjacent segments are
246    /// punched in one call.
247    pub fn release_segments(&self, container: usize, segments: &[u32]) {
248        if segments.is_empty() {
249            return;
250        }
251        let mut ordered = segments.to_vec();
252        ordered.sort_unstable();
253        let file = &self.files[container];
254        let mut start = u64::from(ordered[0]) * self.segment_bytes;
255        let mut end = start + self.segment_bytes;
256        for &index in &ordered[1..] {
257            let offset = u64::from(index) * self.segment_bytes;
258            if offset == end {
259                end += self.segment_bytes;
260                continue;
261            }
262            self.report_punch(punch_hole(&file.file, start, end - start));
263            start = offset;
264            end = offset + self.segment_bytes;
265        }
266        self.report_punch(punch_hole(&file.file, start, end - start));
267    }
268
269    /// Says once if the filesystem will not punch holes.
270    ///
271    /// Reclaim failing is not an error -- the container is unlinked wholesale
272    /// at the end regardless -- but it silently costs peak disk, which is most
273    /// of what the container layout is for. HFS+ and some network filesystems
274    /// do not implement it, and macOS rejects a range that is not aligned to
275    /// the filesystem block size. Better to say so than to leave someone
276    /// wondering why the work directory is larger than documented.
277    fn report_punch(&self, punched: bool) {
278        if punched || self.punch_unsupported.swap(true, Ordering::Relaxed) {
279            return;
280        }
281        eprintln!(
282            "cuttlefish: this filesystem will not punch holes, so consumed \
283             bucket space is held until the build ends; peak disk will be higher"
284        );
285    }
286
287    pub fn paths(&self) -> impl Iterator<Item = &Path> {
288        self.files.iter().map(|file| file.path.as_path())
289    }
290}
291
292/// Punches `len` bytes at `offset` out of `file`.
293///
294/// There is no portable interface for this and no crate that abstracts one:
295/// Linux spells it `fallocate(FALLOC_FL_PUNCH_HOLE)` and macOS spells it
296/// `fcntl(F_PUNCHHOLE)`, and both `nix` and `rustix` gate their `fallocate`
297/// behind `target_os = "linux"` with no Apple equivalent offered. So the two
298/// are written out here.
299///
300/// Both interfaces require the range to be filesystem-block aligned -- macOS
301/// returns `EINVAL` for a punch that is not a multiple of the block size --
302/// which every call here satisfies by construction, because offsets and
303/// lengths are whole segments and `bucket_segment_bytes` admits only multiples
304/// of 4096. That is a second, independent reason buckets got segments rather
305/// than the raw 16.1 KiB extents a flush would otherwise produce.
306///
307/// Failure is ignored on purpose. A filesystem without hole punching -- an old
308/// HFS+ volume, or a network mount -- costs disk rather than correctness,
309/// because the container is unlinked wholesale at the end regardless.
310#[cfg(target_os = "linux")]
311fn punch_hole(file: &File, offset: u64, len: u64) -> bool {
312    use std::os::fd::AsRawFd;
313    // SAFETY: the descriptor is owned by `file` and outlives the call, and the
314    // kernel validates the range itself.
315    unsafe {
316        libc::fallocate(
317            file.as_raw_fd(),
318            libc::FALLOC_FL_KEEP_SIZE | libc::FALLOC_FL_PUNCH_HOLE,
319            offset as libc::off_t,
320            len as libc::off_t,
321        ) == 0
322    }
323}
324
325#[cfg(target_vendor = "apple")]
326fn punch_hole(file: &File, offset: u64, len: u64) -> bool {
327    use std::os::fd::AsRawFd;
328    let punch = libc::fpunchhole_t {
329        fp_flags: 0,
330        reserved: 0,
331        fp_offset: offset as libc::off_t,
332        fp_length: len as libc::off_t,
333    };
334    // SAFETY: as above; `punch` outlives the call and F_PUNCHHOLE reads it as
335    // a `*const fpunchhole_t`.
336    unsafe { libc::fcntl(file.as_raw_fd(), libc::F_PUNCHHOLE, &punch) == 0 }
337}
338
339#[cfg(not(any(target_os = "linux", target_vendor = "apple")))]
340fn punch_hole(_file: &File, _offset: u64, _len: u64) -> bool {
341    false
342}
343
344/// Sequential reader over one bucket's segment chain.
345///
346/// `BucketReader` uses its source purely as a `Read`, so presenting the chain
347/// this way leaves the whole decode path -- record iteration, the compressed
348/// block framing, the borrowed-record fast path -- exactly as it was for whole
349/// files. Reads stop at each segment boundary, and the caller's `BufReader`
350/// hides that.
351#[derive(Debug)]
352struct SegmentChainReader {
353    containers: Arc<BucketContainers>,
354    container: usize,
355    segments: Vec<u64>,
356    len: u64,
357    pos: u64,
358}
359
360impl Read for SegmentChainReader {
361    fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
362        if self.pos >= self.len || buf.is_empty() {
363            return Ok(0);
364        }
365        let segment_bytes = self.containers.segment_bytes;
366        let index = (self.pos / segment_bytes) as usize;
367        let within = self.pos % segment_bytes;
368        let room = (segment_bytes - within).min(self.len - self.pos);
369        let take = (buf.len() as u64).min(room) as usize;
370        let offset = self.segments[index] + within;
371        self.containers
372            .read_at(self.container, offset, &mut buf[..take])
373            .map_err(std::io::Error::other)?;
374        self.pos += take as u64;
375        Ok(take)
376    }
377}
378
379#[derive(Debug, Clone, PartialEq, Eq)]
380pub struct BucketEmitStats {
381    pub bucket_dir: PathBuf,
382    pub bucket_files: usize,
383    pub bytes_written: u64,
384}
385
386pub struct BucketEmitter {
387    bucket_dir: PathBuf,
388    k: u16,
389    minimizer_len: u16,
390    graph_count: usize,
391    colored: bool,
392    compress_buckets: bool,
393    label_words: usize,
394    files: BTreeMap<usize, BucketFileMeta>,
395    writers: BTreeMap<usize, BucketFile>,
396    pending: Vec<PendingBucket>,
397    pending_bytes: usize,
398    scratch: CompressionScratch,
399}
400
401pub struct SharedBucketSink {
402    bucket_dir: PathBuf,
403    containers: BucketContainers,
404    k: u16,
405    minimizer_len: u16,
406    graph_count: usize,
407    colored: bool,
408    compress_buckets: bool,
409    label_words: usize,
410    workers: usize,
411    atlases: Vec<Mutex<SharedBucketAtlas>>,
412    flush_calls: AtomicU64,
413    flush_nanos: AtomicU64,
414}
415
416pub struct SharedBucketEmitter {
417    sink: Arc<SharedBucketSink>,
418    pending: Vec<PendingBucket>,
419    uncolored_pending: Vec<PendingColoredAtlas>,
420    colored_pending: Vec<PendingColoredAtlas>,
421    pending_bytes: usize,
422    deferred_uncolored: bool,
423}
424
425#[derive(Debug, Clone, PartialEq, Eq)]
426struct BucketFileMeta {
427    path: PathBuf,
428    records: u64,
429    bytes_written: u64,
430}
431
432#[derive(Debug, Clone, PartialEq, Eq, Default)]
433struct PendingBucket {
434    records: u64,
435    bytes: Vec<u8>,
436}
437
438#[derive(Debug, Clone, PartialEq, Eq, Default)]
439struct PendingColoredAtlas {
440    graph_ids: Vec<u16>,
441    bytes: Vec<u8>,
442}
443
444#[derive(Default)]
445struct SharedBucketFileMeta {
446    buffer: Vec<u8>,
447    buffer_records: u64,
448    total_records: u64,
449    written_records: u64,
450    bytes_written: u64,
451    /// Segment indices this bucket owns, in write order.
452    segments: Vec<u32>,
453    /// Bytes used in the last segment; a flush that overruns it straddles into
454    /// a freshly reserved one, which the reader stitches back because it walks
455    /// the chain in order.
456    segment_used: u64,
457}
458
459struct SharedBucketAtlas {
460    first_graph_id: usize,
461    files: Vec<SharedBucketFileMeta>,
462    buffered_bytes: usize,
463    /// Shared by every flush through this atlas, which the atlas lock already
464    /// serializes, so it costs no contention and saves an allocation per block.
465    scratch: CompressionScratch,
466}
467
468#[derive(Default)]
469struct SharedBucketFlushStats {
470    calls: u64,
471}
472
473/// Where one bucket's bytes live.
474///
475/// Both forms are real: the shared production sink writes containers, and
476/// `BucketEmitter` -- the serial emitter the tests and the legacy path use --
477/// still writes one file per bucket with its header inline. Keeping the enum
478/// lets the reader serve both without duplicating the decode path.
479#[derive(Debug, Clone, PartialEq, Eq)]
480pub enum BucketLocation {
481    /// One whole file, carrying its own 42-byte header.
482    File(PathBuf),
483    /// A chain of segments inside a shared container. The header that a whole
484    /// file would carry lives in the manifest instead, so nothing has to be
485    /// re-read and revalidated per flush.
486    Container {
487        container: usize,
488        /// Segment indices in write order; byte offset is index * segment_bytes.
489        segments: Vec<u32>,
490        /// Logical length, which is the sum of the payload written into those
491        /// segments and is generally less than their reserved capacity.
492        bytes: u64,
493    },
494}
495
496#[derive(Debug, Clone, PartialEq, Eq)]
497/// Manifest entry for one weak-super-k-mer bucket.
498pub struct BucketManifestEntry {
499    pub graph_id: usize,
500    pub records: u64,
501    pub location: BucketLocation,
502}
503
504impl BucketManifestEntry {
505    /// Bytes this bucket occupies, for the longest-bucket-first scheduler.
506    ///
507    /// Containers make this free. The per-file layout had to `stat` all 16,384
508    /// buckets before local contraction could sort them.
509    pub fn stored_bytes(&self) -> Result<u64, BucketError> {
510        match &self.location {
511            BucketLocation::File(path) => {
512                fs::metadata(path)
513                    .map(|meta| meta.len())
514                    .map_err(|source| BucketError::Io {
515                        path: path.clone(),
516                        source,
517                    })
518            }
519            BucketLocation::Container { bytes, .. } => Ok(*bytes),
520        }
521    }
522
523    /// The whole-file path, when the bucket is one.
524    pub fn file_path(&self) -> Option<&Path> {
525        match &self.location {
526            BucketLocation::File(path) => Some(path.as_path()),
527            BucketLocation::Container { .. } => None,
528        }
529    }
530}
531
532#[derive(Debug, Clone, PartialEq, Eq)]
533/// Versioned parameters decoded from a bucket file header.
534pub struct BucketHeader {
535    pub k: u16,
536    pub minimizer_len: u16,
537    pub graph_count: usize,
538    pub graph_id: usize,
539    pub colored: bool,
540    pub compressed: bool,
541    pub interleaved_compression: bool,
542    pub label_words: usize,
543    pub records: u64,
544}
545
546#[derive(Debug, Clone, PartialEq, Eq, Default)]
547/// Decoded weak-super-k-mer record with an ASCII label.
548pub struct BucketRecord {
549    pub graph_id: usize,
550    pub len: usize,
551    pub source_id: Option<u32>,
552    pub left_discontinuous: bool,
553    pub right_discontinuous: bool,
554    pub label: Vec<u8>,
555}
556
557#[derive(Debug, Clone, PartialEq, Eq, Default)]
558/// Decoded weak-super-k-mer record retaining its packed two-bit label.
559pub struct BucketPackedRecord {
560    pub graph_id: usize,
561    pub len: usize,
562    pub source_id: Option<u32>,
563    pub left_discontinuous: bool,
564    pub right_discontinuous: bool,
565    pub words: Vec<u64>,
566}
567
568pub(crate) struct BorrowedBucketPackedRecord<'a> {
569    pub graph_id: usize,
570    pub len: usize,
571    pub source_id: Option<u32>,
572    pub left_discontinuous: bool,
573    pub right_discontinuous: bool,
574    pub words: &'a [u64],
575}
576
577/// Byte source behind a `BucketReader`.
578///
579/// The decode path treats its source as a plain sequential `Read`, so a
580/// segment chain slots in beside a whole file without any of the record
581/// iteration, block framing, or borrowed-record handling needing to know.
582enum BucketSource {
583    File(File),
584    Chain(SegmentChainReader),
585}
586
587impl Read for BucketSource {
588    #[inline]
589    fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
590        match self {
591            Self::File(file) => file.read(buf),
592            Self::Chain(chain) => chain.read(buf),
593        }
594    }
595}
596
597/// Opens buckets, whichever layout the directory uses.
598///
599/// A container directory keeps the shared header in its manifest, so this owns
600/// both and hands out readers; a whole-file directory carries the header in
601/// each bucket and this just opens paths. Threading one of these through the
602/// consumers replaces threading a path per bucket.
603pub struct BucketStore {
604    containers: Option<Arc<BucketContainers>>,
605    header: Option<ContainerManifestHeader>,
606}
607
608impl BucketStore {
609    /// Opens a bucket directory and reads its manifest.
610    pub fn open_dir(
611        bucket_dir: impl AsRef<Path>,
612    ) -> Result<(Self, Vec<BucketManifestEntry>), BucketError> {
613        let bucket_dir = bucket_dir.as_ref();
614        if let Some((header, entries)) = read_container_manifest(bucket_dir)? {
615            let containers =
616                BucketContainers::open(bucket_dir, header.container_count, header.segment_bytes)?;
617            return Ok((
618                Self {
619                    containers: Some(Arc::new(containers)),
620                    header: Some(header),
621                },
622                entries,
623            ));
624        }
625        let entries = read_manifest(bucket_dir)?;
626        Ok((
627            Self {
628                containers: None,
629                header: None,
630            },
631            entries,
632        ))
633    }
634
635    /// A store for a directory that is known to hold whole bucket files.
636    pub fn files_only() -> Self {
637        Self {
638            containers: None,
639            header: None,
640        }
641    }
642
643    pub fn containers(&self) -> Option<&Arc<BucketContainers>> {
644        self.containers.as_ref()
645    }
646
647    /// Opens one bucket for reading.
648    pub fn reader(&self, entry: &BucketManifestEntry) -> Result<BucketReader, BucketError> {
649        match &entry.location {
650            BucketLocation::File(path) => BucketReader::open(path),
651            BucketLocation::Container {
652                container,
653                segments,
654                bytes,
655            } => {
656                let (Some(containers), Some(header)) = (&self.containers, &self.header) else {
657                    return Err(BucketError::MalformedRecord);
658                };
659                BucketReader::open_chain(
660                    Arc::clone(containers),
661                    *container,
662                    segments.clone(),
663                    *bytes,
664                    header.bucket_header(entry.graph_id, entry.records),
665                )
666            }
667        }
668    }
669}
670
671/// Streaming reader for one versioned weak-super-k-mer bucket.
672pub struct BucketReader {
673    path: PathBuf,
674    file: BufReader<BucketSource>,
675    header: BucketHeader,
676    records_read: u64,
677    block_attrs: Vec<u8>,
678    block_labels: Vec<u8>,
679    block_interleaved: Vec<u8>,
680    compressed_block: Vec<u8>,
681    block_records: usize,
682    block_record: usize,
683}
684
685impl BucketReader {
686    pub fn open(path: impl AsRef<Path>) -> Result<Self, BucketError> {
687        let path = path.as_ref().to_path_buf();
688        let file = File::open(&path).map_err(|source| BucketError::Io {
689            path: path.clone(),
690            source,
691        })?;
692        let mut file = BufReader::with_capacity(1024 * 1024, BucketSource::File(file));
693        let header = read_header(&mut file, &path)?;
694
695        Ok(Self {
696            path,
697            file,
698            header,
699            records_read: 0,
700            block_attrs: Vec::new(),
701            block_labels: Vec::new(),
702            block_interleaved: Vec::new(),
703            compressed_block: Vec::new(),
704            block_records: 0,
705            block_record: 0,
706        })
707    }
708
709    /// Opens a bucket stored as a segment chain.
710    ///
711    /// There is no inline header to skip: a container's buckets keep theirs in
712    /// the manifest, so the chain starts at the first record byte and the
713    /// caller supplies the header it would otherwise have read.
714    fn open_chain(
715        containers: Arc<BucketContainers>,
716        container: usize,
717        segments: Vec<u32>,
718        bytes: u64,
719        header: BucketHeader,
720    ) -> Result<Self, BucketError> {
721        let segment_bytes = containers.segment_bytes();
722        let path = containers.files[container].path.clone();
723        let chain = SegmentChainReader {
724            containers,
725            container,
726            segments: segments
727                .iter()
728                .map(|s| u64::from(*s) * segment_bytes)
729                .collect(),
730            len: bytes,
731            pos: 0,
732        };
733        Ok(Self {
734            path,
735            file: BufReader::with_capacity(1024 * 1024, BucketSource::Chain(chain)),
736            header,
737            records_read: 0,
738            block_attrs: Vec::new(),
739            block_labels: Vec::new(),
740            block_interleaved: Vec::new(),
741            compressed_block: Vec::new(),
742            block_records: 0,
743            block_record: 0,
744        })
745    }
746
747    #[inline]
748    pub fn header(&self) -> &BucketHeader {
749        &self.header
750    }
751
752    pub fn next_record(&mut self) -> Result<Option<BucketRecord>, BucketError> {
753        let mut record = BucketRecord::default();
754        if self.next_record_into(&mut record)? {
755            Ok(Some(record))
756        } else {
757            Ok(None)
758        }
759    }
760
761    pub fn next_record_into(&mut self, record: &mut BucketRecord) -> Result<bool, BucketError> {
762        let mut packed = BucketPackedRecord::default();
763        if !self.next_packed_record_into(&mut packed)? {
764            return Ok(false);
765        }
766
767        record.graph_id = packed.graph_id;
768        record.len = packed.len;
769        record.source_id = packed.source_id;
770        record.left_discontinuous = packed.left_discontinuous;
771        record.right_discontinuous = packed.right_discontinuous;
772        decode_label_into(&packed.words, packed.len, &mut record.label)?;
773        Ok(true)
774    }
775
776    pub fn next_packed_record_into(
777        &mut self,
778        record: &mut BucketPackedRecord,
779    ) -> Result<bool, BucketError> {
780        if self.records_read == self.header.records {
781            return Ok(false);
782        }
783
784        let mut fixed = [0u8; 4];
785        let fixed_len = if self.header.colored { 4 } else { 2 };
786        if self.header.compressed {
787            if self.block_record == self.block_records {
788                self.read_compressed_block()?;
789            }
790            let start = self.block_record * fixed_len;
791            if self.header.interleaved_compression {
792                let record_len = record_size(self.header.colored, self.header.label_words);
793                let start = self.block_record * record_len;
794                fixed[..fixed_len]
795                    .copy_from_slice(&self.block_interleaved[start..start + fixed_len]);
796            } else {
797                fixed[..fixed_len].copy_from_slice(&self.block_attrs[start..start + fixed_len]);
798            }
799        } else {
800            self.file
801                .read_exact(&mut fixed[..fixed_len])
802                .map_err(|source| BucketError::Io {
803                    path: self.path.clone(),
804                    source,
805                })?;
806        }
807        let packed_attr = if self.header.colored {
808            u32::from_le_bytes(fixed[0..4].try_into().unwrap())
809        } else {
810            u16::from_le_bytes(fixed[0..2].try_into().unwrap()) as u32
811        };
812
813        record.words.clear();
814        record.words.resize(self.header.label_words, 0);
815        for (word_idx, word) in record.words.iter_mut().enumerate() {
816            if self.header.compressed {
817                if self.header.interleaved_compression {
818                    let record_len = record_size(self.header.colored, self.header.label_words);
819                    let start = self.block_record * record_len + fixed_len + word_idx * 8;
820                    *word = u64::from_le_bytes(
821                        self.block_interleaved[start..start + 8].try_into().unwrap(),
822                    );
823                } else {
824                    let start = (self.block_record * self.header.label_words + word_idx) * 8;
825                    *word =
826                        u64::from_le_bytes(self.block_labels[start..start + 8].try_into().unwrap());
827                }
828            } else {
829                *word = read_u64(&mut self.file, &self.path)?;
830            }
831        }
832
833        let len = (packed_attr & 0xff) as usize;
834        let source_id = self
835            .header
836            .colored
837            .then_some((packed_attr >> 10) & MAX_SOURCE_ID);
838        record.graph_id = self.header.graph_id;
839        record.len = len;
840        record.source_id = source_id;
841        record.left_discontinuous = (packed_attr & (1 << 8)) != 0;
842        record.right_discontinuous = (packed_attr & (1 << 9)) != 0;
843
844        self.records_read += 1;
845        if self.header.compressed {
846            self.block_record += 1;
847        }
848        Ok(true)
849    }
850
851    pub fn try_for_each_packed_record<E, F>(
852        &mut self,
853        record: &mut BucketPackedRecord,
854        mut f: F,
855    ) -> Result<(), E>
856    where
857        E: From<BucketError>,
858        F: FnMut(&BucketPackedRecord) -> Result<(), E>,
859    {
860        let remaining = self.header.records.saturating_sub(self.records_read);
861        if remaining == 0 {
862            return Ok(());
863        }
864        if self.header.compressed {
865            while self.next_packed_record_into(record).map_err(E::from)? {
866                f(record)?;
867            }
868            return Ok(());
869        }
870        let record_size = record_size(self.header.colored, self.header.label_words);
871        let payload_bytes = remaining
872            .checked_mul(record_size as u64)
873            .ok_or(BucketError::TooManyRecords)?;
874        let mut payload = vec![0u8; payload_bytes as usize];
875        self.file
876            .read_exact(&mut payload)
877            .map_err(|source| BucketError::Io {
878                path: self.path.clone(),
879                source,
880            })?;
881
882        for chunk in payload.chunks_exact(record_size) {
883            let (packed_attr, words_start) = if self.header.colored {
884                (u32::from_le_bytes(chunk[0..4].try_into().unwrap()), 4)
885            } else {
886                (
887                    u16::from_le_bytes(chunk[0..2].try_into().unwrap()) as u32,
888                    2,
889                )
890            };
891
892            record.words.clear();
893            record.words.reserve(self.header.label_words);
894            for word_idx in 0..self.header.label_words {
895                let start = words_start + word_idx * 8;
896                record.words.push(u64::from_le_bytes(
897                    chunk[start..start + 8].try_into().unwrap(),
898                ));
899            }
900
901            let len = (packed_attr & 0xff) as usize;
902            let source_id = self
903                .header
904                .colored
905                .then_some((packed_attr >> 10) & MAX_SOURCE_ID);
906            record.graph_id = self.header.graph_id;
907            record.len = len;
908            record.source_id = source_id;
909            record.left_discontinuous = (packed_attr & (1 << 8)) != 0;
910            record.right_discontinuous = (packed_attr & (1 << 9)) != 0;
911
912            self.records_read += 1;
913            f(record)?;
914        }
915        Ok(())
916    }
917
918    pub(crate) fn try_for_each_borrowed_packed_record<E, F>(&mut self, mut f: F) -> Result<(), E>
919    where
920        E: From<BucketError>,
921        F: FnMut(BorrowedBucketPackedRecord<'_>) -> Result<(), E>,
922    {
923        if !self.header.compressed {
924            let mut record = BucketPackedRecord::default();
925            return self.try_for_each_packed_record(&mut record, |record| {
926                f(BorrowedBucketPackedRecord {
927                    graph_id: record.graph_id,
928                    len: record.len,
929                    source_id: record.source_id,
930                    left_discontinuous: record.left_discontinuous,
931                    right_discontinuous: record.right_discontinuous,
932                    words: &record.words,
933                })
934            });
935        }
936        if self.header.label_words > 4 {
937            return Err(E::from(BucketError::MalformedRecord));
938        }
939
940        let fixed_len = if self.header.colored { 4 } else { 2 };
941        while self.records_read < self.header.records {
942            if self.block_record == self.block_records {
943                self.read_compressed_block().map_err(E::from)?;
944            }
945            let record_index = self.block_record;
946            let (packed_attr, words_bytes) = if self.header.interleaved_compression {
947                let record_len = record_size(self.header.colored, self.header.label_words);
948                let start = record_index * record_len;
949                let attr = if self.header.colored {
950                    u32::from_le_bytes(self.block_interleaved[start..start + 4].try_into().unwrap())
951                } else {
952                    u16::from_le_bytes(self.block_interleaved[start..start + 2].try_into().unwrap())
953                        as u32
954                };
955                (
956                    attr,
957                    &self.block_interleaved[start + fixed_len..start + record_len],
958                )
959            } else {
960                let attr_start = record_index * fixed_len;
961                let attr = if self.header.colored {
962                    u32::from_le_bytes(
963                        self.block_attrs[attr_start..attr_start + 4]
964                            .try_into()
965                            .unwrap(),
966                    )
967                } else {
968                    u16::from_le_bytes(
969                        self.block_attrs[attr_start..attr_start + 2]
970                            .try_into()
971                            .unwrap(),
972                    ) as u32
973                };
974                let words_start = record_index * self.header.label_words * 8;
975                (
976                    attr,
977                    &self.block_labels[words_start..words_start + self.header.label_words * 8],
978                )
979            };
980            let mut words = [0u64; 4];
981            for (word, bytes) in words[..self.header.label_words]
982                .iter_mut()
983                .zip(words_bytes.chunks_exact(8))
984            {
985                *word = u64::from_le_bytes(bytes.try_into().unwrap());
986            }
987            self.records_read += 1;
988            self.block_record += 1;
989            f(BorrowedBucketPackedRecord {
990                graph_id: self.header.graph_id,
991                len: (packed_attr & 0xff) as usize,
992                source_id: self
993                    .header
994                    .colored
995                    .then_some((packed_attr >> 10) & MAX_SOURCE_ID),
996                left_discontinuous: packed_attr & (1 << 8) != 0,
997                right_discontinuous: packed_attr & (1 << 9) != 0,
998                words: &words[..self.header.label_words],
999            })?;
1000        }
1001        Ok(())
1002    }
1003
1004    pub fn records(self) -> BucketRecords {
1005        BucketRecords { reader: self }
1006    }
1007
1008    fn read_compressed_block(&mut self) -> Result<(), BucketError> {
1009        let mut header = [0u8; COMPRESSED_BLOCK_HEADER_LEN];
1010        self.file
1011            .read_exact(&mut header)
1012            .map_err(|source| BucketError::Io {
1013                path: self.path.clone(),
1014                source,
1015            })?;
1016        let records = u32::from_le_bytes(header[0..4].try_into().unwrap()) as usize;
1017        let attr_bytes = u32::from_le_bytes(header[4..8].try_into().unwrap()) as usize;
1018        let label_bytes = u32::from_le_bytes(header[8..12].try_into().unwrap()) as usize;
1019        if records == 0
1020            || attr_bytes == 0
1021            || (!self.header.interleaved_compression && label_bytes == 0)
1022            || (self.header.interleaved_compression && label_bytes != 0)
1023        {
1024            return Err(BucketError::MalformedRecord);
1025        }
1026        let remaining = usize::try_from(self.header.records - self.records_read)
1027            .map_err(|_| BucketError::TooManyRecords)?;
1028        if records > remaining {
1029            return Err(BucketError::MalformedRecord);
1030        }
1031        self.compressed_block.resize(attr_bytes + label_bytes, 0);
1032        self.file
1033            .read_exact(&mut self.compressed_block)
1034            .map_err(|source| BucketError::Io {
1035                path: self.path.clone(),
1036                source,
1037            })?;
1038        let fixed_len = if self.header.colored { 4 } else { 2 };
1039        if self.header.interleaved_compression {
1040            self.block_interleaved.resize(
1041                records * record_size(self.header.colored, self.header.label_words),
1042                0,
1043            );
1044            let decoded = lz4_flex::block::decompress_into(
1045                &self.compressed_block[..attr_bytes],
1046                &mut self.block_interleaved,
1047            )
1048            .map_err(|_| BucketError::MalformedRecord)?;
1049            if decoded != self.block_interleaved.len() {
1050                return Err(BucketError::MalformedRecord);
1051            }
1052            self.block_records = records;
1053            self.block_record = 0;
1054            return Ok(());
1055        }
1056        self.block_attrs.resize(records * fixed_len, 0);
1057        self.block_labels
1058            .resize(records * self.header.label_words * 8, 0);
1059        let decoded_attrs = lz4_flex::block::decompress_into(
1060            &self.compressed_block[..attr_bytes],
1061            &mut self.block_attrs,
1062        )
1063        .map_err(|_| BucketError::MalformedRecord)?;
1064        let decoded_labels = lz4_flex::block::decompress_into(
1065            &self.compressed_block[attr_bytes..],
1066            &mut self.block_labels,
1067        )
1068        .map_err(|_| BucketError::MalformedRecord)?;
1069        if decoded_attrs != self.block_attrs.len() || decoded_labels != self.block_labels.len() {
1070            return Err(BucketError::MalformedRecord);
1071        }
1072        self.block_records = records;
1073        self.block_record = 0;
1074        Ok(())
1075    }
1076}
1077
1078/// Iterator over decoded records from a [`BucketReader`].
1079pub struct BucketRecords {
1080    reader: BucketReader,
1081}
1082
1083impl Iterator for BucketRecords {
1084    type Item = Result<BucketRecord, BucketError>;
1085
1086    fn next(&mut self) -> Option<Self::Item> {
1087        self.reader.next_record().transpose()
1088    }
1089}
1090
1091const CONTAINER_MANIFEST_MAGIC: &[u8; 8] = b"CF3WSKC1";
1092const CONTAINER_MANIFEST_NAME: &str = "manifest.bin";
1093
1094/// The parameters every bucket in a container directory shares.
1095///
1096/// This is the 42-byte per-bucket header, hoisted. Under the per-file layout it
1097/// was written once per bucket and then re-read and revalidated on *every*
1098/// 64 KiB flush, because a flush reopened the file to append. Stored once for
1099/// the whole directory, that disappears; the only genuinely per-bucket field
1100/// was the record count, which the manifest already carried.
1101#[derive(Debug, Clone, PartialEq, Eq)]
1102pub struct ContainerManifestHeader {
1103    pub k: u16,
1104    pub minimizer_len: u16,
1105    pub graph_count: usize,
1106    pub colored: bool,
1107    pub label_words: usize,
1108    pub compressed: bool,
1109    pub interleaved_compression: bool,
1110    pub segment_bytes: u64,
1111    pub container_count: usize,
1112}
1113
1114impl ContainerManifestHeader {
1115    fn compression_code(&self) -> u8 {
1116        if self.interleaved_compression {
1117            2
1118        } else {
1119            u8::from(self.compressed)
1120        }
1121    }
1122
1123    /// The per-bucket header a whole-file reader would have found inline.
1124    fn bucket_header(&self, graph_id: usize, records: u64) -> BucketHeader {
1125        BucketHeader {
1126            k: self.k,
1127            minimizer_len: self.minimizer_len,
1128            graph_count: self.graph_count,
1129            graph_id,
1130            colored: self.colored,
1131            label_words: self.label_words,
1132            compressed: self.compressed,
1133            interleaved_compression: self.interleaved_compression,
1134            records,
1135        }
1136    }
1137}
1138
1139fn write_u32_to(out: &mut Vec<u8>, value: u32) {
1140    out.extend_from_slice(&value.to_le_bytes());
1141}
1142
1143fn write_u64_to(out: &mut Vec<u8>, value: u64) {
1144    out.extend_from_slice(&value.to_le_bytes());
1145}
1146
1147/// Writes the container manifest: shared header, then one record per bucket.
1148pub fn write_container_manifest(
1149    bucket_dir: &Path,
1150    header: &ContainerManifestHeader,
1151    entries: &[BucketManifestEntry],
1152) -> Result<(), BucketError> {
1153    let path = bucket_dir.join(CONTAINER_MANIFEST_NAME);
1154    let mut out = Vec::with_capacity(64 + entries.len() * 32);
1155    out.extend_from_slice(CONTAINER_MANIFEST_MAGIC);
1156    out.extend_from_slice(&header.k.to_le_bytes());
1157    out.extend_from_slice(&header.minimizer_len.to_le_bytes());
1158    write_u64_to(&mut out, header.graph_count as u64);
1159    out.push(u8::from(header.colored));
1160    out.push(header.label_words as u8);
1161    out.push(header.compression_code());
1162    out.push(0);
1163    write_u64_to(&mut out, header.segment_bytes);
1164    write_u64_to(&mut out, header.container_count as u64);
1165    write_u64_to(&mut out, entries.len() as u64);
1166    for entry in entries {
1167        let BucketLocation::Container {
1168            container,
1169            segments,
1170            bytes,
1171        } = &entry.location
1172        else {
1173            return Err(BucketError::MalformedManifest(path.clone()));
1174        };
1175        write_u64_to(&mut out, entry.graph_id as u64);
1176        write_u64_to(&mut out, entry.records);
1177        write_u64_to(&mut out, *bytes);
1178        write_u32_to(&mut out, *container as u32);
1179        write_u32_to(&mut out, segments.len() as u32);
1180        for segment in segments {
1181            write_u32_to(&mut out, *segment);
1182        }
1183    }
1184    fs::write(&path, &out).map_err(|source| BucketError::Io {
1185        path: path.clone(),
1186        source,
1187    })
1188}
1189
1190struct ManifestCursor<'a> {
1191    bytes: &'a [u8],
1192    pos: usize,
1193    path: &'a Path,
1194}
1195
1196impl<'a> ManifestCursor<'a> {
1197    fn take(&mut self, len: usize) -> Result<&'a [u8], BucketError> {
1198        let end = self
1199            .pos
1200            .checked_add(len)
1201            .filter(|end| *end <= self.bytes.len())
1202            .ok_or_else(|| BucketError::MalformedManifest(self.path.to_path_buf()))?;
1203        let slice = &self.bytes[self.pos..end];
1204        self.pos = end;
1205        Ok(slice)
1206    }
1207
1208    fn u16(&mut self) -> Result<u16, BucketError> {
1209        Ok(u16::from_le_bytes(self.take(2)?.try_into().unwrap()))
1210    }
1211
1212    fn u32(&mut self) -> Result<u32, BucketError> {
1213        Ok(u32::from_le_bytes(self.take(4)?.try_into().unwrap()))
1214    }
1215
1216    fn u64(&mut self) -> Result<u64, BucketError> {
1217        Ok(u64::from_le_bytes(self.take(8)?.try_into().unwrap()))
1218    }
1219}
1220
1221/// Reads the container manifest, if this directory has one.
1222pub fn read_container_manifest(
1223    bucket_dir: impl AsRef<Path>,
1224) -> Result<Option<(ContainerManifestHeader, Vec<BucketManifestEntry>)>, BucketError> {
1225    let path = bucket_dir.as_ref().join(CONTAINER_MANIFEST_NAME);
1226    let bytes = match fs::read(&path) {
1227        Ok(bytes) => bytes,
1228        Err(source) if source.kind() == std::io::ErrorKind::NotFound => return Ok(None),
1229        Err(source) => {
1230            return Err(BucketError::Io {
1231                path: path.clone(),
1232                source,
1233            });
1234        }
1235    };
1236    let mut cursor = ManifestCursor {
1237        bytes: &bytes,
1238        pos: 0,
1239        path: &path,
1240    };
1241    if cursor.take(8)? != CONTAINER_MANIFEST_MAGIC {
1242        return Err(BucketError::MalformedManifest(path.clone()));
1243    }
1244    let k = cursor.u16()?;
1245    let minimizer_len = cursor.u16()?;
1246    let graph_count = cursor.u64()? as usize;
1247    let flags = cursor.take(4)?;
1248    let (colored, label_words, compression) = (flags[0], flags[1], flags[2]);
1249    if colored > 1 || compression > 2 || flags[3] != 0 {
1250        return Err(BucketError::MalformedManifest(path.clone()));
1251    }
1252    let header = ContainerManifestHeader {
1253        k,
1254        minimizer_len,
1255        graph_count,
1256        colored: colored == 1,
1257        label_words: label_words as usize,
1258        compressed: compression != 0,
1259        interleaved_compression: compression == 2,
1260        segment_bytes: cursor.u64()?,
1261        container_count: cursor.u64()? as usize,
1262    };
1263    if header.segment_bytes == 0
1264        || header.label_words as usize != label_word_count(k, minimizer_len)
1265    {
1266        return Err(BucketError::MalformedManifest(path.clone()));
1267    }
1268    let bucket_count = cursor.u64()? as usize;
1269    let mut entries = Vec::with_capacity(bucket_count);
1270    for _ in 0..bucket_count {
1271        let graph_id = cursor.u64()? as usize;
1272        let records = cursor.u64()?;
1273        let bytes_len = cursor.u64()?;
1274        let container = cursor.u32()? as usize;
1275        let segment_count = cursor.u32()? as usize;
1276        if graph_id >= graph_count || container >= header.container_count {
1277            return Err(BucketError::MalformedManifest(path.clone()));
1278        }
1279        let mut segments = Vec::with_capacity(segment_count);
1280        for _ in 0..segment_count {
1281            segments.push(cursor.u32()?);
1282        }
1283        // The chain must be able to hold the logical length, and must not be
1284        // more than one segment longer than it needs to be.
1285        let capacity = segment_count as u64 * header.segment_bytes;
1286        if bytes_len > capacity || capacity.saturating_sub(bytes_len) >= header.segment_bytes {
1287            return Err(BucketError::MalformedManifest(path.clone()));
1288        }
1289        entries.push(BucketManifestEntry {
1290            graph_id,
1291            records,
1292            location: BucketLocation::Container {
1293                container,
1294                segments,
1295                bytes: bytes_len,
1296            },
1297        });
1298    }
1299    Ok(Some((header, entries)))
1300}
1301
1302/// Reads and validates `manifest.tsv` from a bucket directory.
1303pub fn read_manifest(
1304    bucket_dir: impl AsRef<Path>,
1305) -> Result<Vec<BucketManifestEntry>, BucketError> {
1306    let bucket_dir = bucket_dir.as_ref();
1307    let path = bucket_dir.join("manifest.tsv");
1308    let file = File::open(&path).map_err(|source| BucketError::Io {
1309        path: path.clone(),
1310        source,
1311    })?;
1312    let mut out = Vec::new();
1313
1314    for (line_no, line) in BufReader::new(file).lines().enumerate() {
1315        let line = line.map_err(|source| BucketError::Io {
1316            path: path.clone(),
1317            source,
1318        })?;
1319        if line_no == 0 {
1320            if line != "graph_id\trecords\tpath" {
1321                return Err(BucketError::MalformedManifest(path.clone()));
1322            }
1323            continue;
1324        }
1325        if line.trim().is_empty() {
1326            continue;
1327        }
1328
1329        let mut fields = line.splitn(3, '\t');
1330        let graph_id = fields
1331            .next()
1332            .ok_or_else(|| BucketError::MalformedManifest(path.clone()))?
1333            .parse()
1334            .map_err(|_| BucketError::MalformedManifest(path.clone()))?;
1335        let records = fields
1336            .next()
1337            .ok_or_else(|| BucketError::MalformedManifest(path.clone()))?
1338            .parse()
1339            .map_err(|_| BucketError::MalformedManifest(path.clone()))?;
1340        let bucket_path = fields
1341            .next()
1342            .ok_or_else(|| BucketError::MalformedManifest(path.clone()))?;
1343        out.push(BucketManifestEntry {
1344            graph_id,
1345            records,
1346            location: BucketLocation::File(PathBuf::from(bucket_path)),
1347        });
1348    }
1349
1350    Ok(out)
1351}
1352
1353pub fn coalesce_bucket_manifest(
1354    bucket_dir: &Path,
1355    entries: &[BucketManifestEntry],
1356) -> Result<BucketEmitStats, BucketError> {
1357    coalesce_bucket_manifest_with_threads(bucket_dir, entries, 1)
1358}
1359
1360pub fn coalesce_bucket_manifest_with_threads(
1361    bucket_dir: &Path,
1362    entries: &[BucketManifestEntry],
1363    threads: usize,
1364) -> Result<BucketEmitStats, BucketError> {
1365    let mut by_graph = BTreeMap::<usize, Vec<BucketManifestEntry>>::new();
1366    for entry in entries {
1367        by_graph
1368            .entry(entry.graph_id)
1369            .or_default()
1370            .push(entry.clone());
1371    }
1372
1373    let groups = by_graph.into_iter().collect::<Vec<_>>();
1374    let workers = threads.max(1).min(groups.len().max(1));
1375    let mut manifest = if workers == 1 {
1376        let mut manifest = Vec::with_capacity(groups.len());
1377        for group in &groups {
1378            manifest.push(coalesce_bucket_group(bucket_dir, group)?);
1379        }
1380        manifest
1381    } else {
1382        let next_group = AtomicUsize::new(0);
1383        let mut manifest = std::thread::scope(|scope| {
1384            let mut handles = Vec::new();
1385            for _ in 0..workers {
1386                let next_group = &next_group;
1387                let groups = &groups;
1388                handles.push(scope.spawn(move || {
1389                    let mut local = Vec::new();
1390                    loop {
1391                        let group_idx = next_group.fetch_add(1, Ordering::Relaxed);
1392                        let Some(group) = groups.get(group_idx) else {
1393                            break;
1394                        };
1395                        local.push(coalesce_bucket_group(bucket_dir, group)?);
1396                    }
1397                    Ok::<_, BucketError>(local)
1398                }));
1399            }
1400
1401            let mut manifest = Vec::with_capacity(groups.len());
1402            for handle in handles {
1403                manifest.extend(handle.join().map_err(|_| BucketError::WorkerPanic)??);
1404            }
1405            Ok::<_, BucketError>(manifest)
1406        })?;
1407        manifest.sort_by_key(|(graph_id, _, path, _)| (*graph_id, path.clone()));
1408        manifest
1409    };
1410    let total_bytes = manifest.iter().map(|(_, _, _, bytes)| *bytes).sum();
1411    let public_manifest = manifest
1412        .drain(..)
1413        .map(|(graph_id, records, path, _)| (graph_id, records, path))
1414        .collect::<Vec<_>>();
1415
1416    write_manifest(bucket_dir, &public_manifest)?;
1417
1418    Ok(BucketEmitStats {
1419        bucket_dir: bucket_dir.to_path_buf(),
1420        bucket_files: public_manifest.len(),
1421        bytes_written: total_bytes,
1422    })
1423}
1424
1425fn coalesce_bucket_group(
1426    bucket_dir: &Path,
1427    group: &(usize, Vec<BucketManifestEntry>),
1428) -> Result<(usize, u64, PathBuf, u64), BucketError> {
1429    let (graph_id, graph_entries) = group;
1430    // Coalescing concatenates raw payloads, which only whole uncompressed
1431    // bucket files expose; a container's buckets are read through the store.
1432    let entry_path = |entry: &BucketManifestEntry| {
1433        entry
1434            .file_path()
1435            .map(Path::to_path_buf)
1436            .ok_or_else(|| BucketError::MalformedManifest(bucket_dir.to_path_buf()))
1437    };
1438    let first_header = read_bucket_header(&entry_path(&graph_entries[0])?)?;
1439    let mut writer = BucketFile::create(
1440        bucket_dir,
1441        first_header.k,
1442        first_header.minimizer_len,
1443        first_header.graph_count,
1444        *graph_id,
1445        first_header.colored,
1446        first_header.label_words,
1447        first_header.compressed,
1448    )?;
1449
1450    for entry in graph_entries {
1451        let path = entry_path(entry)?;
1452        let copied_records = copy_bucket_payload(&path, &first_header, *graph_id, &mut writer)?;
1453        if copied_records != entry.records {
1454            return Err(BucketError::MalformedHeader(path));
1455        }
1456    }
1457
1458    writer.finish()?;
1459    Ok((
1460        *graph_id,
1461        writer.records,
1462        writer.path.clone(),
1463        writer.bytes_written,
1464    ))
1465}
1466
1467fn read_bucket_header(path: &Path) -> Result<BucketHeader, BucketError> {
1468    let mut file = File::open(path).map_err(|source| BucketError::Io {
1469        path: path.to_path_buf(),
1470        source,
1471    })?;
1472    read_header(&mut file, path)
1473}
1474
1475fn copy_bucket_payload(
1476    path: &Path,
1477    expected: &BucketHeader,
1478    graph_id: usize,
1479    writer: &mut BucketFile,
1480) -> Result<u64, BucketError> {
1481    let mut file = File::open(path).map_err(|source| BucketError::Io {
1482        path: path.to_path_buf(),
1483        source,
1484    })?;
1485    let header = read_header(&mut file, path)?;
1486    if header.k != expected.k
1487        || header.minimizer_len != expected.minimizer_len
1488        || header.graph_count != expected.graph_count
1489        || header.graph_id != graph_id
1490        || header.colored != expected.colored
1491        || header.label_words != expected.label_words
1492    {
1493        return Err(BucketError::MalformedHeader(path.to_path_buf()));
1494    }
1495
1496    let payload_bytes = header
1497        .records
1498        .checked_mul(record_size(header.colored, header.label_words) as u64)
1499        .ok_or(BucketError::TooManyRecords)?;
1500    let actual_len = file
1501        .metadata()
1502        .map_err(|source| BucketError::Io {
1503            path: path.to_path_buf(),
1504            source,
1505        })?
1506        .len();
1507    if actual_len != HEADER_LEN + payload_bytes {
1508        return Err(BucketError::MalformedHeader(path.to_path_buf()));
1509    }
1510
1511    file.seek(SeekFrom::Start(HEADER_LEN))
1512        .map_err(|source| BucketError::Io {
1513            path: path.to_path_buf(),
1514            source,
1515        })?;
1516    let copied =
1517        std::io::copy(&mut file.take(payload_bytes), &mut writer.file).map_err(|source| {
1518            BucketError::Io {
1519                path: path.to_path_buf(),
1520                source,
1521            }
1522        })?;
1523    if copied != payload_bytes {
1524        return Err(BucketError::MalformedHeader(path.to_path_buf()));
1525    }
1526    writer.records = writer
1527        .records
1528        .checked_add(header.records)
1529        .ok_or(BucketError::TooManyRecords)?;
1530    writer.bytes_written += copied;
1531    Ok(header.records)
1532}
1533
1534impl BucketEmitter {
1535    pub fn create(params: &BuildParams, graph_count: usize) -> Result<Self, BucketError> {
1536        Self::create_in_dir(params, graph_count, bucket_dir(params))
1537    }
1538
1539    pub fn create_in_dir(
1540        params: &BuildParams,
1541        graph_count: usize,
1542        bucket_dir: PathBuf,
1543    ) -> Result<Self, BucketError> {
1544        if graph_count > u64::MAX as usize {
1545            return Err(BucketError::GraphCountTooLarge(graph_count));
1546        }
1547
1548        if bucket_dir.exists() {
1549            fs::remove_dir_all(&bucket_dir).map_err(|source| BucketError::Io {
1550                path: bucket_dir.clone(),
1551                source,
1552            })?;
1553        }
1554        fs::create_dir_all(&bucket_dir).map_err(|source| BucketError::Io {
1555            path: bucket_dir.clone(),
1556            source,
1557        })?;
1558
1559        Ok(Self {
1560            bucket_dir,
1561            k: params.k,
1562            minimizer_len: params.minimizer_len,
1563            graph_count,
1564            colored: params.color,
1565            compress_buckets: params.color || params.compress_buckets,
1566            label_words: label_word_count(params.k, params.minimizer_len),
1567            files: BTreeMap::new(),
1568            scratch: CompressionScratch::default(),
1569            writers: BTreeMap::new(),
1570            pending: vec![PendingBucket::default(); graph_count],
1571            pending_bytes: 0,
1572        })
1573    }
1574
1575    pub fn add(&mut self, superkmer: &WeakSuperKmer, seq: &[u8]) -> Result<(), BucketError> {
1576        if superkmer.graph_id >= self.graph_count {
1577            return Err(BucketError::InvalidGraphId(superkmer.graph_id));
1578        }
1579        if seq.len() > u8::MAX as usize {
1580            return Err(BucketError::LabelTooLong(seq.len()));
1581        }
1582
1583        let attr = if self.colored {
1584            let source_id = superkmer.source_id.ok_or(BucketError::MissingSourceId)?;
1585            if source_id > MAX_SOURCE_ID {
1586                return Err(BucketError::SourceIdTooLarge(source_id));
1587            }
1588            pack_colored_attr(
1589                seq.len(),
1590                source_id,
1591                superkmer.left_discontinuous,
1592                superkmer.right_discontinuous,
1593            )
1594        } else {
1595            pack_uncolored_attr(
1596                seq.len(),
1597                superkmer.left_discontinuous,
1598                superkmer.right_discontinuous,
1599            )
1600        };
1601        let graph_id = superkmer.graph_id;
1602        let record_len = record_size(self.colored, self.label_words);
1603        let pending = &mut self.pending[graph_id];
1604        pending.records = pending
1605            .records
1606            .checked_add(1)
1607            .ok_or(BucketError::TooManyRecords)?;
1608        append_record(
1609            &mut pending.bytes,
1610            attr,
1611            graph_id,
1612            seq,
1613            self.label_words,
1614            self.colored,
1615        )?;
1616        self.pending_bytes += record_len;
1617
1618        if self.pending[graph_id].bytes.len() >= MAX_PENDING_BUCKET_BYTES {
1619            self.flush_pending_bucket(graph_id)?;
1620        } else if self.pending_bytes >= MAX_TOTAL_PENDING_BYTES {
1621            self.flush_largest_pending_bucket()?;
1622        }
1623        Ok(())
1624    }
1625
1626    fn flush_largest_pending_bucket(&mut self) -> Result<(), BucketError> {
1627        let Some((graph_id, _)) = self
1628            .pending
1629            .iter()
1630            .enumerate()
1631            .max_by_key(|(_, pending)| pending.bytes.len())
1632        else {
1633            return Ok(());
1634        };
1635        self.flush_pending_bucket(graph_id)
1636    }
1637
1638    fn flush_pending_bucket(&mut self, graph_id: usize) -> Result<(), BucketError> {
1639        if self.pending[graph_id].bytes.is_empty() {
1640            return Ok(());
1641        }
1642        let pending = std::mem::take(&mut self.pending[graph_id]);
1643        self.pending_bytes -= pending.bytes.len();
1644
1645        let (records, bytes_written) = {
1646            self.ensure_writer(graph_id)?;
1647            let Self {
1648                writers, scratch, ..
1649            } = self;
1650            let writer = writers.get_mut(&graph_id).expect("writer just ensured");
1651            writer.write_records(&pending.bytes, pending.records, scratch)?
1652        };
1653        let meta = self.files.get_mut(&graph_id).unwrap();
1654        meta.records = records;
1655        meta.bytes_written = bytes_written;
1656        Ok(())
1657    }
1658
1659    fn ensure_writer(&mut self, graph_id: usize) -> Result<&mut BucketFile, BucketError> {
1660        if !self.writers.contains_key(&graph_id) {
1661            self.evict_writer_if_needed(graph_id)?;
1662            let writer = match self.files.get(&graph_id) {
1663                Some(meta) => {
1664                    BucketFile::open_existing(&meta.path, meta.records, meta.bytes_written)?
1665                }
1666                None => {
1667                    let writer = BucketFile::create(
1668                        &self.bucket_dir,
1669                        self.k,
1670                        self.minimizer_len,
1671                        self.graph_count,
1672                        graph_id,
1673                        self.colored,
1674                        self.label_words,
1675                        self.compress_buckets,
1676                    )?;
1677                    self.files.insert(
1678                        graph_id,
1679                        BucketFileMeta {
1680                            path: writer.path.clone(),
1681                            records: writer.records,
1682                            bytes_written: writer.bytes_written,
1683                        },
1684                    );
1685                    writer
1686                }
1687            };
1688            self.writers.insert(graph_id, writer);
1689        }
1690
1691        Ok(self.writers.get_mut(&graph_id).unwrap())
1692    }
1693
1694    fn evict_writer_if_needed(&mut self, requested_graph_id: usize) -> Result<(), BucketError> {
1695        if self.writers.len() < MAX_OPEN_BUCKET_WRITERS {
1696            return Ok(());
1697        }
1698
1699        let evict_graph_id = self
1700            .writers
1701            .keys()
1702            .copied()
1703            .find(|&graph_id| graph_id != requested_graph_id)
1704            .unwrap_or(requested_graph_id);
1705        if let Some(mut writer) = self.writers.remove(&evict_graph_id) {
1706            writer.flush()?;
1707        }
1708        Ok(())
1709    }
1710
1711    pub fn finish(mut self) -> Result<BucketEmitStats, BucketError> {
1712        let mut manifest = Vec::new();
1713        for graph_id in 0..self.pending.len() {
1714            if !self.pending[graph_id].bytes.is_empty() {
1715                self.flush_pending_bucket(graph_id)?;
1716            }
1717        }
1718
1719        for (graph_id, meta) in &self.files {
1720            if let Some(writer) = self.writers.get_mut(graph_id) {
1721                writer.finish()?;
1722            } else {
1723                BucketFile::finish_closed(&meta.path, meta.records)?;
1724            }
1725            manifest.push((*graph_id, meta.records, meta.path.clone()));
1726        }
1727
1728        write_manifest(&self.bucket_dir, &manifest)?;
1729
1730        Ok(BucketEmitStats {
1731            bucket_dir: self.bucket_dir,
1732            bucket_files: manifest.len(),
1733            bytes_written: self.files.values().map(|meta| meta.bytes_written).sum(),
1734        })
1735    }
1736}
1737
1738impl SharedBucketSink {
1739    pub fn create(params: &BuildParams, graph_count: usize) -> Result<Arc<Self>, BucketError> {
1740        let bucket_dir = bucket_dir(params);
1741        if bucket_dir.exists() {
1742            fs::remove_dir_all(&bucket_dir).map_err(|source| BucketError::Io {
1743                path: bucket_dir.clone(),
1744                source,
1745            })?;
1746        }
1747        fs::create_dir_all(&bucket_dir).map_err(|source| BucketError::Io {
1748            path: bucket_dir.clone(),
1749            source,
1750        })?;
1751
1752        // One container per atlas where descriptors allow. An atlas
1753        // serializes its own writes with a mutex, so at one-to-one a container
1754        // has a single writer; when several atlases share one, the atomic
1755        // segment cursor keeps that safe.
1756        let atlas_count = graph_count.div_ceil(ATLAS_GRAPH_COUNT);
1757        let container_count = BucketContainers::plan_container_count(atlas_count);
1758        if container_count < atlas_count {
1759            eprintln!(
1760                "cuttlefish: descriptor budget allows {container_count} bucket container(s) rather than {atlas_count}"
1761            );
1762        }
1763        let containers = BucketContainers::create(&bucket_dir, container_count)?;
1764
1765        Ok(Arc::new(Self {
1766            bucket_dir,
1767            containers,
1768            k: params.k,
1769            minimizer_len: params.minimizer_len,
1770            graph_count,
1771            colored: params.color,
1772            compress_buckets: params.color || params.compress_buckets,
1773            label_words: label_word_count(params.k, params.minimizer_len),
1774            workers: params.threads.max(1),
1775            atlases: (0..graph_count.div_ceil(ATLAS_GRAPH_COUNT))
1776                .map(|atlas_id| {
1777                    let first_graph_id = atlas_id * ATLAS_GRAPH_COUNT;
1778                    let file_count = (graph_count - first_graph_id).min(ATLAS_GRAPH_COUNT);
1779                    Mutex::new(SharedBucketAtlas {
1780                        first_graph_id,
1781                        files: (0..file_count)
1782                            .map(|_| SharedBucketFileMeta::default())
1783                            .collect(),
1784                        buffered_bytes: 0,
1785                        scratch: CompressionScratch::default(),
1786                    })
1787                })
1788                .collect(),
1789            flush_calls: AtomicU64::new(0),
1790            flush_nanos: AtomicU64::new(0),
1791        }))
1792    }
1793
1794    pub fn emitter(self: &Arc<Self>) -> SharedBucketEmitter {
1795        self.make_emitter(false)
1796    }
1797
1798    pub fn deferred_uncolored_emitter(self: &Arc<Self>) -> SharedBucketEmitter {
1799        self.make_emitter(true)
1800    }
1801
1802    fn make_emitter(self: &Arc<Self>, deferred_uncolored: bool) -> SharedBucketEmitter {
1803        SharedBucketEmitter {
1804            sink: Arc::clone(self),
1805            pending: if self.colored || deferred_uncolored {
1806                Vec::new()
1807            } else {
1808                vec![PendingBucket::default(); self.graph_count]
1809            },
1810            uncolored_pending: if !self.colored && deferred_uncolored {
1811                (0..self.atlases.len())
1812                    .map(|_| PendingColoredAtlas::default())
1813                    .collect()
1814            } else {
1815                Vec::new()
1816            },
1817            colored_pending: if self.colored {
1818                (0..self.atlases.len())
1819                    .map(|_| PendingColoredAtlas::default())
1820                    .collect()
1821            } else {
1822                Vec::new()
1823            },
1824            pending_bytes: 0,
1825            deferred_uncolored,
1826        }
1827    }
1828
1829    pub fn flush_uncolored_emitters(
1830        &self,
1831        emitters: Vec<SharedBucketEmitter>,
1832    ) -> Result<(), BucketError> {
1833        let started = Instant::now();
1834        let mut by_atlas = (0..self.atlases.len())
1835            .map(|_| Mutex::new(Vec::<PendingColoredAtlas>::new()))
1836            .collect::<Vec<_>>();
1837        for mut emitter in emitters {
1838            for (atlas_id, pending) in emitter.uncolored_pending.drain(..).enumerate() {
1839                if !pending.bytes.is_empty() {
1840                    by_atlas[atlas_id]
1841                        .get_mut()
1842                        .map_err(|_| BucketError::WorkerPanic)?
1843                        .push(pending);
1844                }
1845            }
1846        }
1847        let next = AtomicUsize::new(0);
1848        let workers = self.workers.min(self.atlases.len().max(1));
1849        let calls = std::thread::scope(|scope| {
1850            let mut handles = Vec::with_capacity(workers);
1851            for _ in 0..workers {
1852                handles.push(scope.spawn(|| {
1853                    let mut calls = 0;
1854                    loop {
1855                        let atlas_id = next.fetch_add(1, Ordering::Relaxed);
1856                        let Some(pending) = by_atlas.get(atlas_id) else {
1857                            break;
1858                        };
1859                        let pending = std::mem::take(
1860                            &mut *pending.lock().map_err(|_| BucketError::WorkerPanic)?,
1861                        );
1862                        if pending.is_empty() {
1863                            continue;
1864                        }
1865                        let mut stats = SharedBucketFlushStats::default();
1866                        for chunk in pending {
1867                            self.append_uncolored_atlas_inner(atlas_id, chunk, &mut stats)?;
1868                        }
1869                        let mut atlas = self.atlases[atlas_id]
1870                            .lock()
1871                            .map_err(|_| BucketError::WorkerPanic)?;
1872                        atlas.flush_all(
1873                            &self.containers,
1874                            false,
1875                            self.label_words,
1876                            self.compress_buckets,
1877                            &mut stats,
1878                        )?;
1879                        calls += stats.calls;
1880                    }
1881                    Ok::<_, BucketError>(calls)
1882                }));
1883            }
1884            let mut calls = 0;
1885            for handle in handles {
1886                calls += handle.join().map_err(|_| BucketError::WorkerPanic)??;
1887            }
1888            Ok::<_, BucketError>(calls)
1889        })?;
1890        self.record_flush_stats(SharedBucketFlushStats { calls }, started.elapsed());
1891        Ok(())
1892    }
1893
1894    fn append_uncolored_atlas(
1895        &self,
1896        atlas_id: usize,
1897        pending: PendingColoredAtlas,
1898    ) -> Result<(), BucketError> {
1899        let started = Instant::now();
1900        let mut stats = SharedBucketFlushStats::default();
1901        self.append_uncolored_atlas_inner(atlas_id, pending, &mut stats)?;
1902        self.record_flush_stats(stats, started.elapsed());
1903        Ok(())
1904    }
1905
1906    fn append_uncolored_atlas_inner(
1907        &self,
1908        atlas_id: usize,
1909        pending: PendingColoredAtlas,
1910        stats: &mut SharedBucketFlushStats,
1911    ) -> Result<(), BucketError> {
1912        if pending.bytes.is_empty() {
1913            return Ok(());
1914        }
1915        let record_len = record_size(false, self.label_words);
1916        if pending.graph_ids.len() * record_len != pending.bytes.len() {
1917            return Err(BucketError::MalformedRecord);
1918        }
1919        let mut atlas = self.atlases[atlas_id]
1920            .lock()
1921            .map_err(|_| BucketError::WorkerPanic)?;
1922        for (&graph_id, record) in pending
1923            .graph_ids
1924            .iter()
1925            .zip(pending.bytes.chunks_exact(record_len))
1926        {
1927            let graph_id = usize::from(graph_id);
1928            let Some(local_graph_id) = graph_id.checked_sub(atlas.first_graph_id) else {
1929                return Err(BucketError::InvalidGraphId(graph_id));
1930            };
1931            let Some(file) = atlas.files.get_mut(local_graph_id) else {
1932                return Err(BucketError::InvalidGraphId(graph_id));
1933            };
1934            file.total_records = file
1935                .total_records
1936                .checked_add(1)
1937                .ok_or(BucketError::TooManyRecords)?;
1938            file.buffer_records = file
1939                .buffer_records
1940                .checked_add(1)
1941                .ok_or(BucketError::TooManyRecords)?;
1942            file.buffer.extend_from_slice(record);
1943        }
1944        atlas.buffered_bytes += pending.bytes.len();
1945        for local_graph_id in 0..atlas.files.len() {
1946            if atlas.files[local_graph_id].buffer.len() >= SUBGRAPH_CHUNK_BYTES {
1947                atlas.flush_subgraph(
1948                    local_graph_id,
1949                    &self.containers,
1950                    false,
1951                    self.label_words,
1952                    self.compress_buckets,
1953                    stats,
1954                )?;
1955            }
1956        }
1957        Ok(())
1958    }
1959
1960    fn append_bucket(&self, graph_id: usize, pending: PendingBucket) -> Result<(), BucketError> {
1961        if pending.bytes.is_empty() {
1962            return Ok(());
1963        }
1964        let started = Instant::now();
1965        let atlas_id = graph_id / ATLAS_GRAPH_COUNT;
1966        let local_graph_id = graph_id % ATLAS_GRAPH_COUNT;
1967        let mut atlas = self.atlases[atlas_id]
1968            .lock()
1969            .map_err(|_| BucketError::WorkerPanic)?;
1970        atlas.append_bucket(local_graph_id, pending)?;
1971        if self.colored {
1972            return Ok(());
1973        }
1974        if atlas.files[local_graph_id].buffer.len() < SUBGRAPH_CHUNK_BYTES {
1975            return Ok(());
1976        }
1977
1978        let mut flush_stats = SharedBucketFlushStats::default();
1979        atlas.flush_subgraph(
1980            local_graph_id,
1981            &self.containers,
1982            self.colored,
1983            self.label_words,
1984            self.compress_buckets,
1985            &mut flush_stats,
1986        )?;
1987        self.record_flush_stats(flush_stats, started.elapsed());
1988        Ok(())
1989    }
1990
1991    fn append_colored_atlas(
1992        &self,
1993        atlas_id: usize,
1994        pending: PendingColoredAtlas,
1995    ) -> Result<(), BucketError> {
1996        if pending.bytes.is_empty() {
1997            return Ok(());
1998        }
1999        let record_len = record_size(true, self.label_words);
2000        if pending.graph_ids.len() * record_len != pending.bytes.len() {
2001            return Err(BucketError::MalformedRecord);
2002        }
2003        let started = Instant::now();
2004        let mut atlas = self.atlases[atlas_id]
2005            .lock()
2006            .map_err(|_| BucketError::WorkerPanic)?;
2007        for (&graph_id, record) in pending
2008            .graph_ids
2009            .iter()
2010            .zip(pending.bytes.chunks_exact(record_len))
2011        {
2012            let graph_id = usize::from(graph_id);
2013            let Some(local_graph_id) = graph_id.checked_sub(atlas.first_graph_id) else {
2014                return Err(BucketError::InvalidGraphId(graph_id));
2015            };
2016            let Some(file) = atlas.files.get_mut(local_graph_id) else {
2017                return Err(BucketError::InvalidGraphId(graph_id));
2018            };
2019            file.total_records = file
2020                .total_records
2021                .checked_add(1)
2022                .ok_or(BucketError::TooManyRecords)?;
2023            file.buffer_records = file
2024                .buffer_records
2025                .checked_add(1)
2026                .ok_or(BucketError::TooManyRecords)?;
2027            file.buffer.extend_from_slice(record);
2028        }
2029        atlas.buffered_bytes += pending.bytes.len();
2030        let mut flush_stats = SharedBucketFlushStats::default();
2031        for local_graph_id in 0..atlas.files.len() {
2032            if atlas.files[local_graph_id].buffer.len() >= SUBGRAPH_CHUNK_BYTES {
2033                atlas.flush_subgraph(
2034                    local_graph_id,
2035                    &self.containers,
2036                    true,
2037                    self.label_words,
2038                    self.compress_buckets,
2039                    &mut flush_stats,
2040                )?;
2041            }
2042        }
2043        self.record_flush_stats(flush_stats, started.elapsed());
2044        Ok(())
2045    }
2046
2047    pub fn flush_stats(&self) -> (u64, Duration) {
2048        (
2049            self.flush_calls.load(Ordering::Relaxed),
2050            Duration::from_nanos(self.flush_nanos.load(Ordering::Relaxed)),
2051        )
2052    }
2053
2054    pub fn flush_colored_window(
2055        &self,
2056        source_min: u32,
2057        source_max: u32,
2058    ) -> Result<(), BucketError> {
2059        if !self.colored || source_min > source_max {
2060            return Err(BucketError::MalformedRecord);
2061        }
2062        let started = Instant::now();
2063        let record_len = record_size(true, self.label_words);
2064        let next_atlas = AtomicUsize::new(0);
2065        let workers = self.workers.min(self.atlases.len().max(1));
2066        let flush_calls = std::thread::scope(|scope| {
2067            let mut handles = Vec::with_capacity(workers);
2068            for _ in 0..workers {
2069                handles.push(scope.spawn(|| {
2070                    let mut stats = SharedBucketFlushStats::default();
2071                    loop {
2072                        let atlas_id = next_atlas.fetch_add(1, Ordering::Relaxed);
2073                        let Some(atlas) = self.atlases.get(atlas_id) else {
2074                            break;
2075                        };
2076                        let mut atlas = atlas.lock().map_err(|_| BucketError::WorkerPanic)?;
2077                        for local_graph_id in 0..atlas.files.len() {
2078                            sort_colored_payload_by_source(
2079                                &mut atlas.files[local_graph_id].buffer,
2080                                record_len,
2081                                source_min,
2082                                source_max,
2083                            )?;
2084                            atlas.flush_subgraph(
2085                                local_graph_id,
2086                                &self.containers,
2087                                true,
2088                                self.label_words,
2089                                true,
2090                                &mut stats,
2091                            )?;
2092                        }
2093                    }
2094                    Ok::<_, BucketError>(stats.calls)
2095                }));
2096            }
2097            let mut calls = 0u64;
2098            for handle in handles {
2099                calls += handle.join().map_err(|_| BucketError::WorkerPanic)??;
2100            }
2101            Ok::<_, BucketError>(calls)
2102        })?;
2103        self.record_flush_stats(
2104            SharedBucketFlushStats { calls: flush_calls },
2105            started.elapsed(),
2106        );
2107        Ok(())
2108    }
2109
2110    pub fn flush_colored_emitters(
2111        &self,
2112        emitters: Vec<SharedBucketEmitter>,
2113    ) -> Result<(), BucketError> {
2114        if !self.colored {
2115            return Err(BucketError::MalformedRecord);
2116        }
2117
2118        let emitter_count = emitters.len();
2119        let mut pending_by_atlas = (0..self.atlases.len())
2120            .map(|_| Vec::with_capacity(emitter_count))
2121            .collect::<Vec<Vec<PendingColoredAtlas>>>();
2122        for emitter in emitters {
2123            if !std::ptr::eq(Arc::as_ptr(&emitter.sink), self) {
2124                return Err(BucketError::MalformedRecord);
2125            }
2126            if !emitter.pending.is_empty()
2127                || emitter.colored_pending.len() != pending_by_atlas.len()
2128            {
2129                return Err(BucketError::MalformedRecord);
2130            }
2131            for (atlas_id, pending) in emitter.colored_pending.into_iter().enumerate() {
2132                if !pending.bytes.is_empty() {
2133                    pending_by_atlas[atlas_id].push(pending);
2134                }
2135            }
2136        }
2137
2138        let started = Instant::now();
2139        let workers = self.workers.min(pending_by_atlas.len().max(1));
2140        let chunk_size = pending_by_atlas.len().div_ceil(workers);
2141        let flush_calls = std::thread::scope(|scope| {
2142            let mut handles = Vec::with_capacity(workers);
2143            for (chunk_id, atlas_work) in pending_by_atlas.chunks_mut(chunk_size).enumerate() {
2144                handles.push(scope.spawn(move || {
2145                    let mut stats = SharedBucketFlushStats::default();
2146                    let first_atlas = chunk_id * chunk_size;
2147                    for (offset, worker_buckets) in atlas_work.iter_mut().enumerate() {
2148                        let atlas_id = first_atlas + offset;
2149                        let mut atlas = self.atlases[atlas_id]
2150                            .lock()
2151                            .map_err(|_| BucketError::WorkerPanic)?;
2152                        let record_len = record_size(true, self.label_words);
2153                        for pending in worker_buckets.iter() {
2154                            if pending.graph_ids.len() * record_len != pending.bytes.len() {
2155                                return Err(BucketError::MalformedRecord);
2156                            }
2157                            for (&graph_id, record) in pending
2158                                .graph_ids
2159                                .iter()
2160                                .zip(pending.bytes.chunks_exact(record_len))
2161                            {
2162                                append_colored_atlas_record(&mut atlas, graph_id, record)?;
2163                            }
2164                        }
2165                        atlas.buffered_bytes += worker_buckets
2166                            .iter()
2167                            .map(|pending| pending.bytes.len())
2168                            .sum::<usize>();
2169                        for local_graph_id in 0..atlas.files.len() {
2170                            atlas.flush_subgraph(
2171                                local_graph_id,
2172                                &self.containers,
2173                                true,
2174                                self.label_words,
2175                                true,
2176                                &mut stats,
2177                            )?;
2178                            // This consumes the complete colored emitter set, so these
2179                            // buffers will not be reused. Releasing their capacity here
2180                            // avoids retaining every uncompressed subgraph bucket while
2181                            // the remaining worker atlas chunks are still resident.
2182                            atlas.files[local_graph_id].buffer = Vec::new();
2183                        }
2184                    }
2185                    Ok::<_, BucketError>(stats.calls)
2186                }));
2187            }
2188            let mut calls = 0u64;
2189            for handle in handles {
2190                calls += handle.join().map_err(|_| BucketError::WorkerPanic)??;
2191            }
2192            Ok::<_, BucketError>(calls)
2193        })?;
2194        self.record_flush_stats(
2195            SharedBucketFlushStats { calls: flush_calls },
2196            started.elapsed(),
2197        );
2198        Ok(())
2199    }
2200
2201    pub fn finish(&self) -> Result<BucketEmitStats, BucketError> {
2202        let mut entries = Vec::new();
2203        let mut finish_flush_stats = SharedBucketFlushStats::default();
2204        let finish_started = Instant::now();
2205        let mut bytes_written = 0u64;
2206        for atlas in &self.atlases {
2207            let mut atlas = atlas.lock().map_err(|_| BucketError::WorkerPanic)?;
2208            atlas.flush_all(
2209                &self.containers,
2210                self.colored,
2211                self.label_words,
2212                self.compress_buckets,
2213                &mut finish_flush_stats,
2214            )?;
2215            let first_graph_id = atlas.first_graph_id;
2216            let container = (first_graph_id / ATLAS_GRAPH_COUNT) % self.containers.len();
2217            for (local_graph_id, meta) in atlas.files.iter_mut().enumerate() {
2218                let meta = std::mem::take(meta);
2219                if meta.total_records == 0 {
2220                    continue;
2221                }
2222                bytes_written += meta.bytes_written;
2223                entries.push(BucketManifestEntry {
2224                    graph_id: first_graph_id + local_graph_id,
2225                    records: meta.total_records,
2226                    location: BucketLocation::Container {
2227                        container,
2228                        segments: meta.segments,
2229                        bytes: meta.bytes_written,
2230                    },
2231                });
2232            }
2233        }
2234        self.record_flush_stats(finish_flush_stats, finish_started.elapsed());
2235        entries.sort_unstable_by_key(|entry| entry.graph_id);
2236
2237        // Nothing to patch and nothing to reopen. The per-file layout finished
2238        // by reopening all 16,384 buckets to write the record count into each
2239        // header, parallelised across workers because it was slow enough to
2240        // matter; the count now lives in the manifest that has to be written
2241        // anyway.
2242        let header = ContainerManifestHeader {
2243            k: self.k,
2244            minimizer_len: self.minimizer_len,
2245            graph_count: self.graph_count,
2246            colored: self.colored,
2247            label_words: self.label_words,
2248            compressed: self.compress_buckets,
2249            interleaved_compression: self.compress_buckets
2250                && !self.colored
2251                && !force_split_compression(),
2252            segment_bytes: self.containers.segment_bytes(),
2253            // The number of containers actually created, which is not the
2254            // atlas count when the descriptor budget narrowed it.
2255            container_count: self.containers.len(),
2256        };
2257        write_container_manifest(&self.bucket_dir, &header, &entries)?;
2258        Ok(BucketEmitStats {
2259            bucket_dir: self.bucket_dir.clone(),
2260            bucket_files: entries.len(),
2261            bytes_written,
2262        })
2263    }
2264
2265    fn record_flush_stats(&self, stats: SharedBucketFlushStats, elapsed: Duration) {
2266        if stats.calls == 0 {
2267            return;
2268        }
2269        self.flush_calls.fetch_add(stats.calls, Ordering::Relaxed);
2270        self.flush_nanos.fetch_add(
2271            u64::try_from(elapsed.as_nanos()).unwrap_or(u64::MAX),
2272            Ordering::Relaxed,
2273        );
2274    }
2275}
2276
2277fn append_colored_atlas_record(
2278    atlas: &mut SharedBucketAtlas,
2279    graph_id: u16,
2280    record: &[u8],
2281) -> Result<(), BucketError> {
2282    let graph_id = usize::from(graph_id);
2283    let local_graph_id = graph_id
2284        .checked_sub(atlas.first_graph_id)
2285        .ok_or(BucketError::InvalidGraphId(graph_id))?;
2286    let file = atlas
2287        .files
2288        .get_mut(local_graph_id)
2289        .ok_or(BucketError::InvalidGraphId(graph_id))?;
2290    file.total_records = file
2291        .total_records
2292        .checked_add(1)
2293        .ok_or(BucketError::TooManyRecords)?;
2294    file.buffer_records = file
2295        .buffer_records
2296        .checked_add(1)
2297        .ok_or(BucketError::TooManyRecords)?;
2298    file.buffer.extend_from_slice(record);
2299    Ok(())
2300}
2301
2302fn sort_colored_payload_by_source(
2303    payload: &mut Vec<u8>,
2304    record_len: usize,
2305    source_min: u32,
2306    source_max: u32,
2307) -> Result<(), BucketError> {
2308    if payload.is_empty() {
2309        return Ok(());
2310    }
2311    if payload.len() % record_len != 0 {
2312        return Err(BucketError::MalformedRecord);
2313    }
2314    let source_count =
2315        usize::try_from(source_max - source_min + 1).map_err(|_| BucketError::MalformedRecord)?;
2316    let mut offsets = vec![0usize; source_count];
2317    for record in payload.chunks_exact(record_len) {
2318        let attr = u32::from_le_bytes(record[..4].try_into().expect("colored attribute"));
2319        let source = attr >> 10;
2320        if source < source_min || source > source_max {
2321            return Err(BucketError::MalformedRecord);
2322        }
2323        offsets[(source - source_min) as usize] += 1;
2324    }
2325    let mut prefix = 0usize;
2326    for offset in &mut offsets {
2327        let count = *offset;
2328        *offset = prefix;
2329        prefix += count;
2330    }
2331    let mut sorted = vec![0u8; payload.len()];
2332    for record in payload.chunks_exact(record_len) {
2333        let attr = u32::from_le_bytes(record[..4].try_into().expect("colored attribute"));
2334        let source = ((attr >> 10) - source_min) as usize;
2335        let output = offsets[source] * record_len;
2336        sorted[output..output + record_len].copy_from_slice(record);
2337        offsets[source] += 1;
2338    }
2339    *payload = sorted;
2340    Ok(())
2341}
2342
2343impl SharedBucketAtlas {
2344    fn append_bucket(
2345        &mut self,
2346        local_graph_id: usize,
2347        pending: PendingBucket,
2348    ) -> Result<(), BucketError> {
2349        let file = &mut self.files[local_graph_id];
2350        file.total_records = file
2351            .total_records
2352            .checked_add(pending.records)
2353            .ok_or(BucketError::TooManyRecords)?;
2354        file.buffer_records = file
2355            .buffer_records
2356            .checked_add(pending.records)
2357            .ok_or(BucketError::TooManyRecords)?;
2358        self.buffered_bytes += pending.bytes.len();
2359        file.buffer.extend_from_slice(&pending.bytes);
2360        Ok(())
2361    }
2362
2363    fn flush_all(
2364        &mut self,
2365        containers: &BucketContainers,
2366        colored: bool,
2367        label_words: usize,
2368        compress_buckets: bool,
2369        stats: &mut SharedBucketFlushStats,
2370    ) -> Result<(), BucketError> {
2371        for local_graph_id in 0..self.files.len() {
2372            self.flush_subgraph(
2373                local_graph_id,
2374                containers,
2375                colored,
2376                label_words,
2377                compress_buckets,
2378                stats,
2379            )?;
2380        }
2381        Ok(())
2382    }
2383
2384    /// Writes one bucket's staged records into its container.
2385    ///
2386    /// This is the whole syscall saving. The per-file path reached here with
2387    /// an `openat`, seven unbuffered reads to re-read the 42-byte header, a
2388    /// revalidation, an `lseek` to the end and a `close` -- around eleven
2389    /// syscalls for every 64 KiB flush, and a full-corpus build performs 14.4
2390    /// million of them. A container flush is the `pwrite` and nothing else:
2391    /// the header is in the manifest, the descriptor is already open, and the
2392    /// offset is known rather than sought.
2393    fn flush_subgraph(
2394        &mut self,
2395        local_graph_id: usize,
2396        containers: &BucketContainers,
2397        colored: bool,
2398        label_words: usize,
2399        compress_buckets: bool,
2400        stats: &mut SharedBucketFlushStats,
2401    ) -> Result<(), BucketError> {
2402        let container = (self.first_graph_id / ATLAS_GRAPH_COUNT) % containers.len();
2403        let Self { files, scratch, .. } = self;
2404        let file = &mut files[local_graph_id];
2405        if file.buffer.is_empty() {
2406            return Ok(());
2407        }
2408        let flushed_bytes = file.buffer.len();
2409        let record_size = record_size(colored, label_words) as u64;
2410
2411        let written = if compress_buckets {
2412            let interleaved = !colored && !force_split_compression();
2413            let len = encode_compressed_block(
2414                &file.buffer,
2415                file.buffer_records,
2416                record_size,
2417                label_words,
2418                interleaved,
2419                scratch,
2420            )?;
2421            let block = std::mem::take(&mut scratch.block);
2422            let written = append_to_chain(containers, container, file, &block)?;
2423            scratch.block = block;
2424            debug_assert_eq!(written, len as u64);
2425            written
2426        } else {
2427            let buffer = std::mem::take(&mut file.buffer);
2428            let written = append_to_chain(containers, container, file, &buffer);
2429            file.buffer = buffer;
2430            written?
2431        };
2432
2433        file.written_records += file.buffer_records;
2434        file.bytes_written += written;
2435        file.buffer.clear();
2436        file.buffer_records = 0;
2437        self.buffered_bytes -= flushed_bytes;
2438        stats.calls += 1;
2439        Ok(())
2440    }
2441}
2442
2443/// Appends `bytes` to a bucket's segment chain, reserving as it goes.
2444///
2445/// A block may straddle a segment boundary rather than starting a fresh
2446/// segment. That costs a second `pwrite` for the rare block that spans one,
2447/// and saves refusing to fill the tail of every segment -- which at a 64 KiB
2448/// flush and a 256 KiB segment would waste up to a quarter of the directory.
2449/// The reader concatenates the chain in order, so a split block reassembles.
2450fn append_to_chain(
2451    containers: &BucketContainers,
2452    container: usize,
2453    file: &mut SharedBucketFileMeta,
2454    bytes: &[u8],
2455) -> Result<u64, BucketError> {
2456    let segment_bytes = containers.segment_bytes();
2457    let mut written = 0usize;
2458    while written < bytes.len() {
2459        if file.segments.is_empty() || file.segment_used == segment_bytes {
2460            let offset = containers.reserve_segment(container);
2461            let index =
2462                u32::try_from(offset / segment_bytes).map_err(|_| BucketError::TooManyRecords)?;
2463            file.segments.push(index);
2464            file.segment_used = 0;
2465        }
2466        let room = (segment_bytes - file.segment_used) as usize;
2467        let take = room.min(bytes.len() - written);
2468        let offset = u64::from(*file.segments.last().unwrap()) * segment_bytes + file.segment_used;
2469        containers.write_at(container, offset, &bytes[written..written + take])?;
2470        written += take;
2471        file.segment_used += take as u64;
2472    }
2473    Ok(written as u64)
2474}
2475
2476impl SharedBucketEmitter {
2477    pub fn flush_colored_worker_if_required(&mut self) -> Result<(), BucketError> {
2478        if !self.sink.colored {
2479            return Err(BucketError::MalformedRecord);
2480        }
2481        for atlas_id in 0..self.colored_pending.len() {
2482            if self.colored_pending[atlas_id].bytes.len() >= SUBGRAPH_CHUNK_BYTES {
2483                self.flush_pending_colored_atlas(atlas_id)?;
2484            }
2485        }
2486        Ok(())
2487    }
2488
2489    pub fn add(&mut self, superkmer: &WeakSuperKmer, seq: &[u8]) -> Result<(), BucketError> {
2490        self.add_impl(superkmer, seq, true)
2491    }
2492
2493    pub fn add_valid(&mut self, superkmer: &WeakSuperKmer, seq: &[u8]) -> Result<(), BucketError> {
2494        debug_assert!(seq.iter().all(|&base| ascii_base_bits(base).is_some()));
2495        self.add_impl(superkmer, seq, false)
2496    }
2497
2498    fn add_impl(
2499        &mut self,
2500        superkmer: &WeakSuperKmer,
2501        seq: &[u8],
2502        check_bases: bool,
2503    ) -> Result<(), BucketError> {
2504        if superkmer.graph_id >= self.sink.graph_count {
2505            return Err(BucketError::InvalidGraphId(superkmer.graph_id));
2506        }
2507        if seq.len() > u8::MAX as usize {
2508            return Err(BucketError::LabelTooLong(seq.len()));
2509        }
2510
2511        let attr = if self.sink.colored {
2512            let source_id = superkmer.source_id.ok_or(BucketError::MissingSourceId)?;
2513            if source_id > MAX_SOURCE_ID {
2514                return Err(BucketError::SourceIdTooLarge(source_id));
2515            }
2516            pack_colored_attr(
2517                seq.len(),
2518                source_id,
2519                superkmer.left_discontinuous,
2520                superkmer.right_discontinuous,
2521            )
2522        } else {
2523            pack_uncolored_attr(
2524                seq.len(),
2525                superkmer.left_discontinuous,
2526                superkmer.right_discontinuous,
2527            )
2528        };
2529        let graph_id = superkmer.graph_id;
2530        let record_len = record_size(self.sink.colored, self.sink.label_words);
2531        if self.sink.colored {
2532            let atlas_id = graph_id / ATLAS_GRAPH_COUNT;
2533            let pending = &mut self.colored_pending[atlas_id];
2534            pending.graph_ids.push(graph_id as u16);
2535            if check_bases {
2536                append_record(
2537                    &mut pending.bytes,
2538                    attr,
2539                    graph_id,
2540                    seq,
2541                    self.sink.label_words,
2542                    true,
2543                )?;
2544            } else {
2545                append_record_valid(
2546                    &mut pending.bytes,
2547                    attr,
2548                    graph_id,
2549                    seq,
2550                    self.sink.label_words,
2551                    true,
2552                )?;
2553            }
2554        } else if self.deferred_uncolored {
2555            let atlas_id = graph_id / ATLAS_GRAPH_COUNT;
2556            let pending = &mut self.uncolored_pending[atlas_id];
2557            pending.graph_ids.push(graph_id as u16);
2558            if !check_bases {
2559                append_uncolored_record_valid(
2560                    &mut pending.bytes,
2561                    attr as u16,
2562                    graph_id as u16,
2563                    seq,
2564                    self.sink.label_words,
2565                )?;
2566            } else {
2567                append_record(
2568                    &mut pending.bytes,
2569                    attr,
2570                    graph_id,
2571                    seq,
2572                    self.sink.label_words,
2573                    false,
2574                )?;
2575            }
2576        } else {
2577            let pending = &mut self.pending[graph_id];
2578            pending.records = pending
2579                .records
2580                .checked_add(1)
2581                .ok_or(BucketError::TooManyRecords)?;
2582            if !check_bases {
2583                append_uncolored_record_valid(
2584                    &mut pending.bytes,
2585                    attr as u16,
2586                    graph_id as u16,
2587                    seq,
2588                    self.sink.label_words,
2589                )?;
2590            } else {
2591                append_record(
2592                    &mut pending.bytes,
2593                    attr,
2594                    graph_id,
2595                    seq,
2596                    self.sink.label_words,
2597                    false,
2598                )?;
2599            }
2600        }
2601        self.pending_bytes += record_len;
2602
2603        // As in C++, colored worker-atlas chunks are checked and handed to the
2604        // shared atlas at the source boundary, not in the middle of a source.
2605        if self.deferred_uncolored {
2606            let atlas_id = graph_id / ATLAS_GRAPH_COUNT;
2607            if self.uncolored_pending[atlas_id].bytes.len() >= SUBGRAPH_CHUNK_BYTES {
2608                self.flush_pending_uncolored_atlas(atlas_id)?;
2609            }
2610        } else if !self.sink.colored
2611            && !self.deferred_uncolored
2612            && self.pending[graph_id].bytes.len() >= MAX_PENDING_BUCKET_BYTES
2613        {
2614            self.flush_pending_bucket(graph_id)?;
2615        } else if !self.sink.colored
2616            && !self.deferred_uncolored
2617            && self.pending_bytes >= MAX_TOTAL_PENDING_BYTES
2618        {
2619            self.flush_largest_pending_bucket()?;
2620        }
2621        Ok(())
2622    }
2623
2624    fn flush_largest_pending_bucket(&mut self) -> Result<(), BucketError> {
2625        if self.sink.colored {
2626            let Some((atlas_id, _)) = self
2627                .colored_pending
2628                .iter()
2629                .enumerate()
2630                .max_by_key(|(_, pending)| pending.bytes.len())
2631            else {
2632                return Ok(());
2633            };
2634            return self.flush_pending_colored_atlas(atlas_id);
2635        }
2636        let Some((graph_id, _)) = self
2637            .pending
2638            .iter()
2639            .enumerate()
2640            .max_by_key(|(_, pending)| pending.bytes.len())
2641        else {
2642            return Ok(());
2643        };
2644        self.flush_pending_bucket(graph_id)
2645    }
2646
2647    fn flush_pending_bucket(&mut self, graph_id: usize) -> Result<(), BucketError> {
2648        if self.pending[graph_id].bytes.is_empty() {
2649            return Ok(());
2650        }
2651        let pending = std::mem::take(&mut self.pending[graph_id]);
2652        self.pending_bytes -= pending.bytes.len();
2653        self.sink.append_bucket(graph_id, pending)
2654    }
2655
2656    fn flush_pending_colored_atlas(&mut self, atlas_id: usize) -> Result<(), BucketError> {
2657        if self.colored_pending[atlas_id].bytes.is_empty() {
2658            return Ok(());
2659        }
2660        let pending = std::mem::take(&mut self.colored_pending[atlas_id]);
2661        self.pending_bytes -= pending.bytes.len();
2662        self.sink.append_colored_atlas(atlas_id, pending)
2663    }
2664
2665    fn flush_pending_uncolored_atlas(&mut self, atlas_id: usize) -> Result<(), BucketError> {
2666        if self.uncolored_pending[atlas_id].bytes.is_empty() {
2667            return Ok(());
2668        }
2669        let pending = std::mem::take(&mut self.uncolored_pending[atlas_id]);
2670        self.pending_bytes -= pending.bytes.len();
2671        self.sink.append_uncolored_atlas(atlas_id, pending)
2672    }
2673
2674    pub fn finish(mut self) -> Result<(), BucketError> {
2675        if self.sink.colored {
2676            for atlas_id in 0..self.colored_pending.len() {
2677                if !self.colored_pending[atlas_id].bytes.is_empty() {
2678                    self.flush_pending_colored_atlas(atlas_id)?;
2679                }
2680            }
2681        } else if self.deferred_uncolored {
2682            for atlas_id in 0..self.uncolored_pending.len() {
2683                if !self.uncolored_pending[atlas_id].bytes.is_empty() {
2684                    self.flush_pending_uncolored_atlas(atlas_id)?;
2685                }
2686            }
2687        } else {
2688            for graph_id in 0..self.pending.len() {
2689                if !self.pending[graph_id].bytes.is_empty() {
2690                    self.flush_pending_bucket(graph_id)?;
2691                }
2692            }
2693        }
2694        Ok(())
2695    }
2696}
2697
2698pub fn bucket_dir(params: &BuildParams) -> PathBuf {
2699    let output_name = Path::new(&params.output_prefix)
2700        .file_name()
2701        .and_then(|s| s.to_str())
2702        .filter(|s| !s.is_empty())
2703        .unwrap_or("cuttlefish3");
2704    PathBuf::from(&params.work_dir).join(format!("{output_name}.cf3rs.wsk"))
2705}
2706
2707pub fn write_manifest(
2708    bucket_dir: &Path,
2709    entries: &[(usize, u64, PathBuf)],
2710) -> Result<(), BucketError> {
2711    let path = bucket_dir.join("manifest.tsv");
2712    let mut out = File::create(&path).map_err(|source| BucketError::Io {
2713        path: path.clone(),
2714        source,
2715    })?;
2716    writeln!(out, "graph_id\trecords\tpath").map_err(|source| BucketError::Io {
2717        path: path.clone(),
2718        source,
2719    })?;
2720    for (graph_id, records, bucket_path) in entries {
2721        writeln!(out, "{graph_id}\t{records}\t{}", bucket_path.display()).map_err(|source| {
2722            BucketError::Io {
2723                path: path.clone(),
2724                source,
2725            }
2726        })?;
2727    }
2728    Ok(())
2729}
2730
2731/// Encodes one compressed block into `scratch.block`, returning its length.
2732///
2733/// Assembling the whole block -- 12-byte header plus one or two LZ4 streams --
2734/// before it is written keeps a flush to a single `write` on a file that has
2735/// no buffering, and lets a container flush reuse the identical framing.
2736fn encode_compressed_block(
2737    bytes: &[u8],
2738    records: u64,
2739    record_size: u64,
2740    label_words: usize,
2741    interleaved: bool,
2742    scratch: &mut CompressionScratch,
2743) -> Result<usize, BucketError> {
2744    let records_u32 = u32::try_from(records).map_err(|_| BucketError::TooManyRecords)?;
2745    if records_u32 == 0 {
2746        scratch.block.clear();
2747        return Ok(0);
2748    }
2749    let (attr_len, label_len) = if interleaved {
2750        (CompressionScratch::encode(bytes, &mut scratch.encoded)?, 0)
2751    } else {
2752        let record_len = usize::try_from(record_size).unwrap();
2753        let fixed_len = record_len - label_words * 8;
2754        scratch.attrs.clear();
2755        scratch.labels.clear();
2756        scratch.attrs.reserve(records as usize * fixed_len);
2757        scratch.labels.reserve(records as usize * label_words * 8);
2758        for record in bytes.chunks_exact(record_len) {
2759            scratch.attrs.extend_from_slice(&record[..fixed_len]);
2760            scratch.labels.extend_from_slice(&record[fixed_len..]);
2761        }
2762        let attrs = std::mem::take(&mut scratch.attrs);
2763        let labels = std::mem::take(&mut scratch.labels);
2764        let attr_len = CompressionScratch::encode(&attrs, &mut scratch.encoded)?;
2765        let label_len = CompressionScratch::encode(&labels, &mut scratch.encoded_labels)?;
2766        scratch.attrs = attrs;
2767        scratch.labels = labels;
2768        (attr_len, label_len)
2769    };
2770    let attr_len_u32 = u32::try_from(attr_len).map_err(|_| BucketError::TooManyRecords)?;
2771    let label_len_u32 = u32::try_from(label_len).map_err(|_| BucketError::TooManyRecords)?;
2772
2773    scratch.block.clear();
2774    scratch.block.extend_from_slice(&records_u32.to_le_bytes());
2775    scratch.block.extend_from_slice(&attr_len_u32.to_le_bytes());
2776    scratch
2777        .block
2778        .extend_from_slice(&label_len_u32.to_le_bytes());
2779    scratch
2780        .block
2781        .extend_from_slice(&scratch.encoded[..attr_len]);
2782    if label_len != 0 {
2783        scratch
2784            .block
2785            .extend_from_slice(&scratch.encoded_labels[..label_len]);
2786    }
2787    Ok(scratch.block.len())
2788}
2789
2790fn append_record(
2791    out: &mut Vec<u8>,
2792    packed_attr: u32,
2793    _graph_id: usize,
2794    seq: &[u8],
2795    label_words: usize,
2796    colored: bool,
2797) -> Result<(), BucketError> {
2798    let mut words = [0u64; 4];
2799    if label_words > words.len() {
2800        return Err(BucketError::MalformedRecord);
2801    }
2802    for (idx, &ch) in seq.iter().enumerate() {
2803        let base_bits = ascii_base_bits(ch).ok_or(BucketError::InvalidBase(ch))?;
2804        let word_idx = idx / 32;
2805        let shift = 2 * (31 - (idx % 32));
2806        words[word_idx] |= (base_bits as u64) << shift;
2807    }
2808
2809    if colored {
2810        out.extend_from_slice(&packed_attr.to_le_bytes());
2811    } else {
2812        out.extend_from_slice(&(packed_attr as u16).to_le_bytes());
2813    }
2814    for &word in &words[..label_words] {
2815        out.extend_from_slice(&word.to_le_bytes());
2816    }
2817    Ok(())
2818}
2819
2820fn append_record_valid(
2821    out: &mut Vec<u8>,
2822    packed_attr: u32,
2823    _graph_id: usize,
2824    seq: &[u8],
2825    label_words: usize,
2826    colored: bool,
2827) -> Result<(), BucketError> {
2828    let words = pack_valid_label(seq, label_words)?;
2829
2830    let mut record = [0u8; MAX_RECORD_BYTES];
2831    let attr_len = if colored {
2832        record[..4].copy_from_slice(&packed_attr.to_le_bytes());
2833        4
2834    } else {
2835        record[..2].copy_from_slice(&(packed_attr as u16).to_le_bytes());
2836        2
2837    };
2838    for (idx, &word) in words[..label_words].iter().enumerate() {
2839        let at = attr_len + idx * 8;
2840        record[at..at + 8].copy_from_slice(&word.to_le_bytes());
2841    }
2842    out.extend_from_slice(&record[..attr_len + label_words * 8]);
2843    Ok(())
2844}
2845
2846#[inline]
2847fn append_uncolored_record_valid(
2848    out: &mut Vec<u8>,
2849    packed_attr: u16,
2850    _graph_id: u16,
2851    seq: &[u8],
2852    label_words: usize,
2853) -> Result<(), BucketError> {
2854    let words = pack_valid_label(seq, label_words)?;
2855
2856    // Assemble the record on the stack and append it once. Reserving and then
2857    // extending per field costs a capacity check and a `memcpy` call for each of
2858    // the attribute and every label word; C++ writes its equivalent with plain
2859    // indexed stores into pre-reserved arrays.
2860    let mut record = [0u8; MAX_RECORD_BYTES];
2861    record[..2].copy_from_slice(&packed_attr.to_le_bytes());
2862    for (idx, &word) in words[..label_words].iter().enumerate() {
2863        let at = 2 + idx * 8;
2864        record[at..at + 8].copy_from_slice(&word.to_le_bytes());
2865    }
2866    out.extend_from_slice(&record[..2 + label_words * 8]);
2867    Ok(())
2868}
2869
2870#[inline(always)]
2871/// Packs 32 valid ACGT bases into one 64-bit word, first base in the high bits.
2872///
2873/// The scalar form carries a loop-carried dependency (`word = (word << 2) | c`),
2874/// so it neither vectorizes nor pipelines: 32 serial iterations for one word,
2875/// which dominates record packing. Each 2-bit code is a pure bitwise function of
2876/// its byte, so eight bases can be reduced at a time inside a `u64` and the
2877/// codes gathered with a single `PEXT`.
2878fn pack_valid_word_32(seq: &[u8]) -> u64 {
2879    debug_assert!(seq.len() >= 32);
2880    #[cfg(all(target_arch = "x86_64", target_feature = "bmi2"))]
2881    {
2882        let mut word = 0u64;
2883        for chunk in 0..4 {
2884            let bytes = u64::from_le_bytes(
2885                seq[chunk * 8..chunk * 8 + 8]
2886                    .try_into()
2887                    .expect("eight bases"),
2888            );
2889            // Per byte: ((b >> 2) ^ (b >> 1)) & 0b11, evaluated eight at a time.
2890            let codes = ((bytes >> 2) ^ (bytes >> 1)) & 0x0303_0303_0303_0303;
2891            // Gather the eight 2-bit codes, lowest byte first, into 16 bits.
2892            let packed = unsafe { core::arch::x86_64::_pext_u64(codes, 0x0303_0303_0303_0303) };
2893            // The scalar form shifts the first base furthest left, and PEXT emits
2894            // the first (lowest-address) base in the least significant bits, so
2895            // reverse the 2-bit groups within this 16-bit lane.
2896            let reversed = reverse_2bit_groups_16(packed as u16);
2897            word = (word << 16) | u64::from(reversed);
2898        }
2899        word
2900    }
2901    #[cfg(not(all(target_arch = "x86_64", target_feature = "bmi2")))]
2902    {
2903        let mut word = 0u64;
2904        for &base in &seq[..32] {
2905            word = (word << 2) | u64::from(valid_ascii_base_bits(base));
2906        }
2907        word
2908    }
2909}
2910
2911/// Reverses the order of the eight 2-bit groups in a 16-bit value.
2912#[cfg(all(target_arch = "x86_64", target_feature = "bmi2"))]
2913#[inline]
2914fn reverse_2bit_groups_16(value: u16) -> u16 {
2915    let v = value as u32;
2916    // Swap adjacent 2-bit pairs, then nibbles, then bytes.
2917    let v = ((v & 0x3333) << 2) | ((v >> 2) & 0x3333);
2918    let v = ((v & 0x0f0f) << 4) | ((v >> 4) & 0x0f0f);
2919    (((v & 0x00ff) << 8) | ((v >> 8) & 0x00ff)) as u16
2920}
2921
2922#[inline]
2923fn pack_valid_label(seq: &[u8], label_words: usize) -> Result<[u64; 4], BucketError> {
2924    let mut words = [0u64; 4];
2925    if label_words > words.len() || seq.len() > label_words * 32 {
2926        return Err(BucketError::MalformedRecord);
2927    }
2928
2929    let full_words = seq.len() / 32;
2930    for word_idx in 0..full_words {
2931        words[word_idx] = pack_valid_word_32(&seq[word_idx * 32..]);
2932    }
2933    let tail = &seq[full_words * 32..];
2934    if !tail.is_empty() {
2935        let mut word = 0u64;
2936        for &base in tail {
2937            word = (word << 2) | u64::from(valid_ascii_base_bits(base));
2938        }
2939        words[full_words] = word << (2 * (32 - tail.len()));
2940    }
2941    Ok(words)
2942}
2943
2944struct BucketFile {
2945    path: PathBuf,
2946    file: File,
2947    records: u64,
2948    bytes_written: u64,
2949    record_size: u64,
2950    compressed: bool,
2951    interleaved_compression: bool,
2952    label_words: usize,
2953}
2954
2955/// Reusable staging for compressed bucket writes.
2956///
2957/// A `BucketFile` is opened per flush, so these buffers belong to the emitter
2958/// that owns the flush loop; held there, a 64 KiB block costs no allocation.
2959#[derive(Default)]
2960pub(crate) struct CompressionScratch {
2961    /// Fixed-size attributes, de-interleaved out of the record stream.
2962    attrs: Vec<u8>,
2963    /// Label words, de-interleaved out of the record stream.
2964    labels: Vec<u8>,
2965    /// LZ4 output for the attribute stream, or for the whole block.
2966    encoded: Vec<u8>,
2967    /// LZ4 output for the label stream.
2968    encoded_labels: Vec<u8>,
2969    /// Header and payload assembled for a single `write_all`.
2970    block: Vec<u8>,
2971}
2972
2973impl CompressionScratch {
2974    /// Compresses `input` into `out`, returning the encoded length.
2975    ///
2976    /// `out` only ever grows, so the zero-fill a resize implies is paid once
2977    /// per writer rather than once per block.
2978    fn encode(input: &[u8], out: &mut Vec<u8>) -> Result<usize, BucketError> {
2979        let bound = lz4_flex::block::get_maximum_output_size(input.len());
2980        if out.len() < bound {
2981            out.resize(bound, 0);
2982        }
2983        lz4_flex::block::compress_into(input, &mut out[..bound])
2984            .map_err(|_| BucketError::MalformedRecord)
2985    }
2986}
2987
2988/// Whether uncolored buckets compress attributes and labels as separate
2989/// streams, the way the colored path and C++ both do.
2990///
2991/// The interleaved default compresses whole records in one stream. The header
2992/// records which was used, so readers stay correct under either setting.
2993fn force_split_compression() -> bool {
2994    static SPLIT: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
2995    *SPLIT.get_or_init(|| std::env::var_os("CF3_RS_SPLIT_COMPRESSION").is_some())
2996}
2997
2998impl BucketFile {
2999    #[allow(clippy::too_many_arguments)]
3000    fn create(
3001        bucket_dir: &Path,
3002        k: u16,
3003        minimizer_len: u16,
3004        graph_count: usize,
3005        graph_id: usize,
3006        colored: bool,
3007        label_words: usize,
3008        compressed: bool,
3009    ) -> Result<Self, BucketError> {
3010        let path = bucket_dir.join(format!("{graph_id:05}.wsk"));
3011        let mut file = File::create(&path).map_err(|source| BucketError::Io {
3012            path: path.clone(),
3013            source,
3014        })?;
3015
3016        file.write_all(MAGIC).map_err(|source| BucketError::Io {
3017            path: path.clone(),
3018            source,
3019        })?;
3020        write_u16(&mut file, &path, k)?;
3021        write_u16(&mut file, &path, minimizer_len)?;
3022        write_u64(&mut file, &path, graph_count as u64)?;
3023        write_u64(&mut file, &path, graph_id as u64)?;
3024        let interleaved_compression = compressed && !colored && !force_split_compression();
3025        file.write_all(&[
3026            u8::from(colored),
3027            label_words as u8,
3028            if interleaved_compression {
3029                2
3030            } else {
3031                u8::from(compressed)
3032            },
3033            0,
3034            0,
3035            0,
3036        ])
3037        .map_err(|source| BucketError::Io {
3038            path: path.clone(),
3039            source,
3040        })?;
3041        write_u64(&mut file, &path, 0)?;
3042
3043        Ok(Self {
3044            path,
3045            file,
3046            records: 0,
3047            bytes_written: HEADER_LEN,
3048            record_size: record_size(colored, label_words) as u64,
3049            compressed,
3050            interleaved_compression,
3051            label_words,
3052        })
3053    }
3054
3055    fn open_existing(path: &Path, records: u64, bytes_written: u64) -> Result<Self, BucketError> {
3056        let mut file = OpenOptions::new()
3057            .read(true)
3058            .write(true)
3059            .open(path)
3060            .map_err(|source| BucketError::Io {
3061                path: path.to_path_buf(),
3062                source,
3063            })?;
3064        let header = read_header(&mut file, path)?;
3065        file.seek(SeekFrom::End(0))
3066            .map_err(|source| BucketError::Io {
3067                path: path.to_path_buf(),
3068                source,
3069            })?;
3070
3071        Ok(Self {
3072            path: path.to_path_buf(),
3073            file,
3074            records,
3075            bytes_written,
3076            record_size: record_size(header.colored, header.label_words) as u64,
3077            compressed: header.compressed,
3078            interleaved_compression: header.interleaved_compression,
3079            label_words: header.label_words,
3080        })
3081    }
3082
3083    fn write_records(
3084        &mut self,
3085        bytes: &[u8],
3086        records: u64,
3087        scratch: &mut CompressionScratch,
3088    ) -> Result<(u64, u64), BucketError> {
3089        debug_assert_eq!(bytes.len() as u64, records * self.record_size);
3090        let written = if self.compressed {
3091            self.write_compressed_block(bytes, records, scratch)?
3092        } else {
3093            self.file
3094                .write_all(bytes)
3095                .map_err(|source| BucketError::Io {
3096                    path: self.path.clone(),
3097                    source,
3098                })?;
3099            bytes.len() as u64
3100        };
3101        self.records = self
3102            .records
3103            .checked_add(records)
3104            .ok_or(BucketError::TooManyRecords)?;
3105        self.bytes_written += written;
3106        Ok((self.records, self.bytes_written))
3107    }
3108
3109    fn write_compressed_block(
3110        &mut self,
3111        bytes: &[u8],
3112        records: u64,
3113        scratch: &mut CompressionScratch,
3114    ) -> Result<u64, BucketError> {
3115        let len = encode_compressed_block(
3116            bytes,
3117            records,
3118            self.record_size,
3119            self.label_words,
3120            self.interleaved_compression,
3121            scratch,
3122        )?;
3123        if len == 0 {
3124            return Ok(0);
3125        }
3126        self.file
3127            .write_all(&scratch.block)
3128            .map_err(|source| BucketError::Io {
3129                path: self.path.clone(),
3130                source,
3131            })?;
3132        Ok(len as u64)
3133    }
3134
3135    fn flush(&mut self) -> Result<(), BucketError> {
3136        self.file.flush().map_err(|source| BucketError::Io {
3137            path: self.path.clone(),
3138            source,
3139        })
3140    }
3141
3142    fn finish(&mut self) -> Result<(), BucketError> {
3143        write_record_count(&mut self.file, &self.path, self.records)?;
3144        self.file
3145            .seek(SeekFrom::End(0))
3146            .map_err(|source| BucketError::Io {
3147                path: self.path.clone(),
3148                source,
3149            })?;
3150        self.file.flush().map_err(|source| BucketError::Io {
3151            path: self.path.clone(),
3152            source,
3153        })
3154    }
3155
3156    fn finish_closed(path: &Path, records: u64) -> Result<(), BucketError> {
3157        let mut file = OpenOptions::new()
3158            .read(true)
3159            .write(true)
3160            .open(path)
3161            .map_err(|source| BucketError::Io {
3162                path: path.to_path_buf(),
3163                source,
3164            })?;
3165        write_record_count(&mut file, path, records)?;
3166        file.flush().map_err(|source| BucketError::Io {
3167            path: path.to_path_buf(),
3168            source,
3169        })
3170    }
3171}
3172
3173fn write_record_count(file: &mut File, path: &Path, records: u64) -> Result<(), BucketError> {
3174    file.seek(SeekFrom::Start(RECORD_COUNT_OFFSET))
3175        .map_err(|source| BucketError::Io {
3176            path: path.to_path_buf(),
3177            source,
3178        })?;
3179    write_u64(file, path, records)
3180}
3181
3182fn label_word_count(k: u16, minimizer_len: u16) -> usize {
3183    let max_weak_superkmer_len = 2 * (usize::from(k) - 1) - usize::from(minimizer_len) + 2;
3184    max_weak_superkmer_len.div_ceil(32)
3185}
3186
3187fn record_size(colored: bool, label_words: usize) -> usize {
3188    (if colored { 4 } else { 2 }) + label_words * 8
3189}
3190
3191fn pack_uncolored_attr(len: usize, left_discontinuous: bool, right_discontinuous: bool) -> u32 {
3192    (len as u32) | ((left_discontinuous as u32) << 8) | ((right_discontinuous as u32) << 9)
3193}
3194
3195fn pack_colored_attr(
3196    len: usize,
3197    source_id: u32,
3198    left_discontinuous: bool,
3199    right_discontinuous: bool,
3200) -> u32 {
3201    pack_uncolored_attr(len, left_discontinuous, right_discontinuous) | (source_id << 10)
3202}
3203
3204fn decode_label_into(words: &[u64], len: usize, seq: &mut Vec<u8>) -> Result<(), BucketError> {
3205    if len > words.len() * 32 {
3206        return Err(BucketError::MalformedRecord);
3207    }
3208
3209    seq.clear();
3210    seq.reserve(len);
3211    for idx in 0..len {
3212        let word_idx = idx / 32;
3213        let shift = 2 * (31 - (idx % 32));
3214        let base = match ((words[word_idx] >> shift) & 0b11) as u8 {
3215            0 => Base::A,
3216            1 => Base::C,
3217            2 => Base::G,
3218            3 => Base::T,
3219            _ => unreachable!(),
3220        };
3221        seq.push(base.to_ascii());
3222    }
3223    Ok(())
3224}
3225
3226fn read_header(file: &mut impl Read, path: &Path) -> Result<BucketHeader, BucketError> {
3227    let mut magic = [0u8; 8];
3228    file.read_exact(&mut magic)
3229        .map_err(|source| BucketError::Io {
3230            path: path.to_path_buf(),
3231            source,
3232        })?;
3233    if &magic != MAGIC {
3234        return Err(BucketError::BadMagic(path.to_path_buf()));
3235    }
3236
3237    let k = read_u16(file, path)?;
3238    let minimizer_len = read_u16(file, path)?;
3239    let graph_count = usize::try_from(read_u64(file, path)?)
3240        .map_err(|_| BucketError::MalformedHeader(path.to_path_buf()))?;
3241    let graph_id = usize::try_from(read_u64(file, path)?)
3242        .map_err(|_| BucketError::MalformedHeader(path.to_path_buf()))?;
3243
3244    let mut flags = [0u8; 6];
3245    file.read_exact(&mut flags)
3246        .map_err(|source| BucketError::Io {
3247            path: path.to_path_buf(),
3248            source,
3249        })?;
3250    let colored = match flags[0] {
3251        0 => false,
3252        1 => true,
3253        _ => return Err(BucketError::MalformedHeader(path.to_path_buf())),
3254    };
3255    let label_words = flags[1] as usize;
3256    if label_words == 0 || label_words != label_word_count(k, minimizer_len) {
3257        return Err(BucketError::MalformedHeader(path.to_path_buf()));
3258    }
3259    let (compressed, interleaved_compression) = match flags[2] {
3260        0 => (false, false),
3261        1 => (true, false),
3262        2 if !colored => (true, true),
3263        _ => return Err(BucketError::MalformedHeader(path.to_path_buf())),
3264    };
3265    if flags[3..].iter().any(|&b| b != 0) {
3266        return Err(BucketError::MalformedHeader(path.to_path_buf()));
3267    }
3268    let records = read_u64(file, path)?;
3269
3270    if graph_id >= graph_count {
3271        return Err(BucketError::MalformedHeader(path.to_path_buf()));
3272    }
3273
3274    Ok(BucketHeader {
3275        k,
3276        minimizer_len,
3277        graph_count,
3278        graph_id,
3279        colored,
3280        compressed,
3281        interleaved_compression,
3282        label_words,
3283        records,
3284    })
3285}
3286
3287fn read_u16(file: &mut impl Read, path: &Path) -> Result<u16, BucketError> {
3288    let mut bytes = [0u8; 2];
3289    file.read_exact(&mut bytes)
3290        .map_err(|source| BucketError::Io {
3291            path: path.to_path_buf(),
3292            source,
3293        })?;
3294    Ok(u16::from_le_bytes(bytes))
3295}
3296
3297fn read_u64(file: &mut impl Read, path: &Path) -> Result<u64, BucketError> {
3298    let mut bytes = [0u8; 8];
3299    file.read_exact(&mut bytes)
3300        .map_err(|source| BucketError::Io {
3301            path: path.to_path_buf(),
3302            source,
3303        })?;
3304    Ok(u64::from_le_bytes(bytes))
3305}
3306
3307fn write_u16(file: &mut File, path: &Path, value: u16) -> Result<(), BucketError> {
3308    file.write_all(&value.to_le_bytes())
3309        .map_err(|source| BucketError::Io {
3310            path: path.to_path_buf(),
3311            source,
3312        })
3313}
3314
3315fn write_u64(file: &mut File, path: &Path, value: u64) -> Result<(), BucketError> {
3316    file.write_all(&value.to_le_bytes())
3317        .map_err(|source| BucketError::Io {
3318            path: path.to_path_buf(),
3319            source,
3320        })
3321}
3322
3323#[derive(Debug)]
3324pub enum BucketError {
3325    Io {
3326        path: PathBuf,
3327        source: std::io::Error,
3328    },
3329    GraphCountTooLarge(usize),
3330    InvalidGraphId(usize),
3331    LabelTooLong(usize),
3332    MissingSourceId,
3333    SourceIdTooLarge(u32),
3334    InvalidBase(u8),
3335    TooManyRecords,
3336    BadMagic(PathBuf),
3337    MalformedHeader(PathBuf),
3338    MalformedManifest(PathBuf),
3339    MalformedRecord,
3340    RecordGraphMismatch {
3341        expected: usize,
3342        got: usize,
3343    },
3344    WorkerPanic,
3345}
3346
3347impl std::fmt::Display for BucketError {
3348    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
3349        match self {
3350            Self::Io { path, source } => write!(f, "{}: {source}", path.display()),
3351            Self::GraphCountTooLarge(count) => write!(f, "graph count is too large: {count}"),
3352            Self::InvalidGraphId(graph_id) => write!(f, "invalid graph id: {graph_id}"),
3353            Self::LabelTooLong(len) => write!(f, "weak super-kmer label is too long: {len}"),
3354            Self::MissingSourceId => write!(f, "colored bucket record is missing source id"),
3355            Self::SourceIdTooLarge(source_id) => {
3356                write!(
3357                    f,
3358                    "source id exceeds 21-bit colored bucket limit: {source_id}"
3359                )
3360            }
3361            Self::InvalidBase(b) => write!(f, "invalid base in bucket label: '{}'", *b as char),
3362            Self::TooManyRecords => write!(f, "too many weak super-kmer records"),
3363            Self::BadMagic(path) => write!(
3364                f,
3365                "not a CF3 Rust weak-superkmer bucket: {}",
3366                path.display()
3367            ),
3368            Self::MalformedHeader(path) => {
3369                write!(
3370                    f,
3371                    "malformed weak-superkmer bucket header: {}",
3372                    path.display()
3373                )
3374            }
3375            Self::MalformedManifest(path) => {
3376                write!(
3377                    f,
3378                    "malformed weak-superkmer bucket manifest: {}",
3379                    path.display()
3380                )
3381            }
3382            Self::MalformedRecord => write!(f, "malformed weak-superkmer bucket record"),
3383            Self::RecordGraphMismatch { expected, got } => {
3384                write!(
3385                    f,
3386                    "bucket record graph id mismatch: expected {expected}, got {got}"
3387                )
3388            }
3389            Self::WorkerPanic => write!(f, "bucket worker thread panicked"),
3390        }
3391    }
3392}
3393
3394impl std::error::Error for BucketError {}
3395
3396#[cfg(test)]
3397mod tests {
3398    /// Reclaim must actually return blocks, not merely be called.
3399    ///
3400    /// A failed punch is ignored by design -- the container is unlinked
3401    /// wholesale at the end regardless -- which means a filesystem or platform
3402    /// where it silently does nothing costs peak disk with no other symptom.
3403    /// That is exactly what happened once already: deferring reclaim left the
3404    /// work directory 24.5 GB larger and presented as an unexplained number.
3405    /// This asserts the blocks come back, so the macOS `fcntl(F_PUNCHHOLE)`
3406    /// path is verified by CI rather than assumed from the fact that it
3407    /// compiles.
3408    #[test]
3409    fn releasing_segments_returns_blocks_to_the_filesystem() {
3410        use std::os::unix::fs::MetadataExt;
3411
3412        let dir = std::env::temp_dir().join(format!(
3413            "cf3-punch-{}-{:?}",
3414            std::process::id(),
3415            std::thread::current().id()
3416        ));
3417        let _ = fs::remove_dir_all(&dir);
3418        fs::create_dir_all(&dir).unwrap();
3419
3420        let containers = BucketContainers::create(&dir, 1).unwrap();
3421        let segment = containers.segment_bytes();
3422        let segments: Vec<u32> = (0..8).collect();
3423        let payload = vec![0xA5u8; segment as usize];
3424        for &index in &segments {
3425            containers
3426                .write_at(0, u64::from(index) * segment, &payload)
3427                .unwrap();
3428        }
3429
3430        let path = dir.join("00000.wskc");
3431        let allocated = || fs::metadata(&path).unwrap().blocks() * 512;
3432        let before = allocated();
3433        assert!(
3434            before >= segments.len() as u64 * segment,
3435            "expected {} bytes of blocks before punching, saw {before}",
3436            segments.len() as u64 * segment
3437        );
3438
3439        containers.release_segments(0, &segments);
3440        let after = allocated();
3441
3442        assert!(
3443            after < before / 2,
3444            "punching {} segments freed {} of {before} bytes; this filesystem may \
3445             not support hole punching, in which case consumed bucket space is \
3446             held until the build ends and peak disk is higher than documented",
3447            segments.len(),
3448            before - after
3449        );
3450        // The file keeps its length; only the blocks behind it go away.
3451        assert_eq!(
3452            fs::metadata(&path).unwrap().len(),
3453            segments.len() as u64 * segment
3454        );
3455
3456        fs::remove_dir_all(&dir).unwrap();
3457    }
3458
3459    #[test]
3460    fn packed_word_matches_scalar_reference() {
3461        fn scalar(seq: &[u8]) -> u64 {
3462            let mut word = 0u64;
3463            for &base in &seq[..32] {
3464                word = (word << 2) | u64::from(valid_ascii_base_bits(base));
3465            }
3466            word
3467        }
3468        let alphabet = b"ACGT";
3469        // Deterministic pseudo-random coverage plus structured edge cases.
3470        let mut state = 0x1234_5678_9abc_def0u64;
3471        for case in 0..2048 {
3472            let mut seq = [0u8; 32];
3473            for (i, slot) in seq.iter_mut().enumerate() {
3474                state = state
3475                    .wrapping_mul(6364136223846793005)
3476                    .wrapping_add(1442695040888963407);
3477                *slot = if case < 4 {
3478                    alphabet[(case + i) % 4]
3479                } else {
3480                    alphabet[(state >> 33) as usize % 4]
3481                };
3482            }
3483            assert_eq!(
3484                pack_valid_word_32(&seq),
3485                scalar(&seq),
3486                "mismatch for {:?}",
3487                std::str::from_utf8(&seq).unwrap()
3488            );
3489        }
3490    }
3491
3492    use super::*;
3493
3494    #[test]
3495    fn deferred_uncolored_atlas_chunks_preserve_graph_buckets() {
3496        let dir = std::env::temp_dir().join(format!(
3497            "cf3-uncolored-atlas-bucket-{}-{:?}",
3498            std::process::id(),
3499            std::thread::current().id()
3500        ));
3501        fs::create_dir_all(&dir).unwrap();
3502        let mut params = BuildParams::new(crate::GraphInput::References, "test".to_string());
3503        params.k = 31;
3504        params.minimizer_len = 15;
3505        params.vertex_partitions = 1;
3506        params.threads = 2;
3507        params.work_dir = dir.to_string_lossy().into_owned();
3508        let sink = SharedBucketSink::create(&params, ATLAS_GRAPH_COUNT + 1).unwrap();
3509        let expected = [(0, b'A'), (7, b'C'), (ATLAS_GRAPH_COUNT, b'G')];
3510        let mut emitters = Vec::new();
3511        for repeat in 0..2 {
3512            let mut emitter = sink.deferred_uncolored_emitter();
3513            for &(graph_id, base) in &expected {
3514                for _ in 0..(repeat + 1) {
3515                    emitter
3516                        .add_valid(
3517                            &WeakSuperKmer {
3518                                graph_id,
3519                                offset: 0,
3520                                len: 31,
3521                                source_id: None,
3522                                left_discontinuous: false,
3523                                right_discontinuous: false,
3524                            },
3525                            &[base; 31],
3526                        )
3527                        .unwrap();
3528                }
3529            }
3530            emitters.push(emitter);
3531        }
3532        sink.flush_uncolored_emitters(emitters).unwrap();
3533        let stats = sink.finish().unwrap();
3534
3535        let (store, entries) = BucketStore::open_dir(&stats.bucket_dir).unwrap();
3536        for &(graph_id, base) in &expected {
3537            let entry = entries
3538                .iter()
3539                .find(|entry| entry.graph_id == graph_id)
3540                .expect("bucket in manifest");
3541            let mut reader = store.reader(entry).unwrap();
3542            let mut record = BucketRecord::default();
3543            let mut count = 0;
3544            while reader.next_record_into(&mut record).unwrap() {
3545                assert_eq!(record.graph_id, graph_id);
3546                assert_eq!(record.label, vec![base; 31]);
3547                count += 1;
3548            }
3549            assert_eq!(count, 3);
3550        }
3551        fs::remove_dir_all(dir).unwrap();
3552    }
3553
3554    #[test]
3555    fn colored_atlas_windows_preserve_global_source_order() {
3556        let dir = std::env::temp_dir().join(format!(
3557            "cf3-colored-bucket-{}-{:?}",
3558            std::process::id(),
3559            std::thread::current().id()
3560        ));
3561        fs::create_dir_all(&dir).unwrap();
3562        let mut params = BuildParams::new(crate::GraphInput::References, "test".to_string());
3563        params.color = true;
3564        params.k = 31;
3565        params.minimizer_len = 15;
3566        params.vertex_partitions = 1;
3567        params.threads = 1;
3568        params.work_dir = dir.to_string_lossy().into_owned();
3569        let sink = SharedBucketSink::create(&params, 1).unwrap();
3570
3571        for (source_min, source_max, sources) in
3572            [(1, 3, vec![3, 1, 3, 2, 1]), (4, 6, vec![6, 4, 5, 4])]
3573        {
3574            let mut emitter = sink.emitter();
3575            for source_id in sources {
3576                emitter
3577                    .add_valid(
3578                        &WeakSuperKmer {
3579                            graph_id: 0,
3580                            offset: 0,
3581                            len: 31,
3582                            source_id: Some(source_id),
3583                            left_discontinuous: false,
3584                            right_discontinuous: false,
3585                        },
3586                        &[b'A'; 31],
3587                    )
3588                    .unwrap();
3589            }
3590            emitter.finish().unwrap();
3591            sink.flush_colored_window(source_min, source_max).unwrap();
3592        }
3593        let stats = sink.finish().unwrap();
3594
3595        let (store, entries) = BucketStore::open_dir(&stats.bucket_dir).unwrap();
3596        let entry = entries
3597            .iter()
3598            .find(|entry| entry.graph_id == 0)
3599            .expect("bucket in manifest");
3600        let mut reader = store.reader(entry).unwrap();
3601        assert!(reader.header().compressed);
3602        let mut sources = Vec::new();
3603        let mut record = BucketPackedRecord::default();
3604        while reader.next_packed_record_into(&mut record).unwrap() {
3605            sources.push(record.source_id.unwrap());
3606        }
3607        assert_eq!(sources, [1, 1, 2, 3, 3, 4, 4, 5, 6]);
3608
3609        // Clipping the container leaves the manifest claiming a length the
3610        // chain can no longer supply, which is the container-shaped version of
3611        // a bucket file cut short by a killed run.
3612        let container_path = stats.bucket_dir.join("00000.wskc");
3613        let len = fs::metadata(&container_path).unwrap().len();
3614        OpenOptions::new()
3615            .write(true)
3616            .open(&container_path)
3617            .unwrap()
3618            .set_len(len - 1)
3619            .unwrap();
3620        let (store, entries) = BucketStore::open_dir(&stats.bucket_dir).unwrap();
3621        let entry = entries
3622            .iter()
3623            .find(|entry| entry.graph_id == 0)
3624            .expect("bucket in manifest");
3625        let mut truncated = store.reader(entry).unwrap();
3626        let mut record = BucketPackedRecord::default();
3627        let mut failed = false;
3628        loop {
3629            match truncated.next_packed_record_into(&mut record) {
3630                Ok(true) => {}
3631                Ok(false) => break,
3632                Err(_) => {
3633                    failed = true;
3634                    break;
3635                }
3636            }
3637        }
3638        assert!(failed, "truncated compressed block must be rejected");
3639        fs::remove_dir_all(dir).unwrap();
3640    }
3641
3642    #[test]
3643    fn colored_worker_tails_preserve_cpp_worker_order() {
3644        let dir = std::env::temp_dir().join(format!(
3645            "cf3-colored-worker-tails-{}-{:?}",
3646            std::process::id(),
3647            std::thread::current().id()
3648        ));
3649        fs::create_dir_all(&dir).unwrap();
3650        let mut params = BuildParams::new(crate::GraphInput::References, "test".to_string());
3651        params.color = true;
3652        params.k = 31;
3653        params.minimizer_len = 15;
3654        params.threads = 2;
3655        params.work_dir = dir.to_string_lossy().into_owned();
3656        let sink = SharedBucketSink::create(&params, 1).unwrap();
3657
3658        let mut emitters = Vec::new();
3659        for sources in [[3, 1], [4, 2]] {
3660            let mut emitter = sink.emitter();
3661            for source_id in sources {
3662                emitter
3663                    .add_valid(
3664                        &WeakSuperKmer {
3665                            graph_id: 0,
3666                            offset: 0,
3667                            len: 31,
3668                            source_id: Some(source_id),
3669                            left_discontinuous: false,
3670                            right_discontinuous: false,
3671                        },
3672                        &[b'A'; 31],
3673                    )
3674                    .unwrap();
3675            }
3676            emitters.push(emitter);
3677        }
3678        sink.flush_colored_emitters(emitters).unwrap();
3679        let stats = sink.finish().unwrap();
3680
3681        let (store, entries) = BucketStore::open_dir(&stats.bucket_dir).unwrap();
3682        let entry = entries
3683            .iter()
3684            .find(|entry| entry.graph_id == 0)
3685            .expect("bucket in manifest");
3686        let mut reader = store.reader(entry).unwrap();
3687        let mut sources = Vec::new();
3688        let mut record = BucketPackedRecord::default();
3689        while reader.next_packed_record_into(&mut record).unwrap() {
3690            sources.push(record.source_id.unwrap());
3691        }
3692        assert_eq!(sources, [3, 1, 4, 2]);
3693        fs::remove_dir_all(dir).unwrap();
3694    }
3695}