ef_rs 0.2.0

use bitvec::prelude::*;
use sucds::bit_vectors::{BitVector, DArray, Select};

/// A compressed representation of a sorted sequence using a two-level encoding scheme.
/// This structure efficiently stores monotonically increasing integers using a combination
/// of high and low bits, providing fast access while maintaining good compression ratios.
///
/// # Examples
///
/// Basic usage:
/// ```
/// use ef_rs::elias_fano::EliasFano;
///
/// // Create from a sorted sequence
/// let values = vec![2, 5, 5, 9];
/// let seq = EliasFano::from(&values);
///
/// // Access elements
/// assert_eq!(seq.get(0), Some(2));
/// assert_eq!(seq.get(1), Some(5));
/// assert_eq!(seq.get(2), Some(5));
/// assert_eq!(seq.get(3), Some(9));
/// assert_eq!(seq.get(4), None);  // Out of bounds
///
/// // Get sequence properties
/// assert_eq!(seq.size(), 4);
/// assert_eq!(seq.universe(), 10);  // max value + 1
/// assert!(!seq.is_empty());
///
/// // Iterate over values
/// let collected: Vec<_> = seq.iter().collect();
/// assert_eq!(collected, values);
/// ```
///
/// Using the builder pattern:
/// ```
/// use ef_rs::elias_fano::EliasFanoBuilder;
///
/// let mut builder = EliasFanoBuilder::new(10, 4);  // universe=10, count=4
/// builder.add(2);
/// builder.add(5);
/// builder.add(5);
/// builder.add(9);
/// let seq = builder.finalize();
///
/// assert_eq!(seq.get(0), Some(2));
/// assert_eq!(seq.get(1), Some(5));
/// assert_eq!(seq.get(2), Some(5));
/// assert_eq!(seq.get(3), Some(9));
/// ```
///
/// Handling empty sequences:
/// ```
/// use ef_rs::elias_fano::EliasFano;
///
/// let seq = EliasFano::from(&[] as &[usize]);
/// assert!(seq.is_empty());
/// assert_eq!(seq.size(), 0);
/// assert_eq!(seq.get(0), None);
/// ```
///
/// # Panics
///
/// The following operations will panic:
///
/// ```should_panic
/// use ef_rs::elias_fano::EliasFano;
///
/// // Input must be sorted
/// let _ = EliasFano::from(&[5, 2, 7]);
/// ```
///
/// ```should_panic
/// use ef_rs::elias_fano::EliasFanoBuilder;
///
/// // Values must be within universe
/// let mut builder = EliasFanoBuilder::new(5, 2);
/// builder.add(2);
/// builder.add(6);  // 6 >= universe (5)
/// ```
///
/// # Performance
///
/// The structure provides O(1) access to elements while using approximately
/// 2n + log(u/n) bits of space, where n is the number of elements and u is
/// the universe size (maximum value + 1).
#[derive(Default, Debug, Clone, PartialEq)]
pub struct EliasFano {
    universe: usize,
    element_count: usize,
    upper_bits: DArray,
    lower_bits: BitVector,
    lower_bit_count: usize,
}

impl EliasFano {
    /// Returns a value that can be used for indexing
    pub fn index(&self, idx: usize) -> usize {
        self.get(idx).expect("index out of bounds")
    }

    /// Retrieves the value at the specified index.
    ///
    /// Returns `None` if the index is out of bounds.
    pub fn get(&self, index: usize) -> Option<usize> {
        if index >= self.element_count {
            None
        } else {
            Some(
                ((self.upper_bits.select1(index).unwrap() - index) << self.lower_bit_count)
                    | self
                        .lower_bits
                        .get_bits(index * self.lower_bit_count, self.lower_bit_count)
                        .unwrap(),
            )
        }
    }

    /// Returns the number of elements in the sequence.
    pub fn size(&self) -> usize {
        self.element_count
    }

    /// Returns the maximum possible value (exclusive) in the sequence.
    pub fn universe(&self) -> usize {
        self.universe
    }

    /// Checks if the sequence is empty.
    pub fn is_empty(&self) -> bool {
        self.element_count == 0
    }

    /// Returns an iterator over the values in the sequence.
    pub fn iter(&self) -> EliasFanoIter<'_> {
        EliasFanoIter { ef: self, index: 0 }
    }
}

