MeiliSearch/milli/src/search/new/small_bitmap.rs

#[derive(Clone)]
pub enum SmallBitmap {
    Tiny(u64),
    Small(Box<[u64]>),
}
impl SmallBitmap {
    pub fn new(universe_length: u16) -> Self {
        if universe_length <= 64 {
            Self::Tiny(0)
        } else {
            Self::Small(vec![0; 1 + universe_length as usize / 64].into_boxed_slice())
        }
    }
    pub fn from_iter(xs: impl Iterator<Item = u16>, universe_length: u16) -> Self {
        let mut s = Self::new(universe_length);
        for x in xs {
            s.insert(x);
        }
        s
    }
    pub fn is_empty(&self) -> bool {
        match self {
            SmallBitmap::Tiny(set) => *set == 0,
            SmallBitmap::Small(sets) => {
                for set in sets.iter() {
                    if *set != 0 {
                        return false;
                    }
                }
                true
            }
        }
    }
    pub fn clear(&mut self) {
        match self {
            SmallBitmap::Tiny(set) => *set = 0,
            SmallBitmap::Small(sets) => {
                for set in sets.iter_mut() {
                    *set = 0;
                }
            }
        }
    }
    pub fn contains(&self, mut x: u16) -> bool {
        let set = match self {
            SmallBitmap::Tiny(set) => *set,
            SmallBitmap::Small(set) => {
                let idx = x / 64;
                x %= 64;
                set[idx as usize]
            }
        };
        set & 0b1 << x != 0
    }
    pub fn insert(&mut self, mut x: u16) {
        let set = match self {
            SmallBitmap::Tiny(set) => set,
            SmallBitmap::Small(set) => {
                let idx = x / 64;
                x %= 64;
                &mut set[idx as usize]
            }
        };
        *set |= 0b1 << x;
    }
    pub fn remove(&mut self, mut x: u16) {
        let set = match self {
            SmallBitmap::Tiny(set) => set,
            SmallBitmap::Small(set) => {
                let idx = x / 64;
                x %= 64;
                &mut set[idx as usize]
            }
        };
        *set &= !(0b1 << x);
    }

    pub fn intersection(&mut self, other: &SmallBitmap) {
        self.apply_op(other, |a, b| *a &= b);
    }
    pub fn union(&mut self, other: &SmallBitmap) {
        self.apply_op(other, |a, b| *a |= b);
    }
    pub fn subtract(&mut self, other: &SmallBitmap) {
        self.apply_op(other, |a, b| *a &= !b);
    }

    pub fn apply_op(&mut self, other: &SmallBitmap, op: impl Fn(&mut u64, u64)) {
        match (self, other) {
            (SmallBitmap::Tiny(a), SmallBitmap::Tiny(b)) => op(a, *b),
            (SmallBitmap::Small(a), SmallBitmap::Small(b)) => {
                assert!(a.len() == b.len(),);
                for (a, b) in a.iter_mut().zip(b.iter()) {
                    op(a, *b);
                }
            }
            _ => {
                panic!();
            }
        }
    }
    pub fn all_satisfy_op(&self, other: &SmallBitmap, op: impl Fn(u64, u64) -> bool) -> bool {
        match (self, other) {
            (SmallBitmap::Tiny(a), SmallBitmap::Tiny(b)) => op(*a, *b),
            (SmallBitmap::Small(a), SmallBitmap::Small(b)) => {
                assert!(a.len() == b.len());
                for (a, b) in a.iter().zip(b.iter()) {
                    if !op(*a, *b) {
                        return false;
                    }
                }
                true
            }
            _ => {
                panic!();
            }
        }
    }
    pub fn any_satisfy_op(&self, other: &SmallBitmap, op: impl Fn(u64, u64) -> bool) -> bool {
        match (self, other) {
            (SmallBitmap::Tiny(a), SmallBitmap::Tiny(b)) => op(*a, *b),
            (SmallBitmap::Small(a), SmallBitmap::Small(b)) => {
                assert!(a.len() == b.len());
                for (a, b) in a.iter().zip(b.iter()) {
                    if op(*a, *b) {
                        return true;
                    }
                }
                false
            }
            _ => {
                panic!();
            }
        }
    }
    pub fn is_subset(&self, other: &SmallBitmap) -> bool {
        self.all_satisfy_op(other, |a, b| a & !b == 0)
    }
    pub fn intersects(&self, other: &SmallBitmap) -> bool {
        self.any_satisfy_op(other, |a, b| a & b != 0)
    }
    pub fn iter(&self) -> SmallBitmapIter<'_> {
        match self {
            SmallBitmap::Tiny(x) => SmallBitmapIter::Tiny(*x),
            SmallBitmap::Small(xs) => {
                SmallBitmapIter::Small { cur: xs[0], next: &xs[1..], base: 0 }
            }
        }
    }
}

pub enum SmallBitmapIter<'b> {
    Tiny(u64),
    Small { cur: u64, next: &'b [u64], base: u16 },
}
impl<'b> Iterator for SmallBitmapIter<'b> {
    type Item = u16;

    fn next(&mut self) -> Option<Self::Item> {
        match self {
            SmallBitmapIter::Tiny(set) => {
                if *set > 0 {
                    let idx = set.trailing_zeros() as u16;
                    *set &= *set - 1;
                    Some(idx)
                } else {
                    None
                }
            }
            SmallBitmapIter::Small { cur, next, base } => {
                if *cur > 0 {
                    let idx = cur.trailing_zeros() as u16;
                    *cur &= *cur - 1;
                    Some(idx + *base)
                } else if next.is_empty() {
                    return None;
                } else {
                    *base += 64;
                    *cur = next[0];
                    *next = &next[1..];
                    self.next()
                }
            }
        }
    }
}

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

    #[test]
    fn test_small_bitmap() {
        let mut bitmap1 = SmallBitmap::new(32);
        for x in 0..16 {
            bitmap1.insert(x * 2);
        }
        let mut bitmap2 = SmallBitmap::new(32);
        for x in 0..=10 {
            bitmap2.insert(x * 3);
        }
        bitmap1.intersection(&bitmap2);
        for v in bitmap1.iter() {
            println!("{v}");
        }
    }
}