cranpose-core 0.1.164

use std::fmt;

pub type SnapshotId = usize;

const BITS_PER_SET: usize = 64;
const SNAPSHOT_ID_SIZE: usize = 64;

#[derive(Clone, PartialEq, Eq)]
pub struct SnapshotIdSet {
    upper_set: u64,
    lower_set: u64,
    lower_bound: SnapshotId,
    below_bound: Option<Box<[SnapshotId]>>,
}

impl SnapshotIdSet {
    /// Empty snapshot ID set.
    pub const EMPTY: SnapshotIdSet = SnapshotIdSet {
        upper_set: 0,
        lower_set: 0,
        lower_bound: 0,
        below_bound: None,
    };

    /// Create a new empty snapshot ID set.
    pub fn new() -> Self {
        Self::EMPTY
    }

    /// Check if a snapshot ID is in the set.
    pub fn get(&self, id: SnapshotId) -> bool {
        let offset = id.wrapping_sub(self.lower_bound);

        if offset < BITS_PER_SET {
            let mask = 1u64 << offset;
            (self.lower_set & mask) != 0
        } else if offset < BITS_PER_SET * 2 {
            let mask = 1u64 << (offset - BITS_PER_SET);
            (self.upper_set & mask) != 0
        } else if id > self.lower_bound {
            false
        } else {
            self.below_bound
                .as_ref()
                .is_some_and(|arr| arr.binary_search(&id).is_ok())
        }
    }

    /// Add a snapshot ID to the set (returns a new set if modified).
    pub fn set(&self, id: SnapshotId) -> Self {
        if id < self.lower_bound {
            if let Some(ref arr) = self.below_bound {
                match arr.binary_search(&id) {
                    Ok(_) => {
                        return self.clone();
                    }
                    Err(insert_pos) => {
                        let mut new_arr = Vec::with_capacity(arr.len() + 1);
                        new_arr.extend_from_slice(&arr[..insert_pos]);
                        new_arr.push(id);
                        new_arr.extend_from_slice(&arr[insert_pos..]);
                        return Self {
                            upper_set: self.upper_set,
                            lower_set: self.lower_set,
                            lower_bound: self.lower_bound,
                            below_bound: Some(new_arr.into_boxed_slice()),
                        };
                    }
                }
            } else {
                return Self {
                    upper_set: self.upper_set,
                    lower_set: self.lower_set,
                    lower_bound: self.lower_bound,
                    below_bound: Some(vec![id].into_boxed_slice()),
                };
            }
        }

        let offset = id - self.lower_bound;

        if offset < BITS_PER_SET {
            let mask = 1u64 << offset;
            if (self.lower_set & mask) == 0 {
                return Self {
                    upper_set: self.upper_set,
                    lower_set: self.lower_set | mask,
                    lower_bound: self.lower_bound,
                    below_bound: self.below_bound.clone(),
                };
            }
        } else if offset < BITS_PER_SET * 2 {
            let mask = 1u64 << (offset - BITS_PER_SET);
            if (self.upper_set & mask) == 0 {
                return Self {
                    upper_set: self.upper_set | mask,
                    lower_set: self.lower_set,
                    lower_bound: self.lower_bound,
                    below_bound: self.below_bound.clone(),
                };
            }
        } else if offset >= BITS_PER_SET * 2 && !self.get(id) {
            return self.shift_and_set(id);
        }

        self.clone()
    }

