Skip to main content

cuttlefish_rs/
colored.rs

1//! End-to-end colored graph construction.
2//!
3//! This module is the high-level colored build driver. It connects partition
4//! buckets to local contraction, external discontinuity processing, FASTA
5//! emission, and color-repository metadata.
6
7use crate::color::ColorError;
8use crate::discontinuity::{
9    DiscontinuityInputError, SerialCollationError, SerialUncoloredCollator,
10    emit_colored_external_discontinuity_inputs_with_threads_in_dir, report_process_memory,
11    spawn_background_dir_removal, trim_process_allocations,
12};
13use crate::input::{InputError, expand_input_paths};
14use crate::kmer::KmerError;
15use crate::params::BuildParams;
16use std::path::{Path, PathBuf};
17use std::time::Instant;
18
19/// Summary of a completed colored graph build.
20#[derive(Debug, Clone, PartialEq, Eq)]
21pub struct ColoredBuildStats {
22    pub input_buckets: usize,
23    pub bucket_records: u64,
24    pub unitigs: u64,
25    pub unitig_bases: u64,
26    pub output_path: PathBuf,
27    pub color_repository: PathBuf,
28}
29
30/// Builds a colored compacted de Bruijn graph from emitted partition buckets.
31///
32/// `K` must equal `params.k`. The output FASTA is written beside
33/// `params.output_prefix`; external intermediates are written under
34/// `params.work_dir` and removed unless intermediate retention is enabled.
35pub fn build_colored_from_buckets<const K: usize>(
36    params: &BuildParams,
37    bucket_dir: impl AsRef<Path>,
38) -> Result<ColoredBuildStats, ColoredBuildError> {
39    if !params.color {
40        return Err(ColoredBuildError::ColorRequired);
41    }
42    let bucket_dir = bucket_dir.as_ref().to_path_buf();
43    let output_name = Path::new(&params.output_prefix)
44        .file_name()
45        .and_then(|name| name.to_str())
46        .filter(|name| !name.is_empty())
47        .unwrap_or("cuttlefish3");
48    let work_dir = PathBuf::from(&params.work_dir);
49    let label_path = work_dir.join(format!("{output_name}.cf3rs.lmtig-labels"));
50    let color_path = work_dir.join(format!("{output_name}.cf3rs.colors"));
51    // The colour repository accompanies the FASTA as part of the build's
52    // output; only the run-sidecar above is scratch.
53    let color_repository_dir =
54        PathBuf::from(format!("{}.cf3rs.color-repository", params.output_prefix));
55    // The hybrid colour encoding selects its regime from the total source
56    // count, so it has to be known before any colour is written.
57    let sources = expand_input_paths(params)?;
58    // Source IDs are one-based, so the regime thresholds and the bitmap width
59    // are sized by the largest ID plus one, as C++ does with max_source_id + 1.
60    let num_colors =
61        u32::try_from(sources.len() + 1).map_err(|_| ColoredBuildError::ColorRequired)?;
62
63    let local_started = Instant::now();
64    report_process_memory("before colored local contraction");
65    let local_threads = params.local_workers();
66    eprintln!("cuttlefish: colored local contraction using {local_threads} worker(s)");
67    let mut inputs = emit_colored_external_discontinuity_inputs_with_threads_in_dir::<K>(
68        &bucket_dir,
69        params.cutoff(),
70        local_threads,
71        &label_path,
72        &color_path,
73        &color_repository_dir,
74        num_colors,
75    )?;
76    eprintln!(
77        "cuttlefish: colored local contraction completed in {:.3}s; {} local unitig(s)",
78        local_started.elapsed().as_secs_f64(),
79        inputs.stats.local_unitigs,
80    );
81    // With buckets in containers this is a bulk delete of 129 files holding
82    // hundreds of gigabytes, not the single leftover manifest the per-file
83    // layout left behind -- local contraction unlinked those as it consumed
84    // them. So it goes to the background unlinkers, like the edge-matrix and
85    // expansion directories, and is joined before the build returns.
86    let bucket_reclaim = (std::env::var_os("CF3_RS_KEEP_INTERMEDIATES").is_none())
87        .then(|| spawn_background_dir_removal(bucket_dir.clone()));
88    trim_process_allocations();
89
90    let output_path = PathBuf::from(format!("{}.fa", params.output_prefix));
91    let coord_dir = work_dir.join(format!("{output_name}.cf3rs.stitch-coords"));
92    let final_dir = work_dir.join(format!("{output_name}.cf3rs.final-unitigs"));
93    let post_local_threads = params.post_local_workers();
94    eprintln!("cuttlefish: colored collation using {post_local_threads} worker(s)");
95    let stats = SerialUncoloredCollator::collate_external_stitched_to_fasta_with_threads_in_dir(
96        &mut inputs,
97        post_local_threads,
98        &coord_dir,
99        &final_dir,
100        &output_path,
101    )?;
102    report_process_memory("after colored collation");
103    let color_repository = inputs
104        .color_repository()
105        .ok_or(ColoredBuildError::MissingColorArtifacts)?
106        .clone();
107    color_repository.write_metadata(params.k, &output_path, &sources)?;
108    if let Some(handle) = bucket_reclaim {
109        let _ = handle.join();
110    }
111    Ok(ColoredBuildStats {
112        input_buckets: inputs.stats.input_buckets,
113        bucket_records: inputs.stats.weak_superkmers,
114        unitigs: stats.emitted_unitigs,
115        unitig_bases: stats.emitted_bases,
116        output_path,
117        color_repository: color_repository.dir,
118    })
119}
120
121#[derive(Debug)]
122pub enum ColoredBuildError {
123    ColorRequired,
124    MissingColorArtifacts,
125    SourceInput(InputError),
126    Input(DiscontinuityInputError),
127    Collation(SerialCollationError),
128    Kmer(KmerError),
129    Color(ColorError),
130    Cleanup {
131        path: PathBuf,
132        source: std::io::Error,
133    },
134}
135
136impl From<DiscontinuityInputError> for ColoredBuildError {
137    fn from(value: DiscontinuityInputError) -> Self {
138        Self::Input(value)
139    }
140}
141
142impl From<SerialCollationError> for ColoredBuildError {
143    fn from(value: SerialCollationError) -> Self {
144        Self::Collation(value)
145    }
146}
147
148impl From<InputError> for ColoredBuildError {
149    fn from(value: InputError) -> Self {
150        Self::SourceInput(value)
151    }
152}
153
154impl From<ColorError> for ColoredBuildError {
155    fn from(value: ColorError) -> Self {
156        Self::Color(value)
157    }
158}
159
160impl std::fmt::Display for ColoredBuildError {
161    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
162        match self {
163            Self::ColorRequired => write!(f, "colored build requires --color"),
164            Self::MissingColorArtifacts => {
165                write!(f, "colored build did not produce color artifacts")
166            }
167            Self::SourceInput(err) => write!(f, "{err}"),
168            Self::Input(err) => write!(f, "{err}"),
169            Self::Collation(err) => write!(f, "{err}"),
170            Self::Kmer(err) => write!(f, "{err}"),
171            Self::Color(err) => write!(f, "{err}"),
172            Self::Cleanup { path, source } => write!(f, "{}: {source}", path.display()),
173        }
174    }
175}
176
177impl std::error::Error for ColoredBuildError {}