impl<T> From<&[T]> for EliasFano
where
    T: TryInto<usize> + Copy + PartialOrd,
    <T as TryInto<usize>>::Error: std::fmt::Debug,
{
    fn from(values: &[T]) -> Self {
        if values.is_empty() {
            return Self::default();
        }

        assert!(
            values.windows(2).all(|w| w[0] <= w[1]),
            "Input sequence must be sorted in ascending order"
        );

        let last_value: usize = (*values.last().unwrap()).try_into().unwrap();
        let mut builder = EliasFanoBuilder::new(last_value + 1, values.len());

        // Convert each value to usize before adding
        for &value in values {
            let value: usize = value.try_into().unwrap();
            builder.add(value);
        }

        builder.finalize()
    }
}

impl<T> From<Vec<T>> for EliasFano
where
    T: TryInto<usize> + Copy + PartialOrd,
    <T as TryInto<usize>>::Error: std::fmt::Debug,
{
    fn from(values: Vec<T>) -> Self {
        Self::from(values.as_slice())
    }
}

impl<T> From<&Vec<T>> for EliasFano
where
    T: TryInto<usize> + Copy + PartialOrd,
    <T as TryInto<usize>>::Error: std::fmt::Debug,
{
    fn from(values: &Vec<T>) -> Self {
        Self::from(values.as_slice())
    }
}

impl<T, const N: usize> From<&[T; N]> for EliasFano
where
    T: TryInto<usize> + Copy + PartialOrd,
    <T as TryInto<usize>>::Error: std::fmt::Debug,
{
    fn from(values: &[T; N]) -> Self {
        Self::from(values.as_slice())
    }
}

impl<T, const N: usize> From<[T; N]> for EliasFano
where
    T: TryInto<usize> + Copy + PartialOrd,
    <T as TryInto<usize>>::Error: std::fmt::Debug,
{
    fn from(values: [T; N]) -> Self {
        Self::from(values.as_slice())
    }
}

/// A builder for creating compressed sequences using the Elias-Fano encoding.
///
/// This builder allows for incremental construction of an Elias-Fano sequence,
/// which is useful when the values are not known in advance or when building
/// the sequence from a stream of values.
///
/// # Examples
///
/// Basic usage:
/// ```
/// use ef_rs::elias_fano::EliasFanoBuilder;
///
/// let mut builder = EliasFanoBuilder::new(10, 4);  // universe=10, count=4
/// builder.add(2);
/// builder.add(5);
/// builder.add(5);
/// builder.add(9);
/// let seq = builder.finalize();
///
/// assert_eq!(seq.get(0), Some(2));
/// assert_eq!(seq.get(1), Some(5));
/// assert_eq!(seq.get(2), Some(5));
/// assert_eq!(seq.get(3), Some(9));
/// ```
///
/// Adding multiple values at once:
/// ```
/// use ef_rs::elias_fano::EliasFanoBuilder;
///
/// let mut builder = EliasFanoBuilder::new(10, 4);
/// builder.add_all(&[2, 5, 5, 9]);
/// let seq = builder.finalize();
///
/// assert_eq!(seq.get(0), Some(2));
/// assert_eq!(seq.get(1), Some(5));
/// assert_eq!(seq.get(2), Some(5));
/// assert_eq!(seq.get(3), Some(9));
/// ```
///
/// # Panics
///
/// The following operations will panic:
///
/// ```should_panic
/// use ef_rs::elias_fano::EliasFanoBuilder;
///
/// // Element count must be greater than zero
/// let _ = EliasFanoBuilder::new(10, 0);
/// ```
///
/// ```should_panic
/// use ef_rs::elias_fano::EliasFanoBuilder;
///
/// // Values must be in ascending order
/// let mut builder = EliasFanoBuilder::new(10, 2);
/// builder.add(5);
/// builder.add(2);  // 2 < 5
/// ```
///
/// ```should_panic
/// use ef_rs::elias_fano::EliasFanoBuilder;
///
/// // Values must be within universe
/// let mut builder = EliasFanoBuilder::new(5, 2);
/// builder.add(2);
/// builder.add(6);  // 6 >= universe (5)
/// ```
///
/// ```should_panic
/// use ef_rs::elias_fano::EliasFanoBuilder;
///
/// // Cannot add more values than specified count
/// let mut builder = EliasFanoBuilder::new(10, 1);
/// builder.add(2);
/// builder.add(5);  // Already added 1 value
/// ```
///
/// # Performance
///
/// The builder pre-allocates space for the sequence based on the universe size
/// and element count, providing O(1) amortized time for adding values. The
/// finalize operation is O(1) as it just returns the constructed sequence.
pub struct EliasFanoBuilder {
    upper_bits: BitVector,
    lower_bits: BitVector,
    universe: usize,
    element_count: usize,
    current_index: usize,
    previous_value: usize,
    lower_bit_count: usize,
}