    /// Remove a snapshot ID from the set (returns a new set if modified).
    pub fn clear(&self, id: SnapshotId) -> Self {
        let offset = id.wrapping_sub(self.lower_bound);

        if offset < BITS_PER_SET {
            let mask = 1u64 << offset;
            if (self.lower_set & mask) != 0 {
                return Self {
                    upper_set: self.upper_set,
                    lower_set: self.lower_set & !mask,
                    lower_bound: self.lower_bound,
                    below_bound: self.below_bound.clone(),
                };
            }
        } else if offset < BITS_PER_SET * 2 {
            let mask = 1u64 << (offset - BITS_PER_SET);
            if (self.upper_set & mask) != 0 {
                return Self {
                    upper_set: self.upper_set & !mask,
                    lower_set: self.lower_set,
                    lower_bound: self.lower_bound,
                    below_bound: self.below_bound.clone(),
                };
            }
        } else if id < self.lower_bound
            && let Some(ref arr) = self.below_bound
            && let Ok(pos) = arr.binary_search(&id)
        {
            let mut new_arr = Vec::with_capacity(arr.len() - 1);
            new_arr.extend_from_slice(&arr[..pos]);
            new_arr.extend_from_slice(&arr[pos + 1..]);
            return Self {
                upper_set: self.upper_set,
                lower_set: self.lower_set,
                lower_bound: self.lower_bound,
                below_bound: if new_arr.is_empty() {
                    None
                } else {
                    Some(new_arr.into_boxed_slice())
                },
            };
        }

        self.clone()
    }

    /// Remove all IDs in `other` from this set (a & ~b).
    pub fn and_not(&self, other: &Self) -> Self {
        if other.is_empty() {
            return self.clone();
        }
        if self.is_empty() {
            return Self::EMPTY;
        }

        if self.lower_bound == other.lower_bound && self.below_bound_equals(&other.below_bound) {
            return Self {
                upper_set: self.upper_set & !other.upper_set,
                lower_set: self.lower_set & !other.lower_set,
                lower_bound: self.lower_bound,
                below_bound: self.below_bound.clone(),
            };
        }

        let mut result = self.clone();
        for id in other.iter() {
            result = result.clear(id);
        }
        result
    }

    /// Union this set with another (a | b).
    pub fn or(&self, other: &Self) -> Self {
        if other.is_empty() {
            return self.clone();
        }
        if self.is_empty() {
            return other.clone();
        }

        if self.lower_bound == other.lower_bound && self.below_bound_equals(&other.below_bound) {
            return Self {
                upper_set: self.upper_set | other.upper_set,
                lower_set: self.lower_set | other.lower_set,
                lower_bound: self.lower_bound,
                below_bound: self.below_bound.clone(),
            };
        }

        let mut result = self.clone();
        for id in other.iter() {
            result = result.set(id);
        }
        result
    }

    /// Find the lowest snapshot ID in the set that is <= upper.
    pub fn lowest(&self, upper: SnapshotId) -> SnapshotId {
        if let Some(ref arr) = self.below_bound
            && let Some(&lowest) = arr.first()
            && lowest <= upper
        {
            return lowest;
        }

        if self.lower_set != 0 {
            let lowest_in_lower = self.lower_bound + self.lower_set.trailing_zeros() as usize;
            if lowest_in_lower <= upper {
                return lowest_in_lower;
            }
        }

        if self.upper_set != 0 {
            let lowest_in_upper =
                self.lower_bound + BITS_PER_SET + self.upper_set.trailing_zeros() as usize;
            if lowest_in_upper <= upper {
                return lowest_in_upper;
            }
        }

        upper
    }

    /// Check if the set is empty.
    pub fn is_empty(&self) -> bool {
        self.lower_set == 0 && self.upper_set == 0 && self.below_bound.is_none()
    }

