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