impl EliasFanoBuilder {
    /// Creates a new builder for a sequence that will store `count` values
    /// with a maximum value of `universe` (exclusive).
    ///
    /// # Arguments
    ///
    /// * `universe` - The maximum possible value (exclusive) in the sequence.
    /// * `count` - The number of values that will be stored in the sequence.
    ///
    /// # Panics
    ///
    /// Panics if `count` is zero.
    ///
    /// # Examples
    ///
    /// ```
    /// use ef_rs::elias_fano::EliasFanoBuilder;
    ///
    /// let builder = EliasFanoBuilder::new(10, 4);
    /// ```
    pub fn new(universe: usize, count: usize) -> Self {
        assert!(count > 0, "Element count must be greater than zero");

        let lower_bit_count = if count > 0 {
            ((universe as f64) / (count as f64)).log2().floor() as usize
        } else {
            0
        };

        let upper_size = count + (universe >> lower_bit_count) + 1;

        Self {
            upper_bits: BitVector::from_bits(bitvec![u64, Lsb0; 0; upper_size]),
            lower_bits: BitVector::new(),
            universe,
            element_count: count,
            current_index: 0,
            previous_value: 0,
            lower_bit_count,
        }
    }

    /// Adds a single value to the sequence.
    ///
    /// # Arguments
    ///
    /// * `value` - The value to add to the sequence.
    ///
    /// # Panics
    ///
    /// Panics if:
    /// - The value is less than the previous value (not in ascending order)
    /// - The value is greater than or equal to the maximum value (universe)
    /// - The sequence is already full (current_index >= element_count)
    ///
    /// # Examples
    ///
    /// ```
    /// use ef_rs::elias_fano::EliasFanoBuilder;
    ///
    /// let mut builder = EliasFanoBuilder::new(10, 2);
    /// builder.add(2);
    /// builder.add(5);
    /// ```
    pub fn add(&mut self, value: usize) {
        assert!(
            self.previous_value <= value,
            "Values must be in ascending order"
        );
        assert!(
            value < self.universe,
            "Value {} exceeds maximum allowed value {}",
            value,
            self.universe
        );
        assert!(
            self.current_index < self.element_count,
            "Sequence is already full"
        );

        self.previous_value = value;

        if self.lower_bit_count > 0 {
            let lower_mask = (1 << self.lower_bit_count) - 1;
            let lower_value = value & lower_mask;
            for i in 0..self.lower_bit_count {
                self.lower_bits.push_bit((lower_value & (1 << i)) != 0);
            }
        }

        let upper_value = value >> self.lower_bit_count;
        _ = self
            .upper_bits
            .set_bit(upper_value + self.current_index, true);
        self.current_index += 1;
    }

    /// Adds multiple values to the sequence.
    ///
    /// # Arguments
    ///
    /// * `values` - An iterator over values to add to the sequence.
    ///
    /// # Panics
    ///
    /// Panics if any of the values violate the ordering or range constraints
    /// (see `add` method for details).
    ///
    /// # Examples
    ///
    /// ```
    /// use ef_rs::elias_fano::EliasFanoBuilder;
    ///
    /// let mut builder = EliasFanoBuilder::new(10, 4);
    /// builder.add_all(&[2, 5, 5, 9]);
    /// ```
    pub fn add_all<'a, I>(&mut self, values: I)
    where
        I: IntoIterator<Item = &'a usize>,
    {
        for &value in values {
            self.add(value);
        }
    }

    /// Finalizes the sequence and returns the compressed representation.
    ///
    /// This method consumes the builder and returns the constructed Elias-Fano
    /// sequence. The sequence is ready for use immediately after finalization.
    ///
    /// # Examples
    ///
    /// ```
    /// use ef_rs::elias_fano::EliasFanoBuilder;
    ///
    /// let mut builder = EliasFanoBuilder::new(10, 2);
    /// builder.add(2);
    /// builder.add(5);
    /// let seq = builder.finalize();
    ///
    /// assert_eq!(seq.get(0), Some(2));
    /// assert_eq!(seq.get(1), Some(5));
    /// ```
    pub fn finalize(self) -> EliasFano {
        EliasFano {
            upper_bits: DArray::from_bits(self.upper_bits.iter()).enable_select0(),
            lower_bits: self.lower_bits,
            lower_bit_count: self.lower_bit_count,
            universe: self.universe,
            element_count: self.element_count,
        }
    }
}

