1use crate::avl::{Iter, NodeHandle, NodeRef, 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 root(&self) -> Option<NodeRef<'_, K, (), SIZE>> {
291 self.0.root()
292 }
293
294 pub fn from_root(root: Option<NodeHandle<K, (), SIZE>>) -> Self {
296 Set(Tree::from_root(root))
297 }
298
299 pub fn downgrade(&self) -> WeakSetRef<K, SIZE> {
301 WeakSetRef(self.0.downgrade())
302 }
303
304 pub fn strong_count(&self) -> usize {
306 self.0.strong_count()
307 }
308
309 pub fn weak_count(&self) -> usize {
311 self.0.weak_count()
312 }
313
314 pub fn insert_many<E: IntoIterator<Item = K>>(&self, elts: E) -> Self {
332 let root = self.0.insert_many(elts.into_iter().map(|k| (k, ())));
333 Set(root)
334 }
335
336 pub fn remove_many<Q, E>(&self, elts: E) -> Self
339 where
340 Q: Ord,
341 K: Borrow<Q>,
342 E: IntoIterator<Item = Q>,
343 {
344 let root = self
345 .0
346 .update_many(elts.into_iter().map(|k| (k, ())), &mut |_, _, _| None);
347 Set(root)
348 }
349
350 pub fn update_many<Q, E, F>(&self, elts: E, mut f: F) -> Self
355 where
356 Q: Ord,
357 K: Borrow<Q>,
358 E: IntoIterator<Item = Q>,
359 F: FnMut(Q, Option<&K>) -> Option<K>,
360 {
361 let root =
362 self.0
363 .update_many(elts.into_iter().map(|k| (k, ())), &mut |q, (), cur| {
364 let cur = cur.map(|(k, ())| k);
365 f(q, cur).map(|k| (k, ()))
366 });
367 Set(root)
368 }
369
370 pub fn insert(&self, k: K) -> (Self, bool) {
375 if self.contains(&k) {
376 (self.clone(), true)
377 } else {
378 (Set(self.0.insert(k, ()).0), false)
379 }
380 }
381
382 pub fn insert_cow(&mut self, k: K) -> bool {
389 self.0.insert_cow(k, ()).is_some()
390 }
391
392 pub fn contains<'a, Q>(&'a self, k: &Q) -> bool
396 where
397 Q: ?Sized + Ord,
398 K: Borrow<Q>,
399 {
400 self.0.get(k).is_some()
401 }
402
403 pub fn get<'a, Q>(&'a self, k: &Q) -> Option<&'a K>
406 where
407 Q: ?Sized + Ord,
408 K: Borrow<Q>,
409 {
410 self.0.get_key(k)
411 }
412
413 pub fn remove<Q: Sized + Ord>(&self, k: &Q) -> (Self, bool)
416 where
417 K: Borrow<Q>,
418 {
419 let (t, prev) = self.0.remove(k);
420 (Set(t), prev.is_some())
421 }
422
423 pub fn remove_cow<Q: Sized + Ord>(&mut self, k: &Q) -> bool
426 where
427 K: Borrow<Q>,
428 {
429 self.0.remove_cow(k).is_some()
430 }
431
432 pub fn union(&self, other: &Set<K, SIZE>) -> Self {
450 Set(Tree::union(&self.0, &other.0, &mut |_, (), ()| Some(())))
451 }
452
453 pub fn intersect(&self, other: &Set<K, SIZE>) -> Self {
474 Set(Tree::intersect(
475 &self.0,
476 &other.0,
477 &mut |_, (), ()| Some(()),
478 ))
479 }
480
481 pub fn diff(&self, other: &Set<K, SIZE>) -> Self
503 where
504 K: Debug,
505 {
506 Set(Tree::diff(&self.0, &other.0, &mut |_, (), ()| None))
507 }
508
509 pub fn len(&self) -> usize {
511 self.0.len()
512 }
513
514 pub fn range<'a, Q, R>(&'a self, r: R) -> SetIter<'a, R, Q, K, SIZE>
523 where
524 Q: Ord + ?Sized + 'a,
525 K: 'a + Clone + Ord + Borrow<Q>,
526 R: RangeBounds<Q> + 'a,
527 {
528 SetIter(self.0.range(r))
529 }
530}
531
532impl<K, const SIZE: usize> Set<K, SIZE>
533where
534 K: Ord + Clone + Debug,
535{
536 #[allow(dead_code)]
537 pub(crate) fn invariant(&self) -> () {
538 self.0.invariant()
539 }
540}