Skip to main content

uqa_core/memory/map/
owned.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7//! Ordinary ordered owners expose exact allocation layouts for later immutable retention.
8
9mod set;
10#[cfg(test)]
11mod tests;
12
13pub use set::{OwnedSet, OwnedSetIntoIter, OwnedSetIter};
14
15use std::{borrow::Borrow, ops::Bound};
16
17use super::{tree, BudgetedMapIter, Link, Node, OwnedNode, MAX_HEIGHT};
18
19/// An ordinary mutable ordered map with exact node ownership. Unlike `BudgetedMap`, insertion has no allowance and cannot be used as an admitted producer by itself. An enclosing controlled producer must reserve `entry_bytes()` before allocating each new entry and retain that lease until the map is dropped. Separately allocated key/value payloads require their own reservations.
20pub struct OwnedMap<K, V> {
21    root: Link<K, V>,
22    len: usize,
23}
24
25impl<K, V> Default for OwnedMap<K, V> {
26    fn default() -> Self {
27        Self { root: None, len: 0 }
28    }
29}
30
31impl<K, V> OwnedMap<K, V> {
32    pub fn new() -> Self {
33        Self::default()
34    }
35
36    /// Exact allocation layout of one node, including its key, value, links and padding. Allocator bookkeeping is outside the payload allowance.
37    pub const fn entry_bytes() -> usize {
38        size_of::<Node<K, V>>()
39    }
40
41    pub fn allocated_bytes(&self) -> usize {
42        self.len * Self::entry_bytes()
43    }
44
45    pub fn len(&self) -> usize {
46        self.len
47    }
48
49    pub fn is_empty(&self) -> bool {
50        self.len == 0
51    }
52
53    pub fn iter(&self) -> BudgetedMapIter<'_, K, V> {
54        BudgetedMapIter::new(&self.root, self.len)
55    }
56
57    pub fn keys(&self) -> impl ExactSizeIterator<Item = &K> + std::iter::FusedIterator {
58        self.iter().map(|(key, _)| key)
59    }
60
61    pub fn values(&self) -> impl ExactSizeIterator<Item = &V> + std::iter::FusedIterator {
62        self.iter().map(|(_, value)| value)
63    }
64}
65
66impl<K: Ord, V> OwnedMap<K, V> {
67    pub fn get<Q: Ord + ?Sized>(&self, key: &Q) -> Option<&V>
68    where
69        K: Borrow<Q>,
70    {
71        tree::get(&self.root, key).map(|node| &node.value)
72    }
73
74    pub fn get_mut<Q: Ord + ?Sized>(&mut self, key: &Q) -> Option<&mut V>
75    where
76        K: Borrow<Q>,
77    {
78        tree::get_mut(&mut self.root, key).map(|node| &mut node.value)
79    }
80
81    pub fn contains_key<Q: Ord + ?Sized>(&self, key: &Q) -> bool
82    where
83        K: Borrow<Q>,
84    {
85        self.get(key).is_some()
86    }
87
88    /// A collision replaces only the value, preserving the original key and node allocation.
89    pub fn insert(&mut self, key: K, value: V) -> Option<V> {
90        if let Some(previous) = self.get_mut(&key) {
91            return Some(std::mem::replace(previous, value));
92        }
93        let entry = OwnedNode {
94            value: Box::new(Node {
95                key,
96                value,
97                left: None,
98                right: None,
99                height: 1,
100            }),
101            memory: None,
102        };
103        let previous = tree::insert(&mut self.root, entry);
104        self.len += 1;
105        previous
106    }
107
108    /// Move entries from `other` without allocating replacement nodes. Collisions retain this map's original key and node, replacing only the value.
109    pub fn append(&mut self, other: Self) {
110        if self.is_empty() {
111            *self = other;
112            return;
113        }
114        let mut entries = other.into_iter();
115        while let Some(node) = entries.next_node() {
116            self.len += usize::from(tree::insert(&mut self.root, node).is_none());
117        }
118    }
119
120    pub fn remove<Q: Ord + ?Sized>(&mut self, key: &Q) -> Option<V>
121    where
122        K: Borrow<Q>,
123    {
124        self.remove_entry(key).map(|(_, value)| value)
125    }
126
127    pub fn remove_entry<Q: Ord + ?Sized>(&mut self, key: &Q) -> Option<(K, V)>
128    where
129        K: Borrow<Q>,
130    {
131        let removed = tree::remove(&mut self.root, key);
132        self.len -= usize::from(removed.is_some());
133        removed
134    }
135
136    /// Borrow the first entry at or after a lower bound in logarithmic time without scratch allocation.
137    pub fn first_from<Q: Ord + ?Sized>(&self, start: Bound<&Q>) -> Option<(&K, &V)>
138    where
139        K: Borrow<Q>,
140    {
141        let mut link = &self.root;
142        let mut candidate = None;
143        while let Some(node) = link {
144            let included = match start {
145                Bound::Unbounded => true,
146                Bound::Included(key) => node.key.borrow() >= key,
147                Bound::Excluded(key) => node.key.borrow() > key,
148            };
149            if included {
150                candidate = Some((&node.key, &node.value.value));
151                link = &node.left;
152            } else {
153                link = &node.right;
154            }
155        }
156        candidate
157    }
158}
159
160impl<K: Ord + Borrow<Q>, V, Q: Ord + ?Sized> std::ops::Index<&Q> for OwnedMap<K, V> {
161    type Output = V;
162
163    fn index(&self, key: &Q) -> &V {
164        self.get(key).expect("missing owned map key")
165    }
166}
167
168impl<K: Ord + Clone, V: Clone> Clone for OwnedMap<K, V> {
169    fn clone(&self) -> Self {
170        self.iter()
171            .map(|(key, value)| (key.clone(), value.clone()))
172            .collect()
173    }
174}
175
176impl<K: std::fmt::Debug, V: std::fmt::Debug> std::fmt::Debug for OwnedMap<K, V> {
177    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
178        f.debug_map().entries(self.iter()).finish()
179    }
180}
181
182impl<K: PartialEq, V: PartialEq> PartialEq for OwnedMap<K, V> {
183    fn eq(&self, other: &Self) -> bool {
184        self.len == other.len && self.iter().eq(other.iter())
185    }
186}
187
188impl<K: Eq, V: Eq> Eq for OwnedMap<K, V> {}
189
190impl<K: Ord, V> FromIterator<(K, V)> for OwnedMap<K, V> {
191    fn from_iter<T: IntoIterator<Item = (K, V)>>(iter: T) -> Self {
192        let mut map = Self::new();
193        for (key, value) in iter {
194            map.insert(key, value);
195        }
196        map
197    }
198}
199
200impl<'a, K, V> IntoIterator for &'a OwnedMap<K, V> {
201    type Item = (&'a K, &'a V);
202    type IntoIter = BudgetedMapIter<'a, K, V>;
203
204    fn into_iter(self) -> Self::IntoIter {
205        self.iter()
206    }
207}
208
209/// Consuming traversal keeps the pending AVL path on the stack and frees each visited node before yielding its entry.
210pub struct OwnedMapIntoIter<K, V> {
211    stack: [Option<OwnedNode<K, V>>; MAX_HEIGHT],
212    depth: usize,
213    remaining: usize,
214}
215
216impl<K, V> OwnedMapIntoIter<K, V> {
217    fn push_left(&mut self, mut link: Link<K, V>) {
218        while let Some(mut node) = link {
219            link = node.value.left.take();
220            self.stack[self.depth] = Some(node);
221            self.depth += 1;
222        }
223    }
224
225    fn next_node(&mut self) -> Option<OwnedNode<K, V>> {
226        self.depth = self.depth.checked_sub(1)?;
227        let mut node = self.stack[self.depth].take().expect("owned traversal node");
228        self.push_left(node.value.right.take());
229        node.value.height = 1;
230        self.remaining -= 1;
231        Some(node)
232    }
233}
234
235impl<K, V> Iterator for OwnedMapIntoIter<K, V> {
236    type Item = (K, V);
237
238    fn next(&mut self) -> Option<Self::Item> {
239        self.next_node().map(tree::into_entry)
240    }
241
242    fn size_hint(&self) -> (usize, Option<usize>) {
243        (self.remaining, Some(self.remaining))
244    }
245}
246
247impl<K, V> ExactSizeIterator for OwnedMapIntoIter<K, V> {}
248impl<K, V> std::iter::FusedIterator for OwnedMapIntoIter<K, V> {}
249
250impl<K, V> IntoIterator for OwnedMap<K, V> {
251    type Item = (K, V);
252    type IntoIter = OwnedMapIntoIter<K, V>;
253
254    fn into_iter(self) -> Self::IntoIter {
255        let mut iter = OwnedMapIntoIter {
256            stack: std::array::from_fn(|_| None),
257            depth: 0,
258            remaining: self.len,
259        };
260        iter.push_left(self.root);
261        iter
262    }
263}