cuttlefish_rs/
state.rs

1//! Compact vertex state and colored-coordinate representations.
2//!
3//! These types are used in the hottest local-contraction tables. Their sizes
4//! and bit allocations are deliberate; avoid adding fields without measuring
5//! memory bandwidth and peak RSS at scale.
6
7use crate::Side;
8use crate::dna::Base;
9use xxhash_rust::xxh3::xxh3_64;
10
11#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
12pub struct EdgeFrequency {
13    packed: u32,
14}
15
16impl EdgeFrequency {
17    const MAX: u32 = 0xF;
18
19    #[inline]
20    pub fn add_edge(&mut self, side: Side, edge: Base) {
21        assert!(matches!(edge, Base::A | Base::C | Base::G | Base::T));
22        let off = Self::offset(side, edge);
23        let mask = Self::MAX << off;
24        let cur = (self.packed & mask) >> off;
25        if cur < Self::MAX {
26            self.packed = (self.packed & !mask) | ((cur + 1) << off);
27        }
28    }
29
30    #[inline]
31    pub fn edge_count(&self, side: Side, cutoff: u32) -> u32 {
32        let side_off = Self::side_offset(side);
33        let packed = (self.packed >> side_off) & 0xffff;
34        u32::from((packed & 0x000f) >= cutoff)
35            + u32::from(((packed >> 4) & 0x000f) >= cutoff)
36            + u32::from(((packed >> 8) & 0x000f) >= cutoff)
37            + u32::from(((packed >> 12) & 0x000f) >= cutoff)
38    }
39
40    #[inline]
41    pub fn edge_at(&self, side: Side, cutoff: u32) -> Base {
42        let side_off = Self::side_offset(side);
43        let packed = (self.packed >> side_off) & 0xffff;
44        let edge_a = u32::from((packed & 0x000f) >= cutoff);
45        let edge_c = u32::from(((packed >> 4) & 0x000f) >= cutoff);
46        let edge_g = u32::from(((packed >> 8) & 0x000f) >= cutoff);
47        let edge_t = u32::from(((packed >> 12) & 0x000f) >= cutoff);
48        match edge_a + edge_c + edge_g + edge_t {
49            0 => Base::E,
50            1 => match edge_c + 2 * edge_g + 3 * edge_t {
51                0 => Base::A,
52                1 => Base::C,
53                2 => Base::G,
54                3 => Base::T,
55                _ => unreachable!(),
56            },
57            _ => Base::N,
58        }
59    }
60
61    #[inline]
62    pub fn frequency(&self, side: Side, base_bits: u32) -> u32 {
63        let off = Self::side_offset(side) + 4 * base_bits;
64        (self.packed >> off) & Self::MAX
65    }
66
67    #[inline]
68    fn side_offset(side: Side) -> u32 {
69        match side {
70            Side::Front => 0,
71            Side::Back => 16,
72        }
73    }
74
75    #[inline]
76    fn offset(side: Side, edge: Base) -> u32 {
77        Self::side_offset(side) + 4 * edge.bits() as u32
78    }
79}
80
81#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
82pub struct VertexState {
83    edges: EdgeFrequency,
84    flags: u32,
85    color_hash: u64,
86}
87
88impl VertexState {
89    const VISITED: u32 = 1 << 0;
90    const DISC_FRONT: u32 = 1 << 1;
91    const DISC_BACK: u32 = 1 << 2;
92    /// Largest source ID representable in the packed last-source field.
93    ///
94    /// Colored builds track the previously seen source per vertex in 21 bits of
95    /// `flags`; partitioning rejects larger source sets before contraction so
96    /// this bound is never reached at run time.
97    pub const MAX_SOURCE_ID: u32 = 0x1F_FFFF;
98
99    const SOURCE_SHIFT: u32 = 11;
100    const SOURCE_MASK: u32 = 0x1F_FFFF << Self::SOURCE_SHIFT;
101
102    #[inline(always)]
103    pub fn update_edges(&mut self, front: Base, back: Base) {
104        if front != Base::E {
105            self.edges.add_edge(Side::Front, front);
106        }
107        if back != Base::E {
108            self.edges.add_edge(Side::Back, back);
109        }
110    }
111
112    #[inline]
113    pub fn edge_at(&self, side: Side, cutoff: u32) -> Base {
114        self.edges.edge_at(side, cutoff)
115    }
116
117    #[inline]
118    pub fn is_branching_side(&self, side: Side, cutoff: u32) -> bool {
119        self.edges.edge_count(side, cutoff) > 1
120    }
121
122    #[inline]
123    pub fn is_empty_side(&self, side: Side, cutoff: u32) -> bool {
124        self.edges.edge_count(side, cutoff) == 0
125    }
126
127    #[inline]
128    pub fn is_isolated(&self, cutoff: u32) -> bool {
129        self.is_empty_side(Side::Front, cutoff) && self.is_empty_side(Side::Back, cutoff)
130    }
131
132    #[inline]
133    pub fn is_discontinuity(&self) -> bool {
134        self.is_discontinuous(Side::Front) || self.is_discontinuous(Side::Back)
135    }
136
137    #[inline]
138    pub fn mark_visited(&mut self) {
139        self.flags |= Self::VISITED;
140    }
141
142    #[inline]
143    pub fn is_visited(&self) -> bool {
144        self.flags & Self::VISITED != 0
145    }
146
147    #[inline]
148    pub fn mark_discontinuous(&mut self, side: Side) {
149        self.flags |= match side {
150            Side::Front => Self::DISC_FRONT,
151            Side::Back => Self::DISC_BACK,
152        };
153    }
154
155    #[inline]
156    pub fn is_discontinuous(&self, side: Side) -> bool {
157        self.flags
158            & match side {
159                Side::Front => Self::DISC_FRONT,
160                Side::Back => Self::DISC_BACK,
161            }
162            != 0
163    }
164
165    #[inline(always)]
166    pub fn add_source(&mut self, source: u32) {
167        self.add_source_hashed(source, source_hash(source));
168    }
169
170    #[inline(always)]
171    pub fn add_source_hashed(&mut self, source: u32, source_hash: u64) {
172        debug_assert!(
173            source <= Self::MAX_SOURCE_ID,
174            "source IDs are bounded during partitioning"
175        );
176        let last = (self.flags & Self::SOURCE_MASK) >> Self::SOURCE_SHIFT;
177        if source != last {
178            self.color_hash = hash_combine(self.color_hash, source_hash);
179            self.flags = (self.flags & !Self::SOURCE_MASK) | (source << Self::SOURCE_SHIFT);
180        }
181    }
182
183    #[inline]
184    pub fn color_hash(&self) -> u64 {
185        self.color_hash
186    }
187}
188
189#[inline]
190pub fn source_hash(source: u32) -> u64 {
191    debug_assert!(source > 0 && source < (1 << 21));
192    xxh3_64(&source.to_le_bytes()[..3])
193}
194
195#[inline]
196pub fn hash_combine(lhs: u64, rhs: u64) -> u64 {
197    lhs ^ rhs
198}
199
200#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
201/// Compact location of a deduplicated source set in the color repository.
202///
203/// Published coordinates use 8 worker bits and 32 worker-local index bits.
204/// The high bit is reserved for concurrent insertion state.
205pub struct ColorCoordinate(u64);
206
207impl ColorCoordinate {
208    const IN_PROCESS: u64 = 1u64 << 63;
209    const INDEX_SHIFT: u32 = 8;
210
211    pub fn in_process(worker: u64) -> Self {
212        assert!(worker < (1u64 << Self::INDEX_SHIFT));
213        Self(Self::IN_PROCESS | worker)
214    }
215
216    pub fn discovered(worker: u64, index: u64) -> Self {
217        assert!(worker < (1u64 << Self::INDEX_SHIFT));
218        assert!(index < (1u64 << 32));
219        Self(worker | (index << Self::INDEX_SHIFT))
220    }
221
222    pub fn from_u40(value: u64) -> Self {
223        assert!(value < (1u64 << 40));
224        Self(value)
225    }
226
227    #[inline]
228    pub fn is_in_process(self) -> bool {
229        self.0 & Self::IN_PROCESS != 0
230    }
231
232    #[inline]
233    pub fn processing_worker(self) -> u64 {
234        assert!(self.is_in_process());
235        self.0 & !Self::IN_PROCESS
236    }
237
238    #[inline]
239    pub fn as_u40(self) -> u64 {
240        assert!(self.0 < (1u64 << 40));
241        self.0
242    }
243
244    #[inline]
245    pub fn worker(self) -> usize {
246        assert!(!self.is_in_process());
247        (self.0 & 0xff) as usize
248    }
249
250    #[inline]
251    pub fn index(self) -> u32 {
252        assert!(!self.is_in_process());
253        (self.0 >> Self::INDEX_SHIFT) as u32
254    }
255}
256
257#[repr(transparent)]
258#[derive(Debug, Clone, Copy, PartialEq, Eq)]
259/// A packed positional color run.
260///
261/// The low 24 bits hold a unitig vertex offset and the upper 40 bits hold a
262/// [`ColorCoordinate`]. This transparent 64-bit layout is written directly to
263/// private intermediate streams.
264pub struct UnitigColor(u64);
265
266impl UnitigColor {
267    /// Packs a run starting at `offset` and referring to `coord`.
268    pub fn new(offset: u32, coord: ColorCoordinate) -> Self {
269        assert!(offset <= 0xFF_FFFF);
270        Self((coord.as_u40() << 24) | offset as u64)
271    }
272
273    /// Returns the zero-based vertex offset where this color run begins.
274    #[inline]
275    pub fn offset(self) -> u32 {
276        (self.0 & 0xFF_FFFF) as u32
277    }
278
279    /// Returns the raw 40-bit color-repository coordinate.
280    #[inline]
281    pub fn coordinate(self) -> u64 {
282        self.0 >> 24
283    }
284
285    /// Returns the complete packed representation.
286    #[inline]
287    pub fn raw(self) -> u64 {
288        self.0
289    }
290}
291
292#[cfg(test)]
293mod tests {
294    use super::*;
295
296    #[test]
297    fn edge_frequency_saturates_and_respects_cutoff() {
298        let mut f = EdgeFrequency::default();
299        f.add_edge(Side::Back, Base::A);
300        assert_eq!(f.edge_at(Side::Back, 1), Base::A);
301        assert_eq!(f.edge_at(Side::Back, 2), Base::E);
302        f.add_edge(Side::Back, Base::C);
303        assert_eq!(f.edge_at(Side::Back, 1), Base::N);
304        for _ in 0..20 {
305            f.add_edge(Side::Back, Base::A);
306        }
307        assert_eq!(f.frequency(Side::Back, Base::A.bits() as u32), 15);
308    }
309
310    #[test]
311    fn color_coordinate_packing_matches_limits() {
312        let c = ColorCoordinate::discovered(7, 42);
313        assert_eq!(c.as_u40(), 7 | (42 << 8));
314        let u = UnitigColor::new(123, c);
315        assert_eq!(u.offset(), 123);
316        assert_eq!(u.coordinate(), c.as_u40());
317    }
318}
cuttlefish_rs/state.rs

cuttlefish_rs/
state.rs
