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
80#[cfg(test)]
81mod tests {
82    use super::*;
83    use crate::Value;
84
85    fn int(n: i64) -> Value {
86        Value::Long(n)
87    }
88
89    #[test]
90    fn test_basic() {
91        let s = SortedSet::empty();
92        let s = s.conj(int(1)).conj(int(2)).conj(int(3));
93        assert_eq!(s.count(), 3);
94        assert!(s.contains(&int(1)));
95        assert!(s.contains(&int(2)));
96        assert!(!s.contains(&int(99)));
97    }
98
99    #[test]
100    fn test_idempotent_conj() {
101        let s = SortedSet::empty().conj(int(1)).conj(int(1));
102        assert_eq!(s.count(), 1);
103    }
104
105    #[test]
106    fn test_disj() {
107        let s = SortedSet::empty().conj(int(1)).conj(int(2));
108        let s2 = s.disj(&int(1));
109        assert!(!s2.contains(&int(1)));
110        assert!(s2.contains(&int(2)));
111        assert_eq!(s2.count(), 1);
112    }
113
114    #[test]
115    fn test_equality_order_independent() {
116        let a = SortedSet::from_iter([int(1), int(2), int(3)]);
117        let b = SortedSet::from_iter([int(3), int(1), int(2)]);
118        assert_eq!(a, b);
119    }
120
121    #[test]
122    fn test_sorted_set_sorted() {
123        let mut a = SortedSet::empty();
124        a = a.conj(int(3));
125        a = a.conj(int(1));
126        a = a.conj(int(2));
127        let v: Vec<Value> = a.iter().cloned().collect();
128        assert!(matches!(&v[0], Value::Long(1)));
129        assert!(matches!(&v[1], Value::Long(2)));
130        assert!(matches!(&v[2], Value::Long(3)));
131    }
132}