/// An iterator over the values in an Elias-Fano sequence.
pub struct EliasFanoIter<'a> {
    ef: &'a EliasFano,
    index: usize,
}

impl<'a> EliasFanoIter<'a> {
    /// Looks up a value in the sequence starting from the current iterator position.
    /// Returns the index of the value if found, or None if not found.
    /// The iterator position is updated to where the search ended.
    ///
    /// # Examples
    ///
    /// ```
    /// use ef_rs::elias_fano::EliasFano;
    ///
    /// let values = vec![2, 5, 5, 9];
    /// let seq = EliasFano::from(&values);
    /// let mut iter = seq.iter();
    ///
    /// assert_eq!(iter.lookup(5), Some(1));  // Found at index 1
    /// iter.next();  // Move to index 1
    /// assert_eq!(iter.lookup(5), Some(2));  // Found at index 2
    /// assert_eq!(iter.lookup(7), None);     // Not found (index should be 3)
    /// ```
    pub fn lookup(&mut self, value: usize) -> Option<usize> {
        if value >= self.ef.universe {
            return None;
        }

        // Use lookup_next to find the position
        let pos = self.lookup_next(value);
        // Update the iterator position to where the search ended
        self.index = pos;
        if pos < self.ef.element_count {
            // Only return the position if the value matches exactly
            if let Some(current) = self.ef.get(pos) {
                if current == value {
                    return Some(pos);
                }
            }
        }
        None
    }

    /// Looks up a value in the sequence starting from the current iterator position.
    /// If the value is not found, returns the index of the next greater value.
    /// If the value is larger than the universe, returns the last index.
    ///
    /// # Examples
    ///
    /// ```
    /// use ef_rs::elias_fano::EliasFano;
    ///
    /// let values = vec![2, 5, 5, 9];
    /// let seq = EliasFano::from(&values);
    /// let mut iter = seq.iter();
    ///
    /// assert_eq!(iter.lookup_next(4), 1);  // Next value is 5 at index 1
    /// iter.next();  // Move to index 1
    /// assert_eq!(iter.lookup_next(6), 3);  // Next value is 9 at index 3
    /// assert_eq!(iter.lookup_next(10), 3); // should return the last index
    /// ```
    pub fn lookup_next(&self, value: usize) -> usize {
        if value >= self.ef.universe {
            return self.ef.element_count - 1;
        }

        // Get current value at iterator position
        if let Some(current) = self.ef.get(self.index) {
            // If searching for current value, return current position
            if value == current {
                return self.index;
            }

            // For small skips, do linear scan
            if value > current {
                const LINEAR_SCAN_THRESHOLD: usize = 8;
                let high_value = value >> self.ef.lower_bit_count;
                let high_current = current >> self.ef.lower_bit_count;
                let high_diff = high_value.saturating_sub(high_current);

                if high_diff <= LINEAR_SCAN_THRESHOLD {
                    // Linear scan for small skips
                    let mut pos = self.index + 1;
                    while pos < self.ef.element_count {
                        if let Some(val) = self.ef.get(pos) {
                            if val >= value {
                                return pos;
                            }
                            pos += 1;
                        } else {
                            break;
                        }
                    }
                    return self.ef.element_count - 1;
                }
            }
        }

        // Fall back to binary search
        let mut left = self.index;
        let mut right = self.ef.element_count;

        while left < right {
            let mid = left + (right - left) / 2;
            if let Some(val) = self.ef.get(mid) {
                if val < value {
                    left = mid + 1;
                } else {
                    right = mid;
                }
            } else {
                break;
            }
        }

        if left < self.ef.element_count {
            if let Some(val) = self.ef.get(left) {
                if val >= value {
                    return left;
                }
            }
        }
        self.ef.element_count - 1
    }
}

impl<'a> Iterator for EliasFanoIter<'a> {
    type Item = usize;

    fn next(&mut self) -> Option<Self::Item> {
        let value = self.ef.get(self.index)?;
        self.index += 1;
        Some(value)
    }

    fn size_hint(&self) -> (usize, Option<usize>) {
        let remaining = self.ef.size().saturating_sub(self.index);
        (remaining, Some(remaining))
    }
}

