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    fn trace(&self, visitor: &mut cljrs_gc::MarkVisitor) {
104        for v in self.inner.iter() {
105            v.trace(visitor);
106        }
107    }
108
109    fn gc_size_extra(&self) -> usize {
110        self.charge.on_alloc(self.full_size_extra())
111    }
112}
113
114#[cfg(test)]
115mod tests {
116    use super::*;
117    use crate::Value;
118
119    fn int(n: i64) -> Value {
120        Value::Long(n)
121    }
122
123    #[test]
124    fn test_basic() {
125        let s = SortedSet::empty();
126        let s = s.conj(int(1)).conj(int(2)).conj(int(3));
127        assert_eq!(s.count(), 3);
128        assert!(s.contains(&int(1)));
129        assert!(s.contains(&int(2)));
130        assert!(!s.contains(&int(99)));
131    }
132
133    #[test]
134    fn test_idempotent_conj() {
135        let s = SortedSet::empty().conj(int(1)).conj(int(1));
136        assert_eq!(s.count(), 1);
137    }
138
139    #[test]
140    fn test_disj() {
141        let s = SortedSet::empty().conj(int(1)).conj(int(2));
142        let s2 = s.disj(&int(1));
143        assert!(!s2.contains(&int(1)));
144        assert!(s2.contains(&int(2)));
145        assert_eq!(s2.count(), 1);
146    }
147
148    #[test]
149    fn test_equality_order_independent() {
150        let a = SortedSet::from_iter([int(1), int(2), int(3)]);
151        let b = SortedSet::from_iter([int(3), int(1), int(2)]);
152        assert_eq!(a, b);
153    }
154
155    #[test]
156    fn test_sorted_set_sorted() {
157        let mut a = SortedSet::empty();
158        a = a.conj(int(3));
159        a = a.conj(int(1));
160        a = a.conj(int(2));
161        let v: Vec<Value> = a.iter().cloned().collect();
162        assert!(matches!(&v[0], Value::Long(1)));
163        assert!(matches!(&v[1], Value::Long(2)));
164        assert!(matches!(&v[2], Value::Long(3)));
165    }
166
167    #[test]
168    fn conj_on_allocated_set_charges_path_bytes() {
169        let _alloc = crate::collections::share_charge::alloc_scope();
170        use cljrs_gc::GcPtr;
171
172        let s = GcPtr::new(SortedSet::from_iter((0..10_000).map(int)));
173        let added = s.get().conj(int(-1));
174        assert_eq!(added.charged(), rb_path_bytes::<Value>(10_001));
175        assert!(added.charged() * 100 < s.get().full_size_extra());
176    }
177}