fst_incremental 1.0.2

#[cfg(feature = "serde")]
use serde::{Deserialize, Serialize};

/// A lightweight handle to a slice of bytes within the arena.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
struct Span {
  pub start: u32,
  pub len: u32,
  /// If true, this key is a "Delete" operation (Tombstone).
  /// If false, this key is an "Insert" operation.
  pub is_tombstone: bool,
}

/// A memory-efficient, sorted set of byte vectors acting as a unified mutation buffer.
#[derive(Debug, Clone, Default, PartialEq, Eq)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
pub struct CompactArenaSet {
  /// The append-only storage arena.
  data: Vec<u8>,
  /// Ordered list of spans. Sorted based on the *content* they point to in `data`.
  spans: Vec<Span>,
}

impl CompactArenaSet {
  pub fn new() -> Self {
    Self {
      data: Vec::new(),
      spans: Vec::new(),
    }
  }

  /// Returns true if the set contains the key and it is NOT a tombstone.
  pub fn contains(&self, key: &[u8]) -> bool {
    match self.binary_search(key) {
      Ok(idx) => !self.spans[idx].is_tombstone,
      Err(_) => false,
    }
  }

  /// Inserts a key. Returns true if the key's state changed (was absent or was a tombstone).
  pub fn insert(&mut self, key: &[u8]) -> bool {
    match self.binary_search(key) {
      Ok(idx) => {
        if self.spans[idx].is_tombstone {
          self.spans[idx].is_tombstone = false;
          true
        } else {
          false
        }
      }
      Err(idx) => {
        let start = self.data.len() as u32;
        let len = key.len() as u32;
        self.data.extend_from_slice(key);

        self.spans.insert(
          idx,
          Span {
            start,
            len,
            is_tombstone: false,
          },
        );
        true
      }
    }
  }

  /// Removes a key (marks it as a tombstone). Returns true if the key's state changed.
  /// Note: This adds the key to the buffer if it wasn't present, explicitly marking it as deleted.
  pub fn remove(&mut self, key: &[u8]) -> bool {
    match self.binary_search(key) {
      Ok(idx) => {
        if !self.spans[idx].is_tombstone {
          self.spans[idx].is_tombstone = true;
          true
        } else {
          false
        }
      }
      Err(idx) => {
        let start = self.data.len() as u32;
        let len = key.len() as u32;
        self.data.extend_from_slice(key);
        self.spans.insert(
          idx,
          Span {
            start,
            len,
            is_tombstone: true,
          },
        );
        true
      }
    }
  }

  /// Returns the number of items in the set (excluding tombstones).
  pub fn len(&self) -> usize {
    self.spans.iter().filter(|s| !s.is_tombstone).count()
  }

  /// Returns the raw number of entries (including tombstones).
  /// Useful for calculating rebuild thresholds.
  pub fn len_raw(&self) -> usize {
    self.spans.len()
  }

  pub fn is_empty(&self) -> bool {
    // If we have spans but all are tombstones, the set is logically empty.
    self.spans.iter().all(|s| s.is_tombstone)
  }

  /// Clears the set and reclaims all memory in the arena.
  pub fn clear(&mut self) {
    self.data.clear();
    self.spans.clear();
  }

  /// Returns an iterator that yields strictly alive keys (skips tombstones).
  pub fn iter(&self) -> ArenaIter<'_> {
    ArenaIter { parent: self, idx: 0 }
  }

  /// Returns an iterator that yields `(&[u8], is_tombstone: bool)`.
  /// This is used by the merge/rebuild logic to process deletions.
  pub fn iter_raw(&self) -> ArenaRawIter<'_> {
    ArenaRawIter { parent: self, idx: 0 }
  }

  /// Internal binary search helper.
  fn binary_search(&self, key: &[u8]) -> Result<usize, usize> {
    self.spans.binary_search_by(|span| {
      let slice = &self.data[span.start as usize..(span.start + span.len) as usize];
      slice.cmp(key)
    })
  }

  /// Returns the status of a key:
  /// - `Some(true)`: Key exists and is Alive.
  /// - `Some(false)`: Key exists and is a Tombstone (Deleted).
  /// - `None`: Key is not in the buffer.
  pub fn get_status(&self, key: &[u8]) -> Option<bool> {
    match self.binary_search(key) {
      Ok(idx) => Some(!self.spans[idx].is_tombstone),
      Err(_) => None,
    }
  }
  
  pub fn size_in_bytes(&self) -> usize {
    self.data.len() + (self.spans.capacity() * std::mem::size_of::<Span>())
  }
}

