use std::collections::{HashMap, HashSet};
use std::hash::Hash;
use std::sync::Arc;
use petgraph::graph::NodeIndex;
type Level<K> = HashMap<K, Option<Vec<NodeIndex>>>;
pub(super) const MAX_LAYER_DEPTH: usize = 32;
#[derive(Debug, Clone)]
pub struct LayeredIndex<K: Eq + Hash + Clone> {
levels: Vec<Arc<Level<K>>>,
}
impl<K: Eq + Hash + Clone> Default for LayeredIndex<K> {
fn default() -> Self {
Self { levels: Vec::new() }
}
}
impl<K: Eq + Hash + Clone> From<HashMap<K, Vec<NodeIndex>>> for LayeredIndex<K> {
fn from(map: HashMap<K, Vec<NodeIndex>>) -> Self {
Self {
levels: vec![Arc::new(
map.into_iter().map(|(k, v)| (k, Some(v))).collect(),
)],
}
}
}
impl<K: Eq + Hash + Clone> FromIterator<(K, Vec<NodeIndex>)> for LayeredIndex<K> {
fn from_iter<I: IntoIterator<Item = (K, Vec<NodeIndex>)>>(iter: I) -> Self {
Self::from(iter.into_iter().collect::<HashMap<K, Vec<NodeIndex>>>())
}
}
impl<K: Eq + Hash + Clone> LayeredIndex<K> {
pub fn get(&self, key: &K) -> Option<&Vec<NodeIndex>> {
match self.levels.as_slice() {
[] => None,
[only] => only.get(key)?.as_ref(),
levels => {
for level in levels.iter().rev() {
if let Some(entry) = level.get(key) {
return entry.as_ref();
}
}
None
}
}
}
#[inline]
pub fn contains_key(&self, key: &K) -> bool {
self.get(key).is_some()
}
pub fn len(&self) -> usize {
match self.levels.as_slice() {
[] => 0,
[only] => only.len(),
_ => self.iter().count(),
}
}
#[inline]
pub fn is_empty(&self) -> bool {
self.len() == 0
}
pub fn iter(&self) -> IndexIter<'_, K> {
match self.levels.as_slice() {
[] => IndexIter::Empty,
[only] => IndexIter::Flat(only.iter()),
levels => {
let mut seen: HashSet<&K> = HashSet::new();
let mut out: Vec<(&K, &Vec<NodeIndex>)> = Vec::new();
for level in levels.iter().rev() {
for (key, entry) in level.iter() {
if !seen.insert(key) {
continue;
}
if let Some(members) = entry {
out.push((key, members));
}
}
}
IndexIter::Merged(out.into_iter())
}
}
}
pub fn get_mut(&mut self, key: &K) -> Option<&mut Vec<NodeIndex>> {
if self.is_flat() {
return Arc::get_mut(&mut self.levels[0])
.expect("is_flat proved sole ownership")
.get_mut(key)
.and_then(|entry| entry.as_mut());
}
self.get(key)?;
Some(self.materialize(key))
}
pub fn entry_or_default(&mut self, key: &K) -> &mut Vec<NodeIndex> {
if self.is_flat() {
return Arc::get_mut(&mut self.levels[0])
.expect("is_flat proved sole ownership")
.entry(key.clone())
.or_insert_with(|| Some(Vec::new()))
.get_or_insert_with(Vec::new);
}
self.materialize(key)
}
pub fn remove(&mut self, key: &K) -> Option<Vec<NodeIndex>> {
if self.is_flat() {
return Arc::get_mut(&mut self.levels[0])
.expect("is_flat proved sole ownership")
.remove(key)
.flatten();
}
let prior = self.get(key).cloned();
prior.as_ref()?;
self.writable_tail().insert(key.clone(), None);
prior
}
pub fn clear(&mut self) {
if self.is_flat() {
Arc::get_mut(&mut self.levels[0])
.expect("is_flat proved sole ownership")
.clear();
return;
}
let keys: Vec<K> = self.iter().map(|(key, _)| key.clone()).collect();
let tail = self.writable_tail();
for key in keys {
tail.insert(key, None);
}
}
pub fn retain_members<F: FnMut(&NodeIndex) -> bool + Copy>(&mut self, predicate: F) {
self.flatten();
let Some(level) = self.levels.first_mut() else {
return;
};
let level = Arc::get_mut(level).expect("a flattened index owns its single level");
for members in level.values_mut().flatten() {
members.retain(predicate);
}
}
pub fn try_compact(&mut self) {
if self.levels.len() <= 1 {
return;
}
let Some(first_owned) = self
.levels
.iter_mut()
.position(|level| Arc::get_mut(level).is_some())
else {
return;
};
if first_owned + 1 >= self.levels.len() {
return;
}
let mut tail = self.levels.split_off(first_owned).into_iter();
let first = tail.next().expect("split_off yields at least one level");
let mut merged = Arc::try_unwrap(first)
.unwrap_or_else(|_| unreachable!("position proved sole ownership"));
for level in tail {
match Arc::try_unwrap(level) {
Ok(owned) => merged.extend(owned),
Err(shared) => {
merged.extend(shared.iter().map(|(k, v)| (k.clone(), v.clone())));
}
}
}
self.levels.push(Arc::new(merged));
}
pub fn to_map(&self) -> HashMap<K, Vec<NodeIndex>> {
self.iter()
.map(|(key, members)| (key.clone(), members.clone()))
.collect()
}
#[inline]
fn is_flat(&mut self) -> bool {
self.levels.len() == 1 && Arc::get_mut(&mut self.levels[0]).is_some()
}
fn writable_tail(&mut self) -> &mut Level<K> {
let tail_is_ours = self
.levels
.last_mut()
.is_some_and(|tail| Arc::get_mut(tail).is_some());
if !tail_is_ours {
if self.levels.len() >= MAX_LAYER_DEPTH {
self.flatten();
} else {
self.levels.push(Arc::new(Level::default()));
}
}
Arc::get_mut(self.levels.last_mut().expect("just ensured a level"))
.expect("the tail is owned by construction")
}
fn materialize(&mut self, key: &K) -> &mut Vec<NodeIndex> {
let tail_has = self
.levels
.last_mut()
.and_then(Arc::get_mut)
.is_some_and(|tail| tail.contains_key(key));
let merged = if tail_has {
None
} else {
self.get(key).cloned()
};
let tail = self.writable_tail();
tail.entry(key.clone())
.or_insert_with(|| Some(merged.unwrap_or_default()))
.get_or_insert_with(Vec::new)
}
fn flatten(&mut self) {
if self.levels.len() <= 1 {
if self.levels.len() == 1 && Arc::get_mut(&mut self.levels[0]).is_none() {
let copy = self.levels[0].as_ref().clone();
self.levels[0] = Arc::new(copy);
}
return;
}
let mut levels = std::mem::take(&mut self.levels).into_iter();
let first = levels.next().expect("length checked above");
let mut merged = Arc::try_unwrap(first).unwrap_or_else(|shared| shared.as_ref().clone());
for level in levels {
match Arc::try_unwrap(level) {
Ok(owned) => merged.extend(owned),
Err(shared) => merged.extend(shared.iter().map(|(k, v)| (k.clone(), v.clone()))),
}
}
merged.retain(|_, entry| entry.is_some());
self.levels.push(Arc::new(merged));
}
#[cfg(test)]
pub(crate) fn depth(&self) -> usize {
self.levels.len()
}
}
pub enum IndexIter<'a, K> {
Empty,
Flat(std::collections::hash_map::Iter<'a, K, Option<Vec<NodeIndex>>>),
Merged(std::vec::IntoIter<(&'a K, &'a Vec<NodeIndex>)>),
}
impl<'a, K> Iterator for IndexIter<'a, K> {
type Item = (&'a K, &'a Vec<NodeIndex>);
fn next(&mut self) -> Option<Self::Item> {
match self {
IndexIter::Empty => None,
IndexIter::Flat(it) => it.find_map(|(key, entry)| Some((key, entry.as_ref()?))),
IndexIter::Merged(it) => it.next(),
}
}
}
#[cfg(test)]
mod tests {
use super::*;
fn idx(i: usize) -> NodeIndex {
NodeIndex::new(i)
}
fn index(pairs: &[(&str, &[usize])]) -> LayeredIndex<String> {
pairs
.iter()
.map(|(key, members)| {
(
(*key).to_string(),
members.iter().copied().map(idx).collect(),
)
})
.collect()
}
fn sorted(index: &LayeredIndex<String>) -> Vec<(String, Vec<usize>)> {
let mut out: Vec<(String, Vec<usize>)> = index
.iter()
.map(|(k, v)| (k.clone(), v.iter().map(|i| i.index()).collect()))
.collect();
out.sort();
out
}
#[test]
fn reversing_an_append_spares_a_pre_statement_occurrence_in_a_shared_level() {
let mut writer = index(&[("v", &[3, 7, 3])]);
let reader = writer.clone();
writer.entry_or_default(&"v".to_string()).push(idx(3));
assert_eq!(sorted(&writer), vec![("v".into(), vec![3, 7, 3, 3])]);
let members = writer.get_mut(&"v".to_string()).expect("bucket present");
let pos = members
.iter()
.rposition(|member| *member == idx(3))
.expect("the appended member");
members.remove(pos);
assert_eq!(
sorted(&writer),
vec![("v".into(), vec![3, 7, 3])],
"both pre-statement occurrences must survive, in order"
);
assert_eq!(
sorted(&reader),
vec![("v".into(), vec![3, 7, 3])],
"the reader's bucket must be untouched by either edit"
);
}
#[test]
fn reversing_an_eviction_restores_the_recorded_position_in_a_shared_bucket() {
let mut writer = index(&[("v", &[10, 11, 12, 13])]);
let reader = writer.clone();
let key = "v".to_string();
let pos = writer
.get(&key)
.unwrap()
.iter()
.position(|member| *member == idx(12))
.expect("member present");
assert_eq!(pos, 2);
writer
.get_mut(&key)
.unwrap()
.retain(|member| *member != idx(12));
assert_eq!(sorted(&writer), vec![("v".into(), vec![10, 11, 13])]);
let members = writer.entry_or_default(&key);
let pos = pos.min(members.len());
members.insert(pos, idx(12));
assert_eq!(
sorted(&writer),
vec![("v".into(), vec![10, 11, 12, 13])],
"the member must return to its own position, not to the end"
);
assert_eq!(sorted(&reader), vec![("v".into(), vec![10, 11, 12, 13])]);
}
#[test]
fn a_fork_shares_the_buckets_and_isolates_the_edits() {
let mut writer = index(&[("a", &[1]), ("b", &[2, 3])]);
let reader = writer.clone();
writer.entry_or_default(&"a".to_string()).push(idx(9));
writer.entry_or_default(&"c".to_string()).push(idx(4));
writer.remove(&"b".to_string());
assert_eq!(
sorted(&writer),
vec![("a".into(), vec![1, 9]), ("c".into(), vec![4])]
);
assert_eq!(
sorted(&reader),
vec![("a".into(), vec![1]), ("b".into(), vec![2, 3])],
"the reader must see its own snapshot, tombstone and all"
);
assert_eq!(writer.len(), 2);
assert_eq!(reader.len(), 2);
assert!(!writer.contains_key(&"b".to_string()));
assert!(reader.contains_key(&"b".to_string()));
}
#[test]
fn a_tombstone_masks_the_base_and_can_be_written_over() {
let mut writer = index(&[("a", &[1])]);
let reader = writer.clone();
writer.remove(&"a".to_string());
assert_eq!(writer.get(&"a".to_string()), None);
assert_eq!(writer.len(), 0);
assert!(writer.is_empty());
writer.entry_or_default(&"a".to_string()).push(idx(5));
assert_eq!(
sorted(&writer),
vec![("a".into(), vec![5])],
"a write over a tombstone must not resurrect the base's members"
);
assert_eq!(sorted(&reader), vec![("a".into(), vec![1])]);
}
#[test]
fn removing_an_absent_value_is_a_no_op() {
let mut writer = index(&[("a", &[1])]);
assert_eq!(writer.remove(&"zz".to_string()), None);
assert_eq!(writer.len(), 1);
assert_eq!(sorted(&writer), vec![("a".into(), vec![1])]);
}
#[test]
fn an_unshared_index_never_grows_a_level() {
let mut solo = index(&[("a", &[1])]);
for i in 0..50 {
solo.entry_or_default(&format!("k{i}")).push(idx(i));
}
assert_eq!(solo.depth(), 1);
assert_eq!(solo.len(), 51);
}
#[test]
fn compaction_waits_for_the_reader_and_then_folds() {
let mut writer = index(&[("a", &[1])]);
let reader = writer.clone();
writer.entry_or_default(&"b".to_string()).push(idx(2));
assert_eq!(writer.depth(), 2);
writer.try_compact();
assert_eq!(writer.depth(), 2, "a live reader must block the fold");
assert_eq!(sorted(&reader), vec![("a".into(), vec![1])]);
drop(reader);
writer.try_compact();
assert_eq!(writer.depth(), 1);
assert_eq!(
sorted(&writer),
vec![("a".into(), vec![1]), ("b".into(), vec![2])]
);
}
#[test]
fn a_never_compacted_stack_stays_bounded_and_correct() {
let mut writer = index(&[("a", &[0])]);
let mut readers = Vec::new();
for round in 1..=(MAX_LAYER_DEPTH * 3) {
readers.push(writer.clone());
writer
.entry_or_default(&format!("k{round}"))
.push(idx(round));
}
assert!(
writer.depth() <= MAX_LAYER_DEPTH,
"level stack {} exceeded the cap",
writer.depth()
);
assert_eq!(writer.len(), MAX_LAYER_DEPTH * 3 + 1);
for round in 1..=(MAX_LAYER_DEPTH * 3) {
assert_eq!(
writer.get(&format!("k{round}")).map(Vec::as_slice),
Some([idx(round)].as_slice()),
"value k{round} lost in the stack"
);
}
assert_eq!(sorted(&readers[0]), vec![("a".into(), vec![0])]);
drop(readers);
writer.try_compact();
assert_eq!(writer.depth(), 1);
assert_eq!(writer.len(), MAX_LAYER_DEPTH * 3 + 1);
}
#[test]
fn the_delete_sweep_flattens_without_disturbing_the_reader() {
let mut writer = index(&[("a", &[1, 2]), ("b", &[2])]);
let reader = writer.clone();
writer.entry_or_default(&"c".to_string()).push(idx(2));
writer.retain_members(|member| *member != idx(2));
assert_eq!(
sorted(&writer),
vec![
("a".into(), vec![1]),
("b".into(), Vec::new()),
("c".into(), Vec::new())
],
"emptied buckets stay, exactly as the plain-map sweep left them"
);
assert_eq!(
sorted(&reader),
vec![("a".into(), vec![1, 2]), ("b".into(), vec![2])]
);
}
}