nordisk/edda

crates/nornir-parquet-delta/src/presence.rs

use crate::Error;

/// **Which (row group × column) chunks the mirror holds** — one bit each, dense.
///
/// Dense and not EWAH on purpose. The population is `row_groups × columns` — 14 807 bits
/// (1.8 kB) for a 1 139-group, 13-column table — and it is written one bit at a time by
/// arriving chunks and read as a whole by the interval builder. A compressed bitmap would
/// be re-encoded on every arrival to save a kilobyte.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct Presence {
    row_groups: usize,
    columns: usize,
    words: Vec<u64>,
}

impl Presence {
    /// An empty presence over a `row_groups × columns` grid.
    #[must_use]
    pub fn empty(row_groups: usize, columns: usize) -> Self {
        let bits = row_groups.saturating_mul(columns);
        Self { row_groups, columns, words: vec![0; bits.div_ceil(64)] }
    }

    /// Rebuild from stored words. Rejects a word count that does not match the grid, so a
    /// truncated sidecar cannot come back as a presence that silently claims fewer chunks.
    pub fn from_words(row_groups: usize, columns: usize, words: Vec<u64>) -> Result<Self, Error> {
        let want = row_groups.saturating_mul(columns).div_ceil(64);
        if words.len() != want {
            return Err(Error::BadState {
                why: format!("bitmap is {} words, a {row_groups}×{columns} grid needs {want}", words.len()),
            });
        }
        Ok(Self { row_groups, columns, words })
    }

    /// The raw words, for the sidecar writer.
    #[must_use]
    pub fn words(&self) -> &[u64] {
        &self.words
    }

    /// Row groups in the grid.
    #[must_use]
    pub fn row_groups(&self) -> usize {
        self.row_groups
    }

    /// Leaf columns in the grid.
    #[must_use]
    pub fn columns(&self) -> usize {
        self.columns
    }

    fn bit(&self, rg: usize, col: usize) -> Option<usize> {
        (rg < self.row_groups && col < self.columns).then(|| rg * self.columns + col)
    }

    /// Is this chunk held? `false` for any index outside the grid — an out-of-range chunk
    /// is one the mirror certainly does not hold.
    #[must_use]
    pub fn holds(&self, rg: usize, col: usize) -> bool {
        match self.bit(rg, col) {
            Some(b) => self.words[b / 64] >> (b % 64) & 1 == 1,
            None => false,
        }
    }

    /// Record a chunk as held. Returns `false` (and changes nothing) for an index outside
    /// the grid.
    pub fn set(&mut self, rg: usize, col: usize) -> bool {
        match self.bit(rg, col) {
            Some(b) => {
                self.words[b / 64] |= 1 << (b % 64);
                true
            }
            None => false,
        }
    }

    /// How many chunks are held.
    #[must_use]
    pub fn count(&self) -> u64 {
        self.words.iter().map(|w| u64::from(w.count_ones())).sum()
    }

    /// How many chunks the grid has.
    #[must_use]
    pub fn total(&self) -> u64 {
        (self.row_groups as u64).saturating_mul(self.columns as u64)
    }

    /// Is every chunk held? The "the mirror is now the file" answer.
    #[must_use]
    pub fn complete(&self) -> bool {
        self.count() == self.total()
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn a_set_bit_is_the_only_bit_that_reads_back() {
        let mut p = Presence::empty(3, 5);
        assert_eq!(p.count(), 0);
        assert!(p.set(1, 3));
        assert!(p.holds(1, 3));
        for rg in 0..3 {
            for c in 0..5 {
                assert_eq!(p.holds(rg, c), (rg, c) == (1, 3), "({rg},{c})");
            }
        }
        assert_eq!(p.count(), 1);
        assert_eq!(p.total(), 15);
        assert!(!p.complete());
    }

    /// RED FIRST: with `bit()` unguarded, `set(3, 0)` on a 3-row grid writes bit 15 —
    /// which is row group 3's column 0 in a grid that has no row group 3, and on a
    /// tighter grid it aliases onto a REAL chunk and claims a hole is held.
    #[test]
    fn an_index_outside_the_grid_is_never_held_and_never_settable() {
        let mut p = Presence::empty(3, 5);
        assert!(!p.set(3, 0), "row group 3 does not exist in a 3-group grid");
        assert!(!p.set(0, 5), "column 5 does not exist in a 5-column grid");
        assert_eq!(p.count(), 0, "an out-of-range set must not mark anything");
        assert!(!p.holds(3, 0));
        assert!(!p.holds(0, 5));
    }

    #[test]
    fn a_bitmap_of_the_wrong_width_is_rejected_rather_than_reinterpreted() {
        let p = Presence::empty(20, 20); // 400 bits = 7 words
        assert_eq!(p.words().len(), 7);
        assert!(Presence::from_words(20, 20, p.words().to_vec()).is_ok());
        assert!(Presence::from_words(20, 20, vec![0; 6]).is_err(), "too few words");
        assert!(Presence::from_words(20, 20, vec![0; 8]).is_err(), "too many words");
        // The word count does NOT pin the grid: 20x21 = 420 bits is also 7 words. That
        // is why `Mirror::open` checks the grid against the FOOTER and does not trust
        // this length alone — see `store::Mirror::open`.
        assert!(Presence::from_words(20, 21, p.words().to_vec()).is_ok());
    }

    #[test]
    fn complete_is_every_chunk_and_not_merely_many() {
        let mut p = Presence::empty(2, 2);
        for rg in 0..2 {
            for c in 0..2 {
                assert!(!p.complete());
                p.set(rg, c);
            }
        }
        assert!(p.complete());
    }
}