impl<T: AsRef<[u8]>> FromIterator<T> for CompactArenaSet {
  fn from_iter<I: IntoIterator<Item = T>>(iter: I) -> Self {
    let mut set = Self::new();
    for item in iter {
      set.insert(item.as_ref());
    }
    set
  }
}

pub struct ArenaIter<'a> {
  parent: &'a CompactArenaSet,
  idx: usize,
}

impl<'a> Iterator for ArenaIter<'a> {
  type Item = &'a [u8];

  fn next(&mut self) -> Option<Self::Item> {
    loop {
      if self.idx >= self.parent.spans.len() {
        return None;
      }

      let span = &self.parent.spans[self.idx];
      self.idx += 1;

      if !span.is_tombstone {
        return Some(&self.parent.data[span.start as usize..(span.start + span.len) as usize]);
      }
    }
  }
}

pub struct ArenaRawIter<'a> {
  parent: &'a CompactArenaSet,
  idx: usize,
}

impl<'a> Iterator for ArenaRawIter<'a> {
  type Item = (&'a [u8], bool); // (Key, IsTombstone)

  fn next(&mut self) -> Option<Self::Item> {
    if self.idx < self.parent.spans.len() {
      let span = &self.parent.spans[self.idx];
      self.idx += 1;
      let slice = &self.parent.data[span.start as usize..(span.start + span.len) as usize];
      Some((slice, span.is_tombstone))
    } else {
      None
    }
  }
}

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

  #[test]
  fn test_insert_and_contains() {
    let mut set = CompactArenaSet::new();
    assert!(set.insert(b"apple"));
    assert!(set.insert(b"banana"));

    assert!(!set.insert(b"apple"));

    assert!(set.contains(b"apple"));
    assert!(set.contains(b"banana"));
    assert!(!set.contains(b"cherry"));
    assert_eq!(set.len(), 2);
  }

  #[test]
  fn test_sorted_iteration() {
    let mut set = CompactArenaSet::new();
    set.insert(b"zebra");
    set.insert(b"apple");
    set.insert(b"mango");
    set.insert(b"banana");

    let collected: Vec<&[u8]> = set.iter().collect();

    // Byte literals are &[u8; N], so the leading cast is what gives the differing lengths a common type.
    let expected: Vec<&[u8]> = vec![b"apple" as &[u8], b"banana", b"mango", b"zebra"];

    assert_eq!(collected, expected);
  }

  #[test]
  fn test_removal() {
    let mut set = CompactArenaSet::new();
    set.insert(b"A");
    set.insert(b"B");
    set.insert(b"C");

    assert!(set.remove(b"B"));
    assert!(!set.remove(b"B"));

    let collected: Vec<&[u8]> = set.iter().collect();
    // Byte literals are &[u8; N], so the leading cast is what gives the differing lengths a common type.
    assert_eq!(collected, vec![b"A" as &[u8], b"C"]);
    assert_eq!(set.len(), 2);
  }

  #[test]
  fn test_tombstone_mechanics() {
    let mut set = CompactArenaSet::new();
    set.insert(b"A");

    set.remove(b"A");
    assert!(!set.contains(b"A"));
    assert_eq!(set.len(), 0);
    let raw: Vec<(&[u8], bool)> = set.iter_raw().collect();
    assert_eq!(raw.len(), 1);
    assert_eq!(raw[0], (b"A" as &[u8], true));

    assert!(set.insert(b"A"));
    assert!(set.contains(b"A"));
    assert_eq!(set.len(), 1);

    let raw_resurrected: Vec<(&[u8], bool)> = set.iter_raw().collect();
    assert_eq!(raw_resurrected[0], (b"A" as &[u8], false));
  }

