uqa_core/memory/map/
shared.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7//! Immutable ordered roots share entries and unchanged subtrees under one allowance.
8
9#[cfg(test)]
10mod tests;
11mod tree;
12
13use std::{borrow::Borrow, cmp::Ordering, ops::Bound, sync::Arc};
14
15use super::{Budgeted, MemoryBudget, MemoryError, MAX_HEIGHT};
16
17type Entry<K, V> = Arc<Budgeted<(K, V)>>;
18type SharedNode<K, V> = Arc<Budgeted<Node<K, V>>>;
19type Link<K, V> = Option<SharedNode<K, V>>;
20
21struct Node<K, V> {
22    entry: Entry<K, V>,
23    left: Link<K, V>,
24    right: Link<K, V>,
25    height: u8,
26}
27
28/// A persistent ordered map with logarithmic lookup and insertion. Cloning shares a root without copying entries or allocating. Updates reserve the complete insertion before mutation, reuse unique nodes and copy shared paths; unchanged entries and subtrees keep their original reservations until their last root is dropped. Keys and values retain their separately owned payloads.
29pub struct BudgetedSharedMap<K, V> {
30    root: Link<K, V>,
31    len: usize,
32    memory: MemoryBudget,
33}
34
35impl<K, V> Clone for BudgetedSharedMap<K, V> {
36    fn clone(&self) -> Self {
37        Self {
38            root: self.root.clone(),
39            len: self.len,
40            memory: self.memory.clone(),
41        }
42    }
43}
44
45impl<K, V> BudgetedSharedMap<K, V> {
46    pub fn new(memory: &MemoryBudget) -> Self {
47        Self {
48            root: None,
49            len: 0,
50            memory: memory.clone(),
51        }
52    }
53
54    pub fn len(&self) -> usize {
55        self.len
56    }
57
58    pub fn is_empty(&self) -> bool {
59        self.len == 0
60    }
61
62    pub fn budget(&self) -> &MemoryBudget {
63        &self.memory
64    }
65
66    pub fn iter(&self) -> BudgetedSharedMapIter<'_, K, V> {
67        let mut iter = BudgetedSharedMapIter::empty();
68        iter.push_left(self.root.as_ref().map(|node| &***node));
69        iter
70    }
71}
72
73impl<K: Ord, V> BudgetedSharedMap<K, V> {
74    pub fn get<Q: Ord + ?Sized>(&self, key: &Q) -> Option<&V>
75    where
76        K: Borrow<Q>,
77    {
78        let mut link = &self.root;
79        while let Some(node) = link {
80            match key.cmp(node.entry.0.borrow()) {
81                Ordering::Less => link = &node.left,
82                Ordering::Greater => link = &node.right,
83                Ordering::Equal => return Some(&node.entry.1),
84            }
85        }
86        None
87    }
88
89    /// Return a new root containing the supplied key and value, leaving this root unchanged even if reservation fails. Matching keys are replaced together with their values; neither type needs to implement `Clone`.
90    pub fn with_insert(&self, key: K, value: V) -> Result<Self, MemoryError> {
91        let mut candidate = self.clone();
92        candidate.try_insert(key, value)?;
93        Ok(candidate)
94    }
95
96    /// Reserve the complete insertion before changing this root, reusing unique nodes and copying only shared paths. Reservation failure preserves this root; existing clones remain unchanged. Matching keys and values are replaced together without requiring `Clone`.
97    pub fn try_insert(&mut self, key: K, value: V) -> Result<(), MemoryError> {
98        let added = tree::insert(&mut self.root, key, value, &self.memory)?;
99        self.len += usize::from(added);
100        Ok(())
101    }
102
103    /// Visit keys in order, seeking the lower bound without traversing earlier entries or allocating a traversal buffer.
104    pub fn range_from<Q: Ord + ?Sized>(&self, start: Bound<&Q>) -> BudgetedSharedMapIter<'_, K, V>
105    where
106        K: Borrow<Q>,
107    {
108        let mut iter = BudgetedSharedMapIter::empty();
109        let mut link = &self.root;
110        while let Some(node) = link {
111            let included = match start {
112                Bound::Unbounded => true,
113                Bound::Included(key) => node.entry.0.borrow() >= key,
114                Bound::Excluded(key) => node.entry.0.borrow() > key,
115            };
116            if included {
117                iter.stack[iter.depth] = Some(node);
118                iter.depth += 1;
119                link = &node.left;
120            } else {
121                link = &node.right;
122            }
123        }
124        iter
125    }
126}
127
128pub struct BudgetedSharedMapIter<'a, K, V> {
129    stack: [Option<&'a Node<K, V>>; MAX_HEIGHT],
130    depth: usize,
131}
132
133impl<'a, K, V> BudgetedSharedMapIter<'a, K, V> {
134    fn empty() -> Self {
135        Self {
136            stack: [None; MAX_HEIGHT],
137            depth: 0,
138        }
139    }
140
141    fn push_left(&mut self, mut node: Option<&'a Node<K, V>>) {
142        while let Some(current) = node {
143            self.stack[self.depth] = Some(current);
144            self.depth += 1;
145            node = current.left.as_ref().map(|node| &***node);
146        }
147    }
148}
149
150impl<'a, K, V> Iterator for BudgetedSharedMapIter<'a, K, V> {
151    type Item = (&'a K, &'a V);
152
153    fn next(&mut self) -> Option<Self::Item> {
154        self.depth = self.depth.checked_sub(1)?;
155        let node = self.stack[self.depth].take().expect("retained map node");
156        self.push_left(node.right.as_ref().map(|node| &***node));
157        Some((&node.entry.0, &node.entry.1))
158    }
159}
160
161impl<K, V> std::iter::FusedIterator for BudgetedSharedMapIter<'_, K, V> {}
162
163impl<'a, K, V> IntoIterator for &'a BudgetedSharedMap<K, V> {
164    type Item = (&'a K, &'a V);
165    type IntoIter = BudgetedSharedMapIter<'a, K, V>;
166
167    fn into_iter(self) -> Self::IntoIter {
168        self.iter()
169    }
170}
uqa_core/memory/map/shared.rs

uqa_core/memory/map/
shared.rs
