Skip to main content

cljrs_value/collections/
sorted_set.rs

1use crate::Value;
2use crate::collections::share_charge::{ShareCharge, rb_path_bytes};
3
4/// An immutable sorted set backed by `rpds::RedBlackTreeSet`.
5///
6/// A set derived by `conj` or `disj` from a GC-allocated set reports only its
7/// path-copied nodes from `gc_size_extra`; see `collections::share_charge`
8/// for the full rule.
9#[derive(Debug, Clone)]
10pub struct SortedSet {
11    inner: rpds::RedBlackTreeSetSync<Value>,
12    charge: ShareCharge,
13}
14
15impl SortedSet {
16    pub fn empty() -> Self {
17        Self {
18            inner: rpds::RedBlackTreeSetSync::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 contains(&self, value: &Value) -> bool {
32        self.inner.contains(value)
33    }
34
35    /// Return a new set with `val` added.
36    pub fn conj(&self, value: Value) -> Self {
37        self.derived(self.inner.insert(value))
38    }
39
40    pub fn conj_mut(&mut self, value: Value) -> &mut Self {
41        self.inner.insert_mut(value);
42        self.charge = ShareCharge::full();
43        self
44    }
45
46    /// Return a new set with `val` removed.
47    pub fn disj(&self, value: &Value) -> Self {
48        self.derived(self.inner.remove(value))
49    }
50
51    /// Wrap `inner`, produced by one path-copying update of `self.inner`.
52    fn derived(&self, inner: rpds::RedBlackTreeSetSync<Value>) -> Self {
53        let n = inner.size().max(self.inner.size());
54        Self {
55            charge: self
56                .charge
57                .derive(self.inner.size(), inner.size(), rb_path_bytes::<Value>(n)),
58            inner,
59        }
60    }
61
62    /// Estimated bytes owned by this set if it shared no nodes.
63    fn full_size_extra(&self) -> usize {
64        self.inner.size() * (40 + std::mem::size_of::<Value>())
65    }
66
67    /// What `gc_size_extra` would report, without marking this version
68    /// allocated.
69    #[cfg(test)]
70    fn charged(&self) -> usize {
71        self.charge.size_extra(self.full_size_extra())
72    }
73
74    /// Iterate over all elements in sorted order
75    pub fn iter(&self) -> impl Iterator<Item = &Value> {
76        self.inner.iter()
77    }
78}
79
80impl FromIterator<Value> for SortedSet {
81    fn from_iter<I: IntoIterator<Item = Value>>(iter: I) -> Self {
82        let mut inner = rpds::RedBlackTreeSetSync::new_sync();
83        for v in iter {
84            inner.insert_mut(v);
85        }
86        SortedSet {
87            inner,
88            charge: ShareCharge::full(),
89        }
90    }
91}
92
93impl PartialEq for SortedSet {
94    fn eq(&self, other: &Self) -> bool {
95        if self.count() != other.count() {
96            return false;
97        }
98        self.inner.iter().all(|v| other.contains(v))
99    }
100}
101
102impl cljrs_gc::Trace for SortedSet {
103    #[cfg(feature = "no-gc")]
104    fn immutable_children(&self) -> bool {
105        true
106    }
107    fn trace(&self, visitor: &mut cljrs_gc::MarkVisitor) {
108        for v in self.inner.iter() {
109            v.trace(visitor);
110        }
111    }
112
113    fn gc_size_extra(&self) -> usize {
114        self.charge.on_alloc(self.full_size_extra())
115    }
116}
117
118#[cfg(test)]
119mod tests {
120    use super::*;
121    use crate::Value;
122
123    fn int(n: i64) -> Value {
124        Value::Long(n)
125    }
126
127    #[test]
128    fn test_basic() {
129        let s = SortedSet::empty();
130        let s = s.conj(int(1)).conj(int(2)).conj(int(3));
131        assert_eq!(s.count(), 3);
132        assert!(s.contains(&int(1)));
133        assert!(s.contains(&int(2)));
134        assert!(!s.contains(&int(99)));
135    }
136
137    #[test]
138    fn test_idempotent_conj() {
139        let s = SortedSet::empty().conj(int(1)).conj(int(1));
140        assert_eq!(s.count(), 1);
141    }
142
143    #[test]
144    fn test_disj() {
145        let s = SortedSet::empty().conj(int(1)).conj(int(2));
146        let s2 = s.disj(&int(1));
147        assert!(!s2.contains(&int(1)));
148        assert!(s2.contains(&int(2)));
149        assert_eq!(s2.count(), 1);
150    }
151
152    #[test]
153    fn test_equality_order_independent() {
154        let a = SortedSet::from_iter([int(1), int(2), int(3)]);
155        let b = SortedSet::from_iter([int(3), int(1), int(2)]);
156        assert_eq!(a, b);
157    }
158
159    #[test]
160    fn test_sorted_set_sorted() {
161        let mut a = SortedSet::empty();
162        a = a.conj(int(3));
163        a = a.conj(int(1));
164        a = a.conj(int(2));
165        let v: Vec<Value> = a.iter().cloned().collect();
166        assert!(matches!(&v[0], Value::Long(1)));
167        assert!(matches!(&v[1], Value::Long(2)));
168        assert!(matches!(&v[2], Value::Long(3)));
169    }
170
171    #[test]
172    fn conj_on_allocated_set_charges_path_bytes() {
173        let _alloc = crate::collections::share_charge::alloc_scope();
174        use cljrs_gc::GcPtr;
175
176        let s = GcPtr::new(SortedSet::from_iter((0..10_000).map(int)));
177        let added = s.get().conj(int(-1));
178        assert_eq!(added.charged(), rb_path_bytes::<Value>(10_001));
179        assert!(added.charged() * 100 < s.get().full_size_extra());
180    }
181}