Skip to main content

cljrs_value/collections/
sorted_set.rs

1use crate::Value;
2
3/// An immutable sorted set backed by `rpds::RedBlackTreeSet`.
4#[derive(Debug, Clone)]
5pub struct SortedSet {
6    inner: rpds::RedBlackTreeSetSync<Value>,
7}
8
9impl SortedSet {
10    pub fn empty() -> Self {
11        Self {
12            inner: rpds::RedBlackTreeSetSync::new_sync(),
13        }
14    }
15
16    pub fn count(&self) -> usize {
17        self.inner.size()
18    }
19
20    pub fn is_empty(&self) -> bool {
21        self.inner.is_empty()
22    }
23
24    pub fn contains(&self, value: &Value) -> bool {
25        self.inner.contains(value)
26    }
27
28    /// Return a new set with `val` added.
29    pub fn conj(&self, value: Value) -> Self {
30        Self {
31            inner: self.inner.insert(value),
32        }
33    }
34
35    pub fn conj_mut(&mut self, value: Value) -> &mut Self {
36        self.inner.insert_mut(value);
37        self
38    }
39
40    /// Return a new set with `val` removed.
41    pub fn disj(&self, value: &Value) -> Self {
42        Self {
43            inner: self.inner.remove(value),
44        }
45    }
46
47    /// Iterate over all elements in sorted order
48    pub fn iter(&self) -> impl Iterator<Item = &Value> {
49        self.inner.iter()
50    }
51}
52
53impl FromIterator<Value> for SortedSet {
54    fn from_iter<I: IntoIterator<Item = Value>>(iter: I) -> Self {
55        let mut inner = rpds::RedBlackTreeSetSync::new_sync();
56        for v in iter {
57            inner.insert_mut(v);
58        }
59        SortedSet { inner }
60    }
61}
62
63impl PartialEq for SortedSet {
64    fn eq(&self, other: &Self) -> bool {
65        if self.count() != other.count() {
66            return false;
67        }
68        self.inner.iter().all(|v| other.contains(v))
69    }
70}
71
72impl cljrs_gc::Trace for SortedSet {
73    fn trace(&self, visitor: &mut cljrs_gc::MarkVisitor) {
74        for v in self.inner.iter() {
75            v.trace(visitor);
76        }
77    }
78
79    fn gc_size_extra(&self) -> usize {
80        let n = self.inner.size();
81        n * (40 + std::mem::size_of::<Value>())
82    }
83}
84
85#[cfg(test)]
86mod tests {
87    use super::*;
88    use crate::Value;
89
90    fn int(n: i64) -> Value {
91        Value::Long(n)
92    }
93
94    #[test]
95    fn test_basic() {
96        let s = SortedSet::empty();
97        let s = s.conj(int(1)).conj(int(2)).conj(int(3));
98        assert_eq!(s.count(), 3);
99        assert!(s.contains(&int(1)));
100        assert!(s.contains(&int(2)));
101        assert!(!s.contains(&int(99)));
102    }
103
104    #[test]
105    fn test_idempotent_conj() {
106        let s = SortedSet::empty().conj(int(1)).conj(int(1));
107        assert_eq!(s.count(), 1);
108    }
109
110    #[test]
111    fn test_disj() {
112        let s = SortedSet::empty().conj(int(1)).conj(int(2));
113        let s2 = s.disj(&int(1));
114        assert!(!s2.contains(&int(1)));
115        assert!(s2.contains(&int(2)));
116        assert_eq!(s2.count(), 1);
117    }
118
119    #[test]
120    fn test_equality_order_independent() {
121        let a = SortedSet::from_iter([int(1), int(2), int(3)]);
122        let b = SortedSet::from_iter([int(3), int(1), int(2)]);
123        assert_eq!(a, b);
124    }
125
126    #[test]
127    fn test_sorted_set_sorted() {
128        let mut a = SortedSet::empty();
129        a = a.conj(int(3));
130        a = a.conj(int(1));
131        a = a.conj(int(2));
132        let v: Vec<Value> = a.iter().cloned().collect();
133        assert!(matches!(&v[0], Value::Long(1)));
134        assert!(matches!(&v[1], Value::Long(2)));
135        assert!(matches!(&v[2], Value::Long(3)));
136    }
137}