    /// Iterate over all snapshot IDs in the set.
    pub fn iter(&self) -> SnapshotIdSetIter<'_> {
        SnapshotIdSetIter::new(self)
    }

    /// Convert to a Vec of snapshot IDs (for testing/debugging).
    pub fn to_list(&self) -> Vec<SnapshotId> {
        self.iter().collect()
    }

    /// Add a contiguous range of IDs [from, until) to the set.
    /// Mirrors AndroidX SnapshotIdSet.addRange semantics used by Snapshot.kt.
    pub fn add_range(&self, from: SnapshotId, until: SnapshotId) -> Self {
        if from >= until {
            return self.clone();
        }
        let mut result = self.clone();
        let mut id = from;
        while id < until {
            result = result.set(id);
            id += 1;
        }
        result
    }

    fn below_bound_equals(&self, other: &Option<Box<[SnapshotId]>>) -> bool {
        match (&self.below_bound, other) {
            (None, None) => true,
            (Some(a), Some(b)) => a == b,
            _ => false,
        }
    }

    fn shift_and_set(&self, id: SnapshotId) -> Self {
        let target_lower_bound = (id / SNAPSHOT_ID_SIZE) * SNAPSHOT_ID_SIZE;

        let mut new_upper_set = self.upper_set;
        let mut new_lower_set = self.lower_set;
        let mut new_lower_bound = self.lower_bound;
        let mut new_below_bound: Vec<SnapshotId> = if let Some(ref arr) = self.below_bound {
            arr.to_vec()
        } else {
            Vec::new()
        };

        while new_lower_bound < target_lower_bound {
            if new_lower_set != 0 {
                for bit_offset in 0..BITS_PER_SET {
                    if (new_lower_set & (1u64 << bit_offset)) != 0 {
                        let id_to_add = new_lower_bound + bit_offset;
                        match new_below_bound.binary_search(&id_to_add) {
                            Ok(_) => {}
                            Err(pos) => new_below_bound.insert(pos, id_to_add),
                        }
                    }
                }
            }

            if new_upper_set == 0 {
                new_lower_bound = target_lower_bound;
                new_lower_set = 0;
                break;
            }

            new_lower_set = new_upper_set;
            new_upper_set = 0;
            new_lower_bound += BITS_PER_SET;
        }

        let result = Self {
            upper_set: new_upper_set,
            lower_set: new_lower_set,
            lower_bound: new_lower_bound,
            below_bound: if new_below_bound.is_empty() {
                None
            } else {
                Some(new_below_bound.into_boxed_slice())
            },
        };

        result.set(id)
    }
}

impl Default for SnapshotIdSet {
    fn default() -> Self {
        Self::EMPTY
    }
}

impl fmt::Debug for SnapshotIdSet {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "SnapshotIdSet{{")?;
        let ids: Vec<_> = self.iter().collect();
        for (i, id) in ids.iter().enumerate() {
            if i > 0 {
                write!(f, ", ")?;
            }
            write!(f, "{id}")?;
        }
        write!(f, "}}")
    }
}

/// Iterator over snapshot IDs in a set.
pub struct SnapshotIdSetIter<'a> {
    set: &'a SnapshotIdSet,
    below_index: usize,
    lower_set: u64,
    upper_set: u64,
    current_offset: usize,
}

impl<'a> SnapshotIdSetIter<'a> {
    fn new(set: &'a SnapshotIdSet) -> Self {
        Self {
            set,
            below_index: 0,
            lower_set: set.lower_set,
            upper_set: set.upper_set,
            current_offset: 0,
        }
    }
}

impl Iterator for SnapshotIdSetIter<'_> {
    type Item = SnapshotId;

    fn next(&mut self) -> Option<Self::Item> {
        if let Some(ref arr) = self.set.below_bound
            && self.below_index < arr.len()
        {
            let id = arr[self.below_index];
            self.below_index += 1;
            return Some(id);
        }

        while self.current_offset < BITS_PER_SET {
            if (self.lower_set & (1u64 << self.current_offset)) != 0 {
                let id = self.set.lower_bound + self.current_offset;
                self.current_offset += 1;
                return Some(id);
            }
            self.current_offset += 1;
        }

        while self.current_offset < BITS_PER_SET * 2 {
            let bit_offset = self.current_offset - BITS_PER_SET;
            if (self.upper_set & (1u64 << bit_offset)) != 0 {
                let id = self.set.lower_bound + self.current_offset;
                self.current_offset += 1;
                return Some(id);
            }
            self.current_offset += 1;
        }

        None
    }
}

#[cfg(test)]
#[path = "tests/snapshot_id_set_tests.rs"]
mod tests;