  #[test]
  fn test_shadowing_deletes() {
    let mut set = CompactArenaSet::new();

    // Remove a key that was NEVER inserted.
    // In a unified buffer model, this must create a tombstone record
    // to ensure we shadow any potential value in the persisted FST.
    assert!(set.remove(b"Ghost"));

    assert!(!set.contains(b"Ghost"));
    assert_eq!(set.len(), 0);

    let raw: Vec<(&[u8], bool)> = set.iter_raw().collect();
    assert_eq!(raw.len(), 1);
    assert_eq!(raw[0], (b"Ghost" as &[u8], true));
  }

  #[test]
  fn test_clear() {
    let mut set = CompactArenaSet::new();
    set.insert(b"foo");
    set.remove(b"bar");
    set.clear();

    assert!(set.is_empty());
    assert_eq!(set.len(), 0);
    assert!(!set.contains(b"foo"));
    assert!(set.data.is_empty());
    assert!(set.spans.is_empty());
    assert_eq!(set.iter_raw().count(), 0);
  }

  #[test]
  #[cfg(feature = "serde")]
  fn test_serde() {
    let mut set = CompactArenaSet::new();
    set.insert(b"hello");
    set.insert(b"world");
    set.remove(b"deleted");

    let serialized = rmp_serde::to_vec(&set).expect("Serialization failed");
    let deserialized: CompactArenaSet = rmp_serde::from_slice(&serialized).expect("Deserialization failed");

    assert_eq!(set, deserialized);
    assert!(deserialized.contains(b"hello"));
    assert!(deserialized.contains(b"world"));
    assert!(!deserialized.contains(b"deleted"));

    let raw: Vec<_> = deserialized.iter_raw().collect();
    assert_eq!(raw.len(), 3); // hello, world, deleted
  }

  #[test]
  fn test_empty_key() {
    let mut set = CompactArenaSet::new();
    assert!(set.insert(b""));
    assert!(set.contains(b""));
    assert!(set.insert(b"a"));

    let collected: Vec<&[u8]> = set.iter().collect();

    // Byte literals are &[u8; N], so the leading cast is what gives the differing lengths a common type.
    assert_eq!(collected, vec![b"" as &[u8], b"a"]);
  }

  // avoid adding 'rand' dependency just for this test block
  struct SimpleRng {
    state: u64,
  }
  impl SimpleRng {
    fn new(seed: u64) -> Self {
      Self { state: seed }
    }
    fn next_u64(&mut self) -> u64 {
      self.state = self.state.wrapping_mul(6364136223846793005).wrapping_add(1);
      self.state
    }
    fn next_bytes(&mut self, max_len: usize) -> Vec<u8> {
      let len = (self.next_u64() as usize) % max_len;
      let mut out = Vec::with_capacity(len);
      for _ in 0..len {
        out.push((self.next_u64() % 256) as u8);
      }
      out
    }
  }

  fn assert_integrity(set: &CompactArenaSet) {
    // Verify Spans are sorted based on Data (regardless of tombstone status)
    for i in 0..set.spans.len().saturating_sub(1) {
      let s1 = &set.spans[i];
      let s2 = &set.spans[i + 1];

      let slice1 = &set.data[s1.start as usize..(s1.start + s1.len) as usize];
      let slice2 = &set.data[s2.start as usize..(s2.start + s2.len) as usize];

      if slice1 >= slice2 {
        panic!(
          "Integrity Failure at index {}: {:?} >= {:?}\nTotal Spans: {}",
          i,
          slice1,
          slice2,
          set.spans.len()
        );
      }
    }

    for span in &set.spans {
      let end = span.start as usize + span.len as usize;
      assert!(
        end <= set.data.len(),
        "Span out of bounds: end={} data_len={}",
        end,
        set.data.len()
      );
    }
  }

