surrealdb_core/idx/trees/
dynamicset.rs

1use crate::idx::trees::hnsw::ElementId;
2use ahash::{HashSet, HashSetExt};
3use std::fmt::Debug;
4
5pub trait DynamicSet: Debug + Send + Sync {
6	fn with_capacity(capacity: usize) -> Self;
7	fn insert(&mut self, v: ElementId) -> bool;
8	fn contains(&self, v: &ElementId) -> bool;
9	fn remove(&mut self, v: &ElementId) -> bool;
10	fn len(&self) -> usize;
11	fn is_empty(&self) -> bool;
12	fn iter(&self) -> impl Iterator<Item = &ElementId>;
13}
14
15#[derive(Debug)]
16pub struct AHashSet(HashSet<ElementId>);
17
18impl DynamicSet for AHashSet {
19	#[inline]
20	fn with_capacity(capacity: usize) -> Self {
21		Self(HashSet::with_capacity(capacity))
22	}
23
24	#[inline]
25	fn insert(&mut self, v: ElementId) -> bool {
26		self.0.insert(v)
27	}
28
29	#[inline]
30	fn contains(&self, v: &ElementId) -> bool {
31		self.0.contains(v)
32	}
33
34	#[inline]
35	fn remove(&mut self, v: &ElementId) -> bool {
36		self.0.remove(v)
37	}
38
39	#[inline]
40	fn len(&self) -> usize {
41		self.0.len()
42	}
43
44	#[inline]
45	fn is_empty(&self) -> bool {
46		self.0.is_empty()
47	}
48
49	#[inline]
50	fn iter(&self) -> impl Iterator<Item = &ElementId> {
51		self.0.iter()
52	}
53}
54
55#[derive(Debug)]
56pub struct ArraySet<const N: usize> {
57	array: [ElementId; N],
58	size: usize,
59}
60
61impl<const N: usize> DynamicSet for ArraySet<N> {
62	fn with_capacity(_capacity: usize) -> Self {
63		#[cfg(debug_assertions)]
64		assert!(_capacity <= N);
65		Self {
66			array: [0; N],
67			size: 0,
68		}
69	}
70
71	#[inline]
72	fn insert(&mut self, v: ElementId) -> bool {
73		if !self.contains(&v) {
74			self.array[self.size] = v;
75			self.size += 1;
76			true
77		} else {
78			false
79		}
80	}
81
82	#[inline]
83	fn contains(&self, v: &ElementId) -> bool {
84		self.array[0..self.size].contains(v)
85	}
86
87	#[inline]
88	fn remove(&mut self, v: &ElementId) -> bool {
89		if let Some(p) = self.array[0..self.size].iter().position(|e| e.eq(v)) {
90			self.array[p..].rotate_left(1);
91			self.size -= 1;
92			true
93		} else {
94			false
95		}
96	}
97
98	#[inline]
99	fn len(&self) -> usize {
100		self.size
101	}
102
103	#[inline]
104	fn is_empty(&self) -> bool {
105		self.size == 0
106	}
107
108	#[inline]
109	fn iter(&self) -> impl Iterator<Item = &ElementId> {
110		self.array[0..self.size].iter()
111	}
112}
113
114#[cfg(test)]
115mod tests {
116	use crate::idx::trees::dynamicset::{AHashSet, ArraySet, DynamicSet};
117	use crate::idx::trees::hnsw::ElementId;
118	use ahash::HashSet;
119
120	fn test_dynamic_set<S: DynamicSet>(capacity: ElementId) {
121		let mut dyn_set = S::with_capacity(capacity as usize);
122		let mut control = HashSet::default();
123		// Test insertions
124		for sample in 0..capacity {
125			assert_eq!(dyn_set.len(), control.len(), "{capacity} - {sample}");
126			let v: HashSet<ElementId> = dyn_set.iter().cloned().collect();
127			assert_eq!(v, control, "{capacity} - {sample}");
128			// We should not have the element yet
129			assert!(!dyn_set.contains(&sample), "{capacity} - {sample}");
130			// The first insertion returns true
131			assert!(dyn_set.insert(sample));
132			assert!(dyn_set.contains(&sample), "{capacity} - {sample}");
133			// The second insertion returns false
134			assert!(!dyn_set.insert(sample));
135			assert!(dyn_set.contains(&sample), "{capacity} - {sample}");
136			// We update the control structure
137			control.insert(sample);
138		}
139		// Test removals
140		for sample in 0..capacity {
141			// The first removal returns true
142			assert!(dyn_set.remove(&sample));
143			assert!(!dyn_set.contains(&sample), "{capacity} - {sample}");
144			// The second removal returns false
145			assert!(!dyn_set.remove(&sample));
146			assert!(!dyn_set.contains(&sample), "{capacity} - {sample}");
147			// We update the control structure
148			control.remove(&sample);
149			// The control structure and the dyn_set should be identical
150			assert_eq!(dyn_set.len(), control.len(), "{capacity} - {sample}");
151			let v: HashSet<ElementId> = dyn_set.iter().cloned().collect();
152			assert_eq!(v, control, "{capacity} - {sample}");
153		}
154	}
155
156	#[test]
157	fn test_dynamic_set_hash() {
158		for capacity in 1..50 {
159			test_dynamic_set::<AHashSet>(capacity);
160		}
161	}
162
163	#[test]
164	fn test_dynamic_set_array() {
165		test_dynamic_set::<ArraySet<1>>(1);
166		test_dynamic_set::<ArraySet<2>>(2);
167		test_dynamic_set::<ArraySet<4>>(4);
168		test_dynamic_set::<ArraySet<10>>(10);
169		test_dynamic_set::<ArraySet<20>>(20);
170		test_dynamic_set::<ArraySet<30>>(30);
171	}
172}