uqa_core/memory/
map.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7//! Ordered maps reserve complete AVL nodes before publication, including links and padding.
8
9mod owned;
10mod shared;
11#[cfg(test)]
12mod tests;
13mod tree;
14
15pub use owned::{OwnedMap, OwnedMapIntoIter, OwnedSet, OwnedSetIntoIter, OwnedSetIter};
16pub use shared::{BudgetedSharedMap, BudgetedSharedMapIter};
17
18use std::borrow::Borrow;
19
20use super::{Budgeted, MemoryBudget, MemoryError, MemoryReservation};
21
22struct OwnedNode<K, V> {
23    value: Box<Node<K, V>>,
24    // BudgetedMap nodes carry a lease; OwnedMap leaves admission to its enclosing owner.
25    memory: Option<MemoryReservation>,
26}
27
28impl<K, V> std::ops::Deref for OwnedNode<K, V> {
29    type Target = Box<Node<K, V>>;
30
31    fn deref(&self) -> &Self::Target {
32        &self.value
33    }
34}
35
36impl<K, V> OwnedNode<K, V> {
37    fn into_parts(self) -> (Box<Node<K, V>>, Option<MemoryReservation>) {
38        (self.value, self.memory)
39    }
40}
41type Link<K, V> = Option<OwnedNode<K, V>>;
42
43struct Node<K, V> {
44    key: K,
45    value: V,
46    left: Link<K, V>,
47    right: Link<K, V>,
48    height: u8,
49}
50
51/// An ordered map with logarithmic lookup, insertion and removal. Each node reserves its entire allocation before construction; keys and values retain their own separately allocated payloads. Removing a node frees its allocation before releasing that reservation. The map never clones elements or retains unused nodes.
52pub struct BudgetedMap<K, V> {
53    root: Link<K, V>,
54    len: usize,
55    memory: MemoryBudget,
56}
57
58/// An unpublished node and its reservation. Preparing several entries allows an owner to finish every fallible reservation before publishing any map mutation.
59pub struct PreparedMapEntry<K, V> {
60    node: OwnedNode<K, V>,
61}
62
63impl<K, V> BudgetedMap<K, V> {
64    pub fn new(memory: &MemoryBudget) -> Self {
65        Self {
66            root: None,
67            len: 0,
68            memory: memory.clone(),
69        }
70    }
71
72    pub fn len(&self) -> usize {
73        self.len
74    }
75
76    pub fn is_empty(&self) -> bool {
77        self.len == 0
78    }
79
80    pub fn budget(&self) -> &MemoryBudget {
81        &self.memory
82    }
83
84    pub fn prepare_entry(&self, key: K, value: V) -> Result<PreparedMapEntry<K, V>, MemoryError> {
85        let memory = self.memory.reserve(size_of::<Node<K, V>>())?;
86        let node = Box::new(Node {
87            key,
88            value,
89            left: None,
90            right: None,
91            height: 1,
92        });
93        Ok(PreparedMapEntry {
94            node: OwnedNode {
95                value: node,
96                memory: Some(memory),
97            },
98        })
99    }
100
101    /// Visit mutable values in key order without allocating a traversal buffer or allowing keys to change.
102    pub fn for_each_mut(&mut self, mut visit: impl FnMut(&K, &mut V)) {
103        tree::for_each_mut(&mut self.root, &mut visit);
104    }
105
106    pub fn iter(&self) -> BudgetedMapIter<'_, K, V> {
107        BudgetedMapIter::new(&self.root, self.len)
108    }
109}
110
111impl<K: Ord, V> BudgetedMap<K, V> {
112    pub fn get<Q: Ord + ?Sized>(&self, key: &Q) -> Option<&V>
113    where
114        K: Borrow<Q>,
115    {
116        tree::get(&self.root, key).map(|node| &node.value)
117    }
118
119    pub fn get_mut<Q: Ord + ?Sized>(&mut self, key: &Q) -> Option<&mut V>
120    where
121        K: Borrow<Q>,
122    {
123        tree::get_mut(&mut self.root, key).map(|node| &mut node.value)
124    }
125
126    pub fn contains_key<Q: Ord + ?Sized>(&self, key: &Q) -> bool
127    where
128        K: Borrow<Q>,
129    {
130        self.get(key).is_some()
131    }
132
133    /// Publish an already reserved entry without allocation. A matching key retains its original key and returns the replaced value. Panics before mutation if the entry belongs to a different allowance.
134    pub fn insert_prepared(&mut self, entry: PreparedMapEntry<K, V>) -> Option<V> {
135        assert!(
136            self.memory.shares_allowance(
137                entry
138                    .node
139                    .memory
140                    .as_ref()
141                    .expect("prepared node is admitted")
142                    .budget()
143            ),
144            "different memory allowances"
145        );
146        let previous = tree::insert(&mut self.root, entry.node);
147        self.len += usize::from(previous.is_none());
148        previous
149    }
150
151    /// Reserve a node only for a new key. Failure preserves the map, its entries and its reservations.
152    pub fn insert(&mut self, key: K, value: V) -> Result<Option<V>, MemoryError> {
153        if let Some(previous) = self.get_mut(&key) {
154            return Ok(Some(std::mem::replace(previous, value)));
155        }
156        let entry = self.prepare_entry(key, value)?;
157        Ok(self.insert_prepared(entry))
158    }
159
160    pub fn remove<Q: Ord + ?Sized>(&mut self, key: &Q) -> Option<V>
161    where
162        K: Borrow<Q>,
163    {
164        let removed = tree::remove(&mut self.root, key);
165        self.len -= usize::from(removed.is_some());
166        removed.map(|(_, value)| value)
167    }
168}
169
170impl<K: Ord + Borrow<Q>, V, Q: Ord + ?Sized> std::ops::Index<&Q> for BudgetedMap<K, V> {
171    type Output = V;
172
173    fn index(&self, key: &Q) -> &V {
174        self.get(key).expect("missing budgeted map key")
175    }
176}
177
178// Every two AVL levels at least double the minimum node count. Even zero-sized keys and values need nonzero links, so this exceeds every addressable tree's height.
179const MAX_HEIGHT: usize = usize::BITS as usize * 2;
180
181pub struct BudgetedMapIter<'a, K, V> {
182    stack: [Option<&'a Node<K, V>>; MAX_HEIGHT],
183    depth: usize,
184    remaining: usize,
185}
186
187impl<'a, K, V> BudgetedMapIter<'a, K, V> {
188    fn new(root: &'a Link<K, V>, len: usize) -> Self {
189        let mut iter = Self {
190            stack: [None; MAX_HEIGHT],
191            depth: 0,
192            remaining: len,
193        };
194        iter.push_left(root.as_ref().map(|node| &***node));
195        iter
196    }
197
198    fn push_left(&mut self, mut node: Option<&'a Node<K, V>>) {
199        while let Some(current) = node {
200            self.stack[self.depth] = Some(current);
201            self.depth += 1;
202            node = current.left.as_ref().map(|node| &***node);
203        }
204    }
205}
206
207impl<'a, K, V> Iterator for BudgetedMapIter<'a, K, V> {
208    type Item = (&'a K, &'a V);
209
210    fn next(&mut self) -> Option<Self::Item> {
211        self.depth = self.depth.checked_sub(1)?;
212        let node = self.stack[self.depth].take().expect("retained map node");
213        self.push_left(node.right.as_ref().map(|node| &***node));
214        self.remaining -= 1;
215        Some((&node.key, &node.value))
216    }
217
218    fn size_hint(&self) -> (usize, Option<usize>) {
219        (self.remaining, Some(self.remaining))
220    }
221}
222
223impl<K, V> ExactSizeIterator for BudgetedMapIter<'_, K, V> {}
224impl<K, V> std::iter::FusedIterator for BudgetedMapIter<'_, K, V> {}
225
226impl<'a, K, V> IntoIterator for &'a BudgetedMap<K, V> {
227    type Item = (&'a K, &'a V);
228    type IntoIter = BudgetedMapIter<'a, K, V>;
229
230    fn into_iter(self) -> Self::IntoIter {
231        self.iter()
232    }
233}
uqa_core/memory/map.rs

uqa_core/memory/
map.rs