  #[test]
  fn test_prefix_and_lexicographical_edge_cases() {
    let mut set = CompactArenaSet::new();

    // Inputs specifically chosen to trip up off-by-one length checks
    // and prefix logic.
    let inputs: Vec<&[u8]> = vec![
      b"a", b"aa", b"aaa", b"ab", b"b", b"ba", b"",     // Empty key
      b"a\0b", // Null bytes
    ];

    // Insert in reverse to force sorting logic to work hard
    for input in inputs.iter().rev() {
      set.insert(input);
    }

    assert_integrity(&set);

    let collected: Vec<Vec<u8>> = set.iter().map(|s| s.to_vec()).collect();

    let expected: Vec<Vec<u8>> = vec![
      b"".to_vec(),
      b"a".to_vec(),
      b"a\0b".to_vec(),
      b"aa".to_vec(),
      b"aaa".to_vec(),
      b"ab".to_vec(),
      b"b".to_vec(),
      b"ba".to_vec(),
    ];

    assert_eq!(collected, expected);
  }

  #[test]
  fn test_fuzz_comparison_against_btreeset() {
    let mut arena_set = CompactArenaSet::new();
    let mut std_set = BTreeSet::new();
    let mut rng = SimpleRng::new(12345);

    let operations = 5000;
    let max_key_len = 16;

    for i in 0..operations {
      if rng.next_u64() % 10 < 7 {
        let key = rng.next_bytes(max_key_len);

        // Semantics: "insert" resurrects tombstones or adds new items.
        // We check if the final state (contains) matches BTreeSet.
        arena_set.insert(&key);
        std_set.insert(key.clone());

        assert!(
          arena_set.contains(&key),
          "Arena should contain key after insert at iter {}",
          i
        );
      } else {
        let key = rng.next_bytes(max_key_len);

        arena_set.remove(&key);
        std_set.remove(&key);

        assert!(
          !arena_set.contains(&key),
          "Arena should NOT contain key after remove at iter {}",
          i
        );
      }

      if i % 500 == 0 {
        assert_integrity(&arena_set);
      }
    }

    // Final equivalence check (Alive items only)
    let arena_vec: Vec<_> = arena_set.iter().collect();
    let std_vec: Vec<_> = std_set.iter().map(|v| v.as_slice()).collect();

    assert_eq!(arena_vec, std_vec, "Sets diverged after fuzzing");
  }

  #[test]
  fn test_append_only_growth_and_churn() {
    let mut set = CompactArenaSet::new();

    set.insert(b"A");
    let data_len_1 = set.data.len();

    set.remove(b"A");
    let data_len_2 = set.data.len();

    // Invariant: Data should NOT shrink (it's append only, or reused)
    assert_eq!(data_len_1, data_len_2);
    assert_eq!(set.len(), 0);

    set.insert(b"A");
    let data_len_3 = set.data.len();

    // OPTIMIZATION CHECK: Data should NOT grow on resurrection.
    // The previous span for "A" was just a tombstone, so insert() should reuse it
    // by flipping the boolean, rather than appending new bytes.
    assert_eq!(data_len_3, data_len_2);

    assert_integrity(&set);
    assert!(set.contains(b"A"));
  }

  #[test]
  fn test_large_keys() {
    let mut set = CompactArenaSet::new();

    let key1 = vec![b'x'; 1024];
    let mut key2 = key1.clone();
    key2.push(b'y');

    set.insert(&key2);
    set.insert(&key1); // Insert smaller second to test sort

    let collected: Vec<_> = set.iter().collect();
    assert_eq!(collected.len(), 2);
    assert_eq!(collected[0], key1.as_slice());
    assert_eq!(collected[1], key2.as_slice());

    assert_integrity(&set);
  }

  #[test]
  fn test_exact_binary_search_boundaries() {
    let mut set = CompactArenaSet::new();
    set.insert(b"10");
    set.insert(b"20");
    set.insert(b"30");

    assert!(set.contains(b"20"));

    assert!(!set.contains(b"00"));

    assert!(!set.contains(b"15"));

    assert!(!set.contains(b"99"));

    // Insert something that falls strictly between existing items in memory
    set.insert(b"15");
    assert_integrity(&set);

    let collected: Vec<_> = set.iter().collect();
    assert_eq!(collected, vec![b"10" as &[u8], b"15", b"20", b"30"]);
  }
}