Skip to main content

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    fn trace(&self, visitor: &mut cljrs_gc::MarkVisitor) {
106        for (k, v) in self.inner.iter() {
107            k.trace(visitor);
108            v.trace(visitor);
109        }
110    }
111
112    fn gc_size_extra(&self) -> usize {
113        self.charge.on_alloc(self.full_size_extra())
114    }
115}
116
117#[cfg(test)]
118mod tests {
119    use super::*;
120    use crate::Value;
121
122    fn int(n: i64) -> Value {
123        Value::Long(n)
124    }
125
126    #[test]
127    fn test_basic_ops() {
128        let m = SortedMap::empty();
129        let m = m.assoc(int(3), int(30));
130        let m = m.assoc(int(1), int(10));
131        let m = m.assoc(int(2), int(20));
132        assert_eq!(m.count(), 3);
133        assert_eq!(m.get(&int(1)), Some(&int(10)));
134        assert_eq!(m.get(&int(2)), Some(&int(20)));
135        assert_eq!(m.get(&int(3)), Some(&int(30)));
136    }
137
138    #[test]
139    fn test_sorted_iteration_order() {
140        let m = SortedMap::empty()
141            .assoc(int(3), int(30))
142            .assoc(int(1), int(10))
143            .assoc(int(2), int(20));
144        let keys: Vec<Value> = m.iter().map(|(k, _)| k.clone()).collect();
145        assert_eq!(keys, vec![int(1), int(2), int(3)]);
146    }
147
148    #[test]
149    fn test_dissoc() {
150        let m = SortedMap::empty()
151            .assoc(int(1), int(10))
152            .assoc(int(2), int(20));
153        let m2 = m.dissoc(&int(1));
154        assert_eq!(m2.count(), 1);
155        assert_eq!(m2.get(&int(1)), None);
156        assert_eq!(m2.get(&int(2)), Some(&int(20)));
157    }
158
159    #[test]
160    fn assoc_on_allocated_map_charges_path_bytes() {
161        let _alloc = crate::collections::share_charge::alloc_scope();
162        use cljrs_gc::GcPtr;
163
164        let m = GcPtr::new(SortedMap::from_pairs((0..10_000).map(|i| (int(i), int(i)))));
165        let added = m.get().assoc(int(-1), int(0));
166        assert_eq!(added.charged(), rb_path_bytes::<(Value, Value)>(10_001));
167        assert!(added.charged() * 100 < m.get().full_size_extra());
168    }
169}