impl<'a> ExactSizeIterator for EliasFanoIter<'a> {}

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

    #[test]
    fn test_basic_operations() {
        let values = vec![2, 5, 5, 9];
        let seq = EliasFano::from(&values);

        assert_eq!(seq.size(), 4);
        assert_eq!(seq.universe(), 10);
        assert!(!seq.is_empty());

        assert_eq!(seq.get(0), Some(2));
        assert_eq!(seq.get(1), Some(5));
        assert_eq!(seq.get(2), Some(5));
        assert_eq!(seq.get(3), Some(9));
        assert_eq!(seq.get(4), None);
    }

    #[test]
    fn test_indexing() {
        let values = vec![2, 5, 5, 9];
        let seq = EliasFano::from(&values);

        // Test indexing
        assert_eq!(seq.index(0), 2);
        assert_eq!(seq.index(1), 5);
        assert_eq!(seq.index(2), 5);
        assert_eq!(seq.index(3), 9);
    }

    #[test]
    #[should_panic(expected = "index out of bounds")]
    fn test_index_out_of_bounds() {
        let values = vec![2, 5, 5, 9];
        let seq = EliasFano::from(&values);
        let _ = seq.index(4);
    }

    #[test]
    fn test_single_value() {
        let values = vec![2];
        let seq = EliasFano::from(&values);
        assert_eq!(seq.get(0), Some(2));
    }

    #[test]
    fn test_consecutive_values() {
        let values = vec![1, 2, 3, 4];
        let seq = EliasFano::from(&values);
        assert_eq!(seq.get(0), Some(1));
        assert_eq!(seq.get(1), Some(2));
        assert_eq!(seq.get(2), Some(3));
        assert_eq!(seq.get(3), Some(4));
    }

    #[test]
    fn test_empty_sequence() {
        let seq = EliasFano::from(&[] as &[usize]);
        assert!(seq.is_empty());
        assert_eq!(seq.size(), 0);
    }

    #[test]
    #[should_panic(expected = "Input sequence must be sorted")]
    fn test_unsorted_input() {
        let values: &[usize] = &[5, 2, 7];
        let _ = EliasFano::from(values);
    }

    #[test]
    fn test_iter() {
        let values = vec![2, 5, 5, 9];
        let seq = EliasFano::from(&values);

        // Test iterator
        let collected: Vec<_> = seq.iter().collect();
        assert_eq!(collected, values);

        // Test size_hint
        let mut iter = seq.iter();
        assert_eq!(iter.size_hint(), (4, Some(4)));
        iter.next();
        assert_eq!(iter.size_hint(), (3, Some(3)));
    }

    #[test]
    fn test_array_implementation() {
        let values = [2usize, 5, 5, 9];
        let seq = EliasFano::from(&values);
        assert_eq!(seq.size(), 4);
        assert_eq!(seq.get(0), Some(2));
        assert_eq!(seq.get(1), Some(5));
        assert_eq!(seq.get(2), Some(5));
        assert_eq!(seq.get(3), Some(9));
    }

    #[test]
    fn test_iter_lookup() {
        let values = vec![2, 5, 5, 9];
        let seq = EliasFano::from(&values);
        let mut iter = seq.iter();

        // Test lookup from start
        assert_eq!(iter.lookup(5), Some(1));
        assert_eq!(iter.lookup(7), None);

        // we should be at index 3 now
        assert_eq!(iter.index, 3);

        // Test lookup after moving iterator
        // move back to index 0
        iter.index = 0;

        iter.next(); // Move to index 1
        assert_eq!(iter.lookup(5), Some(1));
        iter.next(); // Move to index 2
        assert_eq!(iter.lookup(5), Some(2));
        assert_eq!(iter.lookup(9), Some(3));
        assert_eq!(iter.lookup(7), None);

        // move back to index 0
        iter.index = 0;

        // test larger than universe
        assert_eq!(iter.lookup(10), None);
        //check for index
        assert_eq!(iter.index, 0);
    }

    #[test]
    fn test_iter_lookup_next() {
        let values = vec![2, 5, 5, 9];
        let seq = EliasFano::from(&values);
        let mut iter = seq.iter();

        // Test lookup_next from start
        assert_eq!(iter.lookup_next(4), 1); // Next value is 5 at index 1
        assert_eq!(iter.lookup_next(6), 3); // Next value is 9 at index 3
        assert_eq!(iter.lookup_next(10), 3); // should return the last index

        // Test lookup_next after moving iterator
        iter.next(); // Move to index 0
        assert_eq!(iter.lookup_next(6), 3); // Next value is 9
        assert_eq!(iter.lookup_next(8), 3); // Next value is 9

        // go to index 0
        iter.index = 0;
        assert_eq!(iter.lookup_next(1), 0);
    }
}