1use crate::avl::{Iter, Tree, WeakTree};
2pub use crate::chunk::DEFAULT_SIZE;
3use core::{
4 borrow::Borrow,
5 cmp::{Eq, Ord, Ordering, PartialEq, PartialOrd},
6 default::Default,
7 fmt::{self, Debug, Formatter},
8 hash::{Hash, Hasher},
9 iter::FromIterator,
10 ops::{RangeBounds, RangeFull},
11};
12
13#[cfg(feature = "serde")]
14use serde::{
15 de::{SeqAccess, Visitor},
16 ser::SerializeSeq,
17 Deserialize, Deserializer, Serialize, Serializer,
18};
19
20#[cfg(feature = "serde")]
21use core::marker::PhantomData;
22
23#[cfg(feature = "rayon")]
24use rayon::{
25 iter::{FromParallelIterator, IntoParallelIterator},
26 prelude::*,
27};
28
29#[derive(Clone)]
52#[repr(transparent)]
53pub struct Set<K: Ord + Clone, const SIZE: usize>(Tree<K, (), SIZE>);
54
55pub type SetS<K> = Set<K, { DEFAULT_SIZE / 2 }>;
57
58pub type SetM<K> = Set<K, DEFAULT_SIZE>;
60
61pub type SetL<K> = Set<K, { DEFAULT_SIZE * 2 }>;
63
64#[derive(Clone)]
65pub struct WeakSetRef<K: Ord + Clone, const SIZE: usize>(WeakTree<K, (), SIZE>);
66
67pub type WeakSetRefS<K> = WeakSetRef<K, 32>;
68pub type WeakSetRefM<K> = WeakSetRef<K, 128>;
69pub type WeakSetRefL<K> = WeakSetRef<K, 512>;
70
71impl<K, const SIZE: usize> WeakSetRef<K, SIZE>
72where
73 K: Ord + Clone,
74{
75 pub fn upgrade(&self) -> Option<Set<K, SIZE>> {
76 self.0.upgrade().map(Set)
77 }
78}
79
80impl<K, const SIZE: usize> Hash for Set<K, SIZE>
81where
82 K: Hash + Ord + Clone,
83{
84 fn hash<H: Hasher>(&self, state: &mut H) {
85 self.0.hash(state)
86 }
87}
88
89impl<K, const SIZE: usize> Default for Set<K, SIZE>
90where
91 K: Ord + Clone,
92{
93 fn default() -> Set<K, SIZE> {
94 Set::new()
95 }
96}
97
98impl<K, const SIZE: usize> PartialEq for Set<K, SIZE>
99where
100 K: Ord + Clone,
101{
102 fn eq(&self, other: &Set<K, SIZE>) -> bool {
103 self.0 == other.0
104 }
105}
106
107impl<K, const SIZE: usize> Eq for Set<K, SIZE> where K: Eq + Ord + Clone {}
108
109impl<K, const SIZE: usize> PartialOrd for Set<K, SIZE>
110where
111 K: Ord + Clone,
112{
113 fn partial_cmp(&self, other: &Set<K, SIZE>) -> Option<Ordering> {
114 self.0.partial_cmp(&other.0)
115 }
116}
117
118impl<K, const SIZE: usize> Ord for Set<K, SIZE>
119where
120 K: Ord + Clone,
121{
122 fn cmp(&self, other: &Set<K, SIZE>) -> Ordering {
123 self.0.cmp(&other.0)
124 }
125}
126
127impl<K, const SIZE: usize> Debug for Set<K, SIZE>
128where
129 K: Debug + Ord + Clone,
130{
131 fn fmt(&self, f: &mut Formatter) -> fmt::Result {
132 f.debug_set().entries(self.into_iter()).finish()
133 }
134}
135
136impl<K, const SIZE: usize> FromIterator<K> for Set<K, SIZE>
137where
138 K: Ord + Clone,
139{
140 fn from_iter<T: IntoIterator<Item = K>>(iter: T) -> Self {
141 Set::new().insert_many(iter)
142 }
143}
144
145pub struct SetIter<
146 'a,
147 R: RangeBounds<Q> + 'a,
148 Q: Ord + ?Sized,
149 K: 'a + Clone + Ord + Borrow<Q>,
150 const SIZE: usize,
151>(Iter<'a, R, Q, K, (), SIZE>);
152
153impl<'a, R, Q, K, const SIZE: usize> Iterator for SetIter<'a, R, Q, K, SIZE>
154where
155 Q: Ord + ?Sized,
156 R: RangeBounds<Q> + 'a,
157 K: 'a + Clone + Ord + Borrow<Q>,
158{
159 type Item = &'a K;
160 fn next(&mut self) -> Option<Self::Item> {
161 self.0.next().map(|(k, ())| k)
162 }
163}
164
165impl<'a, R, Q, K, const SIZE: usize> DoubleEndedIterator for SetIter<'a, R, Q, K, SIZE>
166where
167 Q: Ord + ?Sized,
168 R: RangeBounds<Q> + 'a,
169 K: 'a + Clone + Ord + Borrow<Q>,
170{
171 fn next_back(&mut self) -> Option<Self::Item> {
172 self.0.next_back().map(|(k, ())| k)
173 }
174}
175
176impl<'a, K, const SIZE: usize> IntoIterator for &'a Set<K, SIZE>
177where
178 K: 'a + Ord + Clone,
179{
180 type Item = &'a K;
181 type IntoIter = SetIter<'a, RangeFull, K, K, SIZE>;
182 fn into_iter(self) -> Self::IntoIter {
183 SetIter(self.0.into_iter())
184 }
185}
186
187#[cfg(feature = "serde")]
188impl<V, const SIZE: usize> Serialize for Set<V, SIZE>
189where
190 V: Serialize + Clone + Ord,
191{
192 fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
193 where
194 S: Serializer,
195 {
196 let mut seq = serializer.serialize_seq(Some(self.len()))?;
197 for v in self {
198 seq.serialize_element(v)?
199 }
200 seq.end()
201 }
202}
203
204#[cfg(feature = "serde")]
205struct SetVisitor<V: Clone + Ord, const SIZE: usize> {
206 marker: PhantomData<fn() -> Set<V, SIZE>>,
207}
208
209#[cfg(feature = "serde")]
210impl<'a, V, const SIZE: usize> Visitor<'a> for SetVisitor<V, SIZE>
211where
212 V: Deserialize<'a> + Clone + Ord,
213{
214 type Value = Set<V, SIZE>;
215
216 fn expecting(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
217 formatter.write_str("expecting an immutable_chunkmap::Set")
218 }
219
220 fn visit_seq<A>(self, mut seq: A) -> Result<Self::Value, A::Error>
221 where
222 A: SeqAccess<'a>,
223 {
224 let mut t = Set::<V, SIZE>::new();
225 while let Some(v) = seq.next_element()? {
226 t.insert_cow(v);
227 }
228 Ok(t)
229 }
230}
231
232#[cfg(feature = "serde")]
233impl<'a, V, const SIZE: usize> Deserialize<'a> for Set<V, SIZE>
234where
235 V: Deserialize<'a> + Clone + Ord,
236{
237 fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
238 where
239 D: Deserializer<'a>,
240 {
241 deserializer.deserialize_seq(SetVisitor {
242 marker: PhantomData,
243 })
244 }
245}
246
247#[cfg(feature = "rayon")]
248impl<'a, V, const SIZE: usize> IntoParallelIterator for &'a Set<V, SIZE>
249where
250 V: 'a + Ord + Clone + Send + Sync,
251{
252 type Item = &'a V;
253 type Iter = rayon::vec::IntoIter<&'a V>;
254
255 fn into_par_iter(self) -> Self::Iter {
256 self.into_iter().collect::<Vec<_>>().into_par_iter()
257 }
258}
259
260#[cfg(feature = "rayon")]
261impl<V, const SIZE: usize> FromParallelIterator<V> for Set<V, SIZE>
262where
263 V: Ord + Clone + Send + Sync,
264{
265 fn from_par_iter<I>(i: I) -> Self
266 where
267 I: IntoParallelIterator<Item = V>,
268 {
269 i.into_par_iter()
270 .fold_with(Set::new(), |mut m, v| {
271 m.insert_cow(v);
272 m
273 })
274 .reduce_with(|m0, m1| m0.union(&m1))
275 .unwrap_or_else(Set::new)
276 }
277}
278
279impl<K, const SIZE: usize> Set<K, SIZE>
280where
281 K: Ord + Clone,
282{
283 pub fn new() -> Self {
285 Set(Tree::new())
286 }
287
288 pub fn downgrade(&self) -> WeakSetRef<K, SIZE> {
290 WeakSetRef(self.0.downgrade())
291 }
292
293 pub fn strong_count(&self) -> usize {
295 self.0.strong_count()
296 }
297
298 pub fn weak_count(&self) -> usize {
300 self.0.weak_count()
301 }
302
303 pub fn insert_many<E: IntoIterator<Item = K>>(&self, elts: E) -> Self {
321 let root = self.0.insert_many(elts.into_iter().map(|k| (k, ())));
322 Set(root)
323 }
324
325 pub fn remove_many<Q, E>(&self, elts: E) -> Self
328 where
329 Q: Ord,
330 K: Borrow<Q>,
331 E: IntoIterator<Item = Q>,
332 {
333 let root = self
334 .0
335 .update_many(elts.into_iter().map(|k| (k, ())), &mut |_, _, _| None);
336 Set(root)
337 }
338
339 pub fn update_many<Q, E, F>(&self, elts: E, mut f: F) -> Self
344 where
345 Q: Ord,
346 K: Borrow<Q>,
347 E: IntoIterator<Item = Q>,
348 F: FnMut(Q, Option<&K>) -> Option<K>,
349 {
350 let root =
351 self.0
352 .update_many(elts.into_iter().map(|k| (k, ())), &mut |q, (), cur| {
353 let cur = cur.map(|(k, ())| k);
354 f(q, cur).map(|k| (k, ()))
355 });
356 Set(root)
357 }
358
359 pub fn insert(&self, k: K) -> (Self, bool) {
364 if self.contains(&k) {
365 (self.clone(), true)
366 } else {
367 (Set(self.0.insert(k, ()).0), false)
368 }
369 }
370
371 pub fn insert_cow(&mut self, k: K) -> bool {
378 self.0.insert_cow(k, ()).is_some()
379 }
380
381 pub fn contains<'a, Q>(&'a self, k: &Q) -> bool
385 where
386 Q: ?Sized + Ord,
387 K: Borrow<Q>,
388 {
389 self.0.get(k).is_some()
390 }
391
392 pub fn get<'a, Q>(&'a self, k: &Q) -> Option<&'a K>
395 where
396 Q: ?Sized + Ord,
397 K: Borrow<Q>,
398 {
399 self.0.get_key(k)
400 }
401
402 pub fn remove<Q: Sized + Ord>(&self, k: &Q) -> (Self, bool)
405 where
406 K: Borrow<Q>,
407 {
408 let (t, prev) = self.0.remove(k);
409 (Set(t), prev.is_some())
410 }
411
412 pub fn remove_cow<Q: Sized + Ord>(&mut self, k: &Q) -> bool
415 where
416 K: Borrow<Q>,
417 {
418 self.0.remove_cow(k).is_some()
419 }
420
421 pub fn union(&self, other: &Set<K, SIZE>) -> Self {
439 Set(Tree::union(&self.0, &other.0, &mut |_, (), ()| Some(())))
440 }
441
442 pub fn intersect(&self, other: &Set<K, SIZE>) -> Self {
463 Set(Tree::intersect(
464 &self.0,
465 &other.0,
466 &mut |_, (), ()| Some(()),
467 ))
468 }
469
470 pub fn diff(&self, other: &Set<K, SIZE>) -> Self
492 where
493 K: Debug,
494 {
495 Set(Tree::diff(&self.0, &other.0, &mut |_, (), ()| None))
496 }
497
498 pub fn len(&self) -> usize {
500 self.0.len()
501 }
502
503 pub fn range<'a, Q, R>(&'a self, r: R) -> SetIter<'a, R, Q, K, SIZE>
512 where
513 Q: Ord + ?Sized + 'a,
514 K: 'a + Clone + Ord + Borrow<Q>,
515 R: RangeBounds<Q> + 'a,
516 {
517 SetIter(self.0.range(r))
518 }
519}
520
521impl<K, const SIZE: usize> Set<K, SIZE>
522where
523 K: Ord + Clone + Debug,
524{
525 #[allow(dead_code)]
526 pub(crate) fn invariant(&self) -> () {
527 self.0.invariant()
528 }
529}