cljrs_value/collections/
sorted_map.rs

1use crate::Value;
2use crate::collections::share_charge::{ShareCharge, rb_path_bytes};
3
4/// An immutable sorted map backed by `rpds::RedBlackTreeMapSync`.
5///
6/// A map derived by `assoc` or `dissoc` from a GC-allocated map reports only
7/// its path-copied nodes from `gc_size_extra`; see
8/// `collections::share_charge` for the full rule.
9#[derive(Debug, Clone)]
10pub struct SortedMap {
11    inner: rpds::RedBlackTreeMapSync<Value, Value>,
12    charge: ShareCharge,
13}
14
15impl SortedMap {
16    pub fn empty() -> Self {
17        Self {
18            inner: rpds::RedBlackTreeMapSync::new_sync(),
19            charge: ShareCharge::full(),
20        }
21    }
22
23    pub fn count(&self) -> usize {
24        self.inner.size()
25    }
26
27    pub fn is_empty(&self) -> bool {
28        self.inner.is_empty()
29    }
30
31    pub fn get(&self, key: &Value) -> Option<&Value> {
32        self.inner.get(key)
33    }
34
35    pub fn contains_key(&self, key: &Value) -> bool {
36        self.inner.contains_key(key)
37    }
38
39    pub fn assoc(&self, key: Value, value: Value) -> Self {
40        self.derived(self.inner.insert(key, value))
41    }
42
43    pub fn dissoc(&self, key: &Value) -> Self {
44        self.derived(self.inner.remove(key))
45    }
46
47    /// Wrap `inner`, produced by one path-copying update of `self.inner`.
48    fn derived(&self, inner: rpds::RedBlackTreeMapSync<Value, Value>) -> Self {
49        let n = inner.size().max(self.inner.size());
50        Self {
51            charge: self.charge.derive(
52                self.inner.size(),
53                inner.size(),
54                rb_path_bytes::<(Value, Value)>(n),
55            ),
56            inner,
57        }
58    }
59
60    /// Estimated bytes owned by this map if it shared no nodes.
61    fn full_size_extra(&self) -> usize {
62        // Per entry: Arc<Node<K,V>> (16 overhead) + left/right ptrs (16) + color+pad (8).
63        self.inner.size() * (40 + 2 * std::mem::size_of::<Value>())
64    }
65
66    /// What `gc_size_extra` would report, without marking this version
67    /// allocated.
68    #[cfg(test)]
69    fn charged(&self) -> usize {
70        self.charge.size_extra(self.full_size_extra())
71    }
72
73    /// Iterate over all `(key, value)` pairs in sorted key order.
74    pub fn iter(&self) -> impl Iterator<Item = (&Value, &Value)> {
75        self.inner.iter()
76    }
77
78    pub fn keys(&self) -> Vec<Value> {
79        self.inner.keys().cloned().collect()
80    }
81
82    pub fn vals(&self) -> Vec<Value> {
83        self.inner.values().cloned().collect()
84    }
85
86    pub fn from_pairs<I: IntoIterator<Item = (Value, Value)>>(iter: I) -> Self {
87        let mut m = Self::empty();
88        for (k, v) in iter {
89            m = m.assoc(k, v);
90        }
91        m
92    }
93}
94
95impl PartialEq for SortedMap {
96    fn eq(&self, other: &Self) -> bool {
97        if self.count() != other.count() {
98            return false;
99        }
100        self.inner.iter().all(|(k, v)| other.get(k) == Some(v))
101    }
102}
103
104impl cljrs_gc::Trace for SortedMap {
105    #[cfg(feature = "no-gc")]
106    fn immutable_children(&self) -> bool {
107        true
108    }
109    fn trace(&self, visitor: &mut cljrs_gc::MarkVisitor) {
110        for (k, v) in self.inner.iter() {
111            k.trace(visitor);
112            v.trace(visitor);
113        }
114    }
115
116    fn gc_size_extra(&self) -> usize {
117        self.charge.on_alloc(self.full_size_extra())
118    }
119}
120
121#[cfg(test)]
122mod tests {
123    use super::*;
124    use crate::Value;
125
126    fn int(n: i64) -> Value {
127        Value::Long(n)
128    }
129
130    #[test]
131    fn test_basic_ops() {
132        let m = SortedMap::empty();
133        let m = m.assoc(int(3), int(30));
134        let m = m.assoc(int(1), int(10));
135        let m = m.assoc(int(2), int(20));
136        assert_eq!(m.count(), 3);
137        assert_eq!(m.get(&int(1)), Some(&int(10)));
138        assert_eq!(m.get(&int(2)), Some(&int(20)));
139        assert_eq!(m.get(&int(3)), Some(&int(30)));
140    }
141
142    #[test]
143    fn test_sorted_iteration_order() {
144        let m = SortedMap::empty()
145            .assoc(int(3), int(30))
146            .assoc(int(1), int(10))
147            .assoc(int(2), int(20));
148        let keys: Vec<Value> = m.iter().map(|(k, _)| k.clone()).collect();
149        assert_eq!(keys, vec![int(1), int(2), int(3)]);
150    }
151
152    #[test]
153    fn test_dissoc() {
154        let m = SortedMap::empty()
155            .assoc(int(1), int(10))
156            .assoc(int(2), int(20));
157        let m2 = m.dissoc(&int(1));
158        assert_eq!(m2.count(), 1);
159        assert_eq!(m2.get(&int(1)), None);
160        assert_eq!(m2.get(&int(2)), Some(&int(20)));
161    }
162
163    #[test]
164    fn assoc_on_allocated_map_charges_path_bytes() {
165        let _alloc = crate::collections::share_charge::alloc_scope();
166        use cljrs_gc::GcPtr;
167
168        let m = GcPtr::new(SortedMap::from_pairs((0..10_000).map(|i| (int(i), int(i)))));
169        let added = m.get().assoc(int(-1), int(0));
170        assert_eq!(added.charged(), rb_path_bytes::<(Value, Value)>(10_001));
171        assert!(added.charged() * 100 < m.get().full_size_extra());
172    }
173}
cljrs_value/collections/sorted_map.rs

cljrs_value/collections/
sorted_map.rs
