hitree/hiset.rs
1//use std::fmt::{Debug,Display,Formatter};
2use std::borrow::{Borrow, BorrowMut};
3use std::cmp::Ordering;
4use super::tree_height;
5
6/// Ordered set of values, accessible by value or index of value in the set.
7/// Stores values in a balanced binary tree with subtree node count tracking.
8/// Nodes are allocated on the heap using `Box`.
9pub struct HiSet<T: Ord> {
10 root: Ref<T>,
11}
12
13/// Reference to a subtree of Nodes, including node count of subtree pointed to by it.
14struct Ref<T>
15 where T: Ord
16{
17 count: usize,
18 node: Option<Box<Node<T>>>,
19}
20
21/// Node holding a value and references to the left (lesser) and right (greater) subtrees.
22/// Left and right subtrees are always balanced - they may differ by at most one level of depth,
23/// and all the inner nodes of the tree (all levels except the one furthest from the root)
24/// must contain both left and right subtrees that are also balanced.
25struct Node<T>
26 where T: Ord
27{
28 value: T,
29 left: Ref<T>,
30 right: Ref<T>,
31}
32
33
34
35impl <T> HiSet<T>
36 where T: Ord
37{
38 /// Create new empty HiSet.
39 ///
40 /// Does not allocate anything.
41 ///
42 /// # Examples:
43 ///
44 /// ```
45 /// # #[allow(unused_mut)]
46 /// # use hitree::hiset::HiSet;
47 /// let mut set = HiSet::<String>::new();
48 /// ```
49 pub fn new() -> HiSet<T> {
50 HiSet { root: Ref::default() }
51 }
52
53
54
55
56
57
58
59
60 /// Return current number of entries in the set.
61 ///
62 /// Extremely cheap.
63 ///
64 /// # Examples:
65 ///
66 /// ```
67 /// # use hitree::hiset::HiSet;
68 /// let hiset = HiSet::<i32>::new();
69 /// assert_eq!(hiset.len(), 0);
70 /// ```
71 pub fn len(&self) -> usize {
72 self.root.count
73 }
74
75
76 /// Insert a new value into the set.
77 /// If the value was not in the set, return true.
78 /// If the value was already in the set, return false and don't touch the old value.
79 /// Value can be any type that can be converted into the value type using Into trait.
80 ///
81 /// # Examples:
82 ///
83 /// ```
84 /// # use hitree::hiset::HiSet;
85 /// let mut hiset = HiSet::<i32>::new();
86 /// assert_eq!(hiset.insert(1), true);
87 /// assert_eq!(hiset.insert(2), true);
88 /// assert_eq!(hiset.insert(1), false);
89 /// assert_eq!(hiset.len(), 2);
90 /// ```
91 /// You can insert &str into `HiSet<String>` for example:
92 /// ```
93 /// # use hitree::hiset::HiSet;
94 /// let mut hiset = HiSet::<String>::new(); // This is a set of Strings
95 /// assert_eq!(hiset.insert("This can be converted to a String"), true);
96 /// ```
97 pub fn insert(&mut self, value: impl Into<T>) -> bool {
98 self.root.insert(Node::new(value))
99 }
100
101
102 /// Get a shared borrow of value from set by index.
103 /// Values in the set are sorted according to their Ord trait,
104 /// index 0 is the smallest value.
105 /// Borrowed value can be any shared reference type that can be borrowed from T.
106 /// You can use it to borrow `&str` from `HiSet<String>` for example.
107 ///
108 /// # Examples:
109 ///
110 /// ```
111 /// # use hitree::hiset::HiSet;
112 /// let mut hiset = HiSet::<String>::new();
113 /// hiset.insert("This");
114 /// hiset.insert("is");
115 /// hiset.insert("a");
116 /// hiset.insert("test!");
117 ///
118 /// // you can ask for Option<&str>
119 /// assert_eq!(hiset.get_by_index::<str>(0), Some("This"));
120 /// // or Option<&String> if you want, whatever you can borrow from the T type.
121 /// assert_eq!(hiset.get_by_index::<String>(1), Some(&"a".to_string()));
122 /// assert_eq!(hiset.get_by_index::<str>(2), Some("is"));
123 /// assert_eq!(hiset.get_by_index::<str>(3), Some("test!"));
124 /// assert_eq!(hiset.get_by_index::<str>(4), None );
125 /// ```
126 ///
127 pub fn get_by_index<B>(&self, index: usize) -> Option<&B>
128 where T: Borrow<B>,
129 B: ?Sized
130 {
131 let mut index_to_find = index;
132 let mut current_node = self.root.node();
133 loop {
134 match current_node {
135 None => return None,
136 Some(node) => {
137 match node.left.count.cmp(&index_to_find) {
138 Ordering::Greater => {
139 // index must be in the left subtree
140 current_node = node.left.node();
141 },
142 Ordering::Equal => {
143 // found it, its this node
144 return Some(node.borrow_value())
145 },
146 Ordering::Less => {
147 // index must be in the right subtree
148 index_to_find = index_to_find - 1 - node.left.count;
149 current_node = node.right.node();
150 }
151 }
152 }
153 }
154 }
155 }
156
157 /// Get a mutable borrow of value from set by index.
158 /// Values in the set are sorted according to their Ord trait,
159 /// index 0 is the smallest value.
160 /// Borrowed value can be any mutable reference type that can be borrowed from T.
161 /// WARNING: You must never change the borrowed value in a way that would affect its ordering according to
162 /// its Ord trait implementation!
163 ///
164 /// # Examples:
165 ///
166 /// ```
167 /// # use std::cmp::Ordering;
168 /// # use hitree::hiset::HiSet;
169 ///
170 /// struct TestValue {
171 /// ordering: String,
172 /// data: usize,
173 /// }
174 ///
175 /// impl TestValue {
176 /// pub fn new(ordering: impl Into<String>) -> Self { TestValue { ordering: ordering.into(), data: 0 }}
177 /// pub fn touch(&mut self) { self.data += 1; }
178 /// }
179 /// impl PartialEq for TestValue { fn eq(&self, other: &Self) -> bool { self.ordering.eq(&other.ordering) } }
180 /// impl Eq for TestValue {}
181 /// impl PartialOrd for TestValue { fn partial_cmp(&self, other: &Self) -> Option<Ordering> { self.ordering.partial_cmp(&other.ordering) } }
182 /// impl Ord for TestValue { fn cmp(&self, other: &Self) -> Ordering { self.ordering.cmp(&other.ordering) } }
183 ///
184 ///
185 /// let mut hiset = HiSet::<TestValue>::new();
186 /// hiset.insert(TestValue::new("first"));
187 /// hiset.insert(TestValue::new("second"));
188 /// hiset.insert(TestValue::new("third"));
189 ///
190 /// hiset.get_by_index_mut(0).map(|value| value.touch() );
191 /// hiset.get_by_index_mut(2).map(|value| { value.touch(); value.touch();} );
192 ///
193 /// assert_eq!(hiset.get_by_index(0).unwrap().data, 1);
194 /// assert_eq!(hiset.get_by_index(1).unwrap().data, 0);
195 /// assert_eq!(hiset.get_by_index(2).unwrap().data, 2);
196 /// ```
197 ///
198 pub fn get_by_index_mut<B>(&mut self, index: usize) -> Option<&mut B>
199 where T: BorrowMut<B>,
200 B: ?Sized
201 {
202 let mut index_to_find = index;
203 let mut current_node = self.root.node_mut();
204 loop {
205 match current_node {
206 None => return None,
207 Some(node) => {
208 match node.left.count.cmp(&index_to_find) {
209 Ordering::Greater => {
210 // index must be in the left subtree
211 current_node = node.left.node_mut();
212 },
213 Ordering::Equal => {
214 // found it, its this node
215 return Some(node.borrow_value_mut())
216 },
217 Ordering::Less => {
218 // index must be in the right subtree
219 index_to_find = index_to_find - 1 - node.left.count;
220 current_node = node.right.node_mut();
221 }
222 }
223 }
224 }
225 }
226 }
227
228
229 /// Borrow value from set by key reference.
230 /// Reference type of key must have the same `Ord` ordering as `&T`.
231 ///
232 /// # Examples:
233 /// ```
234 /// # use hitree::hiset::HiSet;
235 /// let mut set = HiSet::<String>::new();
236 /// set.insert("This");
237 /// set.insert("is");
238 /// set.insert("a");
239 /// set.insert("test!");
240 ///
241 /// assert_eq!(set.get("test!"), Some(&"test!".to_string()));
242 /// assert_eq!(set.get("not there"), None);
243 /// assert_eq!(set.get(&"This".to_string()), Some(&"This".to_string()));
244 /// ```
245 pub fn get<KEY>(&mut self, key: &KEY) -> Option<&T>
246 where KEY: ?Sized + Ord, T: Borrow<KEY>
247 {
248 let mut current_node = self.root.node();
249 loop {
250 match current_node {
251 None => return None,
252 Some(node) => {
253 match Ord::cmp(node.value.borrow(), key) {
254 Ordering::Greater => {
255 // index must be in the left subtree
256 current_node = node.left.node();
257 },
258 Ordering::Equal => {
259 // found it, its this node
260 return Some(node.borrow_value::<T>())
261 },
262 Ordering::Less => {
263 // index must be in the right subtree
264 current_node = node.right.node();
265 }
266 }
267 }
268 }
269 }
270 }
271
272 /// Borrow mutably value from set by key reference.
273 /// Reference type of key must have the same `Ord` ordering as `&T`.
274 ///
275 /// # Examples:
276 /// ```
277 /// # use hitree::hiset::HiSet;
278 /// let mut set = HiSet::<String>::new();
279 /// set.insert("This");
280 /// set.insert("is");
281 /// set.insert("a");
282 /// set.insert("test!");
283 ///
284 /// assert_eq!(set.get_mut("test!"), Some(&mut "test!".to_string()));
285 /// assert_eq!(set.get_mut("not there"), None);
286 /// assert_eq!(set.get_mut(&"This".to_string()), Some(&mut "This".to_string()));
287 ///```
288 pub fn get_mut<KEY>(&mut self, key: &KEY) -> Option<&mut T>
289 where KEY: ?Sized + Ord, T: Borrow<KEY>
290 {
291 let mut current_node = self.root.node_mut();
292 loop {
293 match current_node {
294 None => return None,
295 Some(node) => {
296 match Ord::cmp(node.value.borrow(), key) {
297 Ordering::Greater => {
298 // index must be in the left subtree
299 current_node = node.left.node_mut();
300 },
301 Ordering::Equal => {
302 // found it, its this node
303 return Some(node.borrow_value_mut::<T>())
304 },
305 Ordering::Less => {
306 // index must be in the right subtree
307 current_node = node.right.node_mut();
308 }
309 }
310 }
311 }
312 }
313 }
314
315
316
317 /// Find index of value given by key reference.
318 ///
319 /// # Examples:
320 /// ```
321 /// # use hitree::hiset::HiSet;
322 /// let mut set = HiSet::<String>::new();
323 /// set.insert("This");
324 /// set.insert("is");
325 /// set.insert("a");
326 /// set.insert("test!");
327 ///
328 /// assert_eq!(set.index_of("This"), Some(0));
329 /// assert_eq!(set.index_of("a"), Some(1));
330 /// assert_eq!(set.index_of("is"), Some(2));
331 /// assert_eq!(set.index_of("test!"), Some(3));
332 /// assert_eq!(set.index_of("nonexistent"), None);
333 ///
334 /// ```
335 pub fn index_of<KEY>(&mut self, key: &KEY) -> Option<usize>
336 where KEY: ?Sized + Ord, T: Borrow<KEY>
337 {
338 let mut current_node = self.root.node();
339 let mut current_index_shift = 0;
340 loop {
341 match current_node {
342 None => return None,
343 Some(node) => {
344 match Ord::cmp(node.value.borrow(), key) {
345 Ordering::Greater => {
346 // index must be in the left subtree
347 current_node = node.left.node();
348 },
349 Ordering::Equal => {
350 // found it, its this node
351 return Some(current_index_shift + node.left.count)
352 },
353 Ordering::Less => {
354 // index must be in the right subtree
355 current_node = node.right.node();
356 current_index_shift += 1 + node.left.count;
357 }
358 }
359 }
360 }
361 }
362 }
363
364
365
366 /// Remove the smallest value from the set and return it.
367 ///
368 /// Examples:
369 ///
370 /// ```
371 /// # use hitree::hiset::HiSet;
372 /// let mut set = HiSet::<i32>::new();
373 /// set.insert(10);
374 /// set.insert(15);
375 /// set.insert(5);
376 ///
377 /// assert_eq!(set.len(), 3);
378 /// assert_eq!(set.take_first(), Some(5));
379 /// assert_eq!(set.len(), 2);
380 /// assert_eq!(set.take_first(), Some(10));
381 /// assert_eq!(set.len(), 1);
382 /// assert_eq!(set.take_first(), Some(15));
383 /// assert_eq!(set.len(), 0);
384 /// assert_eq!(set.take_first(), None);
385 /// ```
386 ///
387 pub fn take_first(&mut self) -> Option<T> {
388 self.root.take_leftmost_node().map(|node| node.value )
389 }
390
391 /// Remove the largest value from the set and return it.
392 ///
393 /// Examples:
394 ///
395 /// ```
396 /// # use hitree::hiset::HiSet;
397 /// let mut set = HiSet::<i32>::new();
398 /// set.insert(10);
399 /// set.insert(15);
400 /// set.insert(5);
401 ///
402 /// assert_eq!(set.len(), 3);
403 /// assert_eq!(set.take_last(), Some(15));
404 /// assert_eq!(set.len(), 2);
405 /// assert_eq!(set.take_last(), Some(10));
406 /// assert_eq!(set.len(), 1);
407 /// assert_eq!(set.take_last(), Some(5));
408 /// assert_eq!(set.len(), 0);
409 /// assert_eq!(set.take_last(), None);
410 /// ```
411 ///
412 pub fn take_last(&mut self) -> Option<T> {
413 self.root.take_rightmost_node().map(|node| node.value )
414 }
415
416
417 /// Take an entry by reference to another value and return it.
418 /// Whatever you use as key must give the same Ord results as Ord on &T!
419 ///
420 /// # Examples:
421 ///
422 /// ```
423 /// # use hitree::hiset::HiSet;
424 /// let mut set = HiSet::<String>::new();
425 /// assert_eq!(set.insert("first"), true);
426 /// assert_eq!(set.insert("second"), true);
427 /// assert_eq!(set.insert("third"), true);
428 /// assert_eq!(set.len(), 3);
429 ///
430 /// assert_eq!(set.take("second").unwrap().as_str(), "second");
431 /// assert_eq!(set.len(), 2);
432 /// assert_eq!(set.take("second"), None);
433 ///
434 /// assert_eq!(set.take(&"third".to_string()).unwrap().as_str(), "third");
435 /// assert_eq!(set.len(), 1);
436 ///
437 /// ```
438 pub fn take<KEY>(&mut self, key: &KEY) -> Option<T>
439 where KEY: ?Sized + Ord, T: Borrow<KEY>
440 {
441 let key = key.borrow();
442 self.root.take_node_by_key(key).map(|node| node.value )
443 }
444
445 /// Take an entry by reference to another value and return it.
446 /// Whatever you use as key must give the same Ord results as Ord on &T!
447 ///
448 /// # Examples:
449 ///
450 /// ```
451 /// # use hitree::hiset::HiSet;
452 /// let mut set = HiSet::<String>::new();
453 /// assert_eq!(set.insert("first"), true);
454 /// assert_eq!(set.insert("second"), true);
455 /// assert_eq!(set.insert("third"), true);
456 /// assert_eq!(set.len(), 3);
457 ///
458 /// assert_eq!(set.take_by_index(2).unwrap().as_str(), "third");
459 /// assert_eq!(set.len(), 2);
460 ///
461 /// assert_eq!(set.take_by_index(3), None);
462 ///
463 /// assert_eq!(set.take_by_index(1).unwrap().as_str(), "second");
464 /// assert_eq!(set.len(), 1);
465 ///
466 /// assert_eq!(set.take_by_index(0).unwrap().as_str(), "first");
467 /// assert_eq!(set.len(), 0);
468 /// ```
469 pub fn take_by_index(&mut self, index: usize) -> Option<T> {
470 self.root.take_node_by_index(index).map(|node| node.value )
471 }
472
473
474
475
476 /// Return iterator over all &T.
477 ///
478 ///
479 pub fn iter(&self) -> HiSetIterator<'_,T> {
480 HiSetIterator { set: self, start: 0, end: self.root.count }
481 }
482
483
484 /// Return double ended iterator over &T in given index range.
485 ///
486 /// # Examples:
487 /// ```
488 /// # use hitree::hiset::HiSet;
489 /// let s = HiSet::<i32>::from([0,1,2,3,4,5,6].into_iter());
490 /// let mut r = s.range_by_index(2..=5).map(|v| *v);
491 /// assert!(r.eq( [2,3,4,5].into_iter() ));
492 /// ```
493 pub fn range_by_index(&self, range: impl std::ops::RangeBounds<usize>) -> HiSetIterator<'_,T> {
494 use std::ops::Bound::*;
495 let start = match range.start_bound() {
496 Included(index) => *index,
497 Excluded(index) => *index + 1,
498 Unbounded => 0
499 };
500 let end = match range.end_bound() {
501 Included(index) => *index + 1,
502 Excluded(index) => *index,
503 Unbounded => self.root.count
504 };
505
506 HiSetIterator { set: self, start, end }
507 }
508
509 /// Return double ended iterator over &mut T in given index range.
510 ///
511 /// # Examples:
512 /// ```
513 /// # use hitree::hiset::HiSet;
514 /// let mut s = HiSet::<i32>::from([0,1,2,3,4,5,6].into_iter());
515 /// let mut r = s.range_by_index_mut(2..=5).map(|v| *v);
516 /// assert!(r.eq( [2,3,4,5].into_iter() ));
517 /// ```
518 pub fn range_by_index_mut(&mut self, range: impl std::ops::RangeBounds<usize>) -> HiSetIteratorMut<'_,T> {
519 use std::ops::Bound::*;
520 let start = match range.start_bound() {
521 Included(index) => *index,
522 Excluded(index) => *index + 1,
523 Unbounded => 0
524 };
525 let end = match range.end_bound() {
526 Included(index) => *index + 1,
527 Excluded(index) => *index,
528 Unbounded => self.root.count
529 };
530
531 HiSetIteratorMut { set: self, start, end }
532 }
533
534
535
536}
537
538#[test]
539fn test_hiset_range() {
540 let s = HiSet::<i32>::from([0,1,2,3,4,5,6].into_iter() );
541 let r = s.range_by_index(2..=5).map(|v| *v);
542 assert!(r.eq( [2,3,4,5].into_iter() ));
543}
544
545pub struct HiSetOwnedIterator<T>
546 where T: Ord
547{
548 root: Ref<T>,
549}
550
551impl <T> Iterator for HiSetOwnedIterator<T>
552 where T: Ord
553{
554 type Item = T;
555
556 fn next(&mut self) -> Option<Self::Item> {
557 // take leftmost node without bothering to re-balance or maintain node counts
558 self.root.consume_next()
559 }
560}
561
562impl <T> IntoIterator for HiSet<T>
563 where T: Ord
564{
565 type Item = T;
566 type IntoIter = HiSetOwnedIterator<T>;
567
568 /// Turn HiSet<T> into an Iterator of owned T
569 /// ```
570 /// # use hitree::hiset::HiSet;
571 /// let mut s = HiSet::<String>::new();
572 /// s.insert("This");
573 /// s.insert("is");
574 /// s.insert("a");
575 /// s.insert("test!");
576 ///
577 /// let mut i = s.into_iter();
578 /// assert_eq!(i.next(), Some("This".to_string()));
579 /// assert_eq!(i.next(), Some("a".to_string()));
580 /// assert_eq!(i.next(), Some("is".to_string()));
581 /// assert_eq!(i.next(), Some("test!".to_string()));
582 /// assert_eq!(i.next(), None);
583 /// ```
584 fn into_iter(self) -> Self::IntoIter {
585 HiSetOwnedIterator { root: self.root }
586 }
587}
588
589
590
591/// Get iterator over &T
592///
593/// # Examples:
594///
595/// ```
596/// # use hitree::hiset::HiSet;
597/// let mut s = HiSet::<String>::new();
598/// s.insert("This");
599/// s.insert("is");
600/// s.insert("a");
601/// s.insert("test!");
602///
603/// let mut i = s.iter();
604///
605/// assert_eq!(i.next(), Some(&"This".to_string()));
606/// assert_eq!(i.next(), Some(&"a".to_string()));
607/// assert_eq!(i.next(), Some(&"is".to_string()));
608/// assert_eq!(i.next(), Some(&"test!".to_string()));
609/// assert_eq!(i.next(), None);
610///
611/// ```
612pub struct HiSetIterator<'set,T>
613 where T: Ord
614{
615 set: &'set HiSet<T>,
616 start: usize,
617 end: usize,
618}
619
620impl <'set,T> Iterator for HiSetIterator<'set,T>
621 where T: Ord
622{
623 type Item = &'set T;
624
625 fn next(&mut self) -> Option<Self::Item> {
626 if self.start >= self.end {
627 None
628 } else {
629 let index_to_return = self.start;
630 self.start += 1;
631 self.set.get_by_index(index_to_return)
632 }
633 }
634}
635
636impl <'set,T> DoubleEndedIterator for HiSetIterator<'set,T>
637 where T: Ord
638{
639 fn next_back(&mut self) -> Option<Self::Item> {
640 if self.start >= self.end {
641 None
642 } else {
643 self.end -= 1;
644 self.set.get_by_index(self.end)
645 }
646 }
647}
648
649
650
651
652impl <'set,T> IntoIterator for &'set HiSet<T>
653 where T: Ord
654{
655 type Item = &'set T;
656 type IntoIter = HiSetIterator<'set,T>;
657
658 fn into_iter(self) -> Self::IntoIter {
659 self.iter()
660 }
661}
662
663
664/// Get iterator over &mut T
665///
666/// # Examples:
667///
668/// ```
669/// # use hitree::hiset::HiSet;
670/// let mut s = HiSet::<String>::new();
671/// s.insert("This");
672/// s.insert("is");
673/// s.insert("a");
674/// s.insert("test!");
675///
676/// let mut i = s.iter_mut();
677///
678/// assert_eq!(i.next(), Some(&mut "This".to_string()));
679/// assert_eq!(i.next(), Some(&mut "a".to_string()));
680/// assert_eq!(i.next(), Some(&mut "is".to_string()));
681/// assert_eq!(i.next(), Some(&mut "test!".to_string()));
682/// assert_eq!(i.next(), None);
683///
684/// ```
685pub struct HiSetIteratorMut<'set,T>
686 where T: Ord
687{
688 set: &'set mut HiSet<T>,
689 start: usize,
690 end: usize,
691}
692
693impl <'set,T> Iterator for HiSetIteratorMut<'set,T>
694 where T: Ord,
695{
696 type Item = &'set mut T;
697
698 fn next<'iter>(&mut self) -> Option<Self::Item> {
699 if self.start >= self.end {
700 None
701 } else {
702 let index_to_return = self.start;
703 self.start += 1;
704 unsafe { std::mem::transmute(self.set.get_by_index_mut(index_to_return)) }
705 }
706 }
707}
708
709impl <'set,T> DoubleEndedIterator for HiSetIteratorMut<'set,T>
710 where T: Ord
711{
712 fn next_back(&mut self) -> Option<Self::Item> {
713 if self.start >= self.end {
714 None
715 } else {
716 self.end -= 1;
717 unsafe { std::mem::transmute(self.set.get_by_index_mut(self.end)) }
718 }
719 }
720}
721
722impl <'set,T> IntoIterator for &'set mut HiSet<T>
723 where T: Ord
724{
725 type Item = &'set mut T;
726 type IntoIter = HiSetIteratorMut<'set,T>;
727
728 fn into_iter(self) -> Self::IntoIter {
729 self.iter_mut()
730 }
731}
732
733impl <T> HiSet<T>
734 where T: Ord
735{
736 pub fn iter_mut(&mut self) -> HiSetIteratorMut<'_,T> {
737 let end = self.root.count;
738 HiSetIteratorMut { set: self, start: 0, end }
739 }
740
741}
742
743
744impl <T,I,X,O> From<I> for HiSet<T>
745 where T: Ord,
746 I: Iterator<Item=X>,
747 O: Into<T>,
748 X: ToOwned<Owned=O>
749{
750 /// Construct HiSet from an Iterator of values that can be made into owned instances of T
751 ///
752 /// # Examples:
753 ///
754 /// ```
755 /// # use hitree::hiset::HiSet;
756 /// let s = HiSet::<String>::from( ["This","is","a","test!"].into_iter() );
757 ///
758 /// assert!(s.iter().eq(["This","a","is","test!"].iter()));
759 /// ```
760 fn from(mut iterator: I) -> Self {
761 let mut s = HiSet::<T>::new();
762 while let Some(value) = iterator.next() {
763 s.insert(value.to_owned());
764 }
765 s
766 }
767}
768
769
770//---------------- Ref -------------------------------------------------------
771
772impl <T> Ref<T>
773 where T: Ord
774{
775
776 pub fn to(node: Box<Node<T>>) -> Ref<T> {
777 let count = 1 + node.left.count + node.right.count;
778 Ref { count, node: Some(node) }
779 }
780
781
782 fn node(&self) -> Option<&Node<T>> {
783 self.node.as_deref()
784 }
785
786 fn node_mut(&mut self) -> Option<&mut Node<T>> {
787 self.node.as_deref_mut()
788 }
789
790
791 fn take(&mut self) -> Ref<T> {
792 std::mem::take(&mut *self)
793 }
794
795 fn take_left_subtree(&mut self) -> Ref<T> {
796 match self.node_mut() {
797 None => Ref::default(),
798 Some(node) => {
799 let left = node.left.take();
800 self.count -= left.count;
801 left
802 },
803 }
804 }
805
806 fn take_right_subtree(&mut self) -> Ref<T> {
807 match self.node_mut() {
808 None => Ref::default(),
809 Some(node) => {
810 let right = node.right.take();
811 self.count -= right.count;
812 right
813 },
814 }
815 }
816
817 fn is_empty(&self) -> bool {
818 self.node.is_none()
819 }
820
821 /*
822 #[inline]
823 fn left_count(&self) -> usize {
824 match self.node() {
825 None => 0,
826 Some(node) => node.left.count,
827 }
828 }
829
830 #[inline]
831 fn right_count(&self) -> usize {
832 match self.node() {
833 None => 0,
834 Some(node) => node.right.count,
835 }
836 }
837 */
838
839 #[inline]
840 fn balance(&self) -> isize {
841 self.node.as_deref().unwrap().balance() // balance only makes sense if there is a node, hence unwrap()
842 }
843
844
845 fn set_left(&mut self, subtree: Ref<T>) {
846 let node = self.node_mut().unwrap();
847 node.left = subtree;
848 self.count = node.count();
849 }
850
851 fn set_right(&mut self, subtree: Ref<T>) {
852 let node = self.node_mut().unwrap();
853 node.right = subtree;
854 self.count = node.count();
855 }
856
857 /*
858 self self
859 | |
860 [old_root] [new_root]
861 / \ ----> / \
862 [left_subtree] [new_root] [old_root] [right_subtree]
863 / \ / \
864 [mid_subtree] [right_subtree] [left_subtree] [mid subtree]
865
866 */
867 #[inline]
868 fn rotate_left(&mut self) {
869 let mut old_root = self.take();
870 let mut new_root = old_root.take_right_subtree();
871 let mid_subtree = new_root.take_left_subtree();
872 old_root.set_right(mid_subtree);
873 new_root.set_left(old_root);
874 *self = new_root;
875 }
876
877 /*
878 self self
879 | |
880 [old_root] [new_root]
881 / \ ----> / \
882 [new_root] [right_subtree] [left_subtree] [old_root]
883 / \ / \
884[left_subtree] [mid_subtree] [mid subtree] [right_subtree]
885
886 */
887 #[inline]
888 fn rotate_right(&mut self) {
889 let mut old_root = self.take();
890 let mut new_root = old_root.take_left_subtree();
891 let mid_subtree = new_root.take_right_subtree();
892 old_root.set_left(mid_subtree);
893 new_root.set_right(old_root);
894 *self = new_root;
895 }
896
897
898
899 /// insert is recursive as it needs to balance the tree on the way back up
900 fn insert(&mut self, new_node: Box<Node<T>>) -> bool {
901 match self.node_mut() {
902 None => { // there are no nodes in subtree rooted at this Ref.
903 *self = Ref::to(new_node);
904 true // we have inserted a value, return true
905 },
906 Some(node) => { // There is at least one node
907 match Ord::cmp(&node.value,&new_node.value) {
908 Ordering::Equal => {
909 false // already in there, return false
910 },
911 Ordering::Less => { // insert into right subtree
912 if node.right.insert(new_node) {
913 self.count += 1; // increase number of entries for subtree
914 if self.balance() > 1 { // too right heavy
915 // difference in height has become greater than 1, rotate subtree left
916 self.rotate_left();
917 }
918 true
919 } else {
920 false
921 }
922 },
923 Ordering::Greater => {
924 if node.left.insert(new_node) {
925 self.count += 1; // increase number of entries for subtree
926 if self.balance() < -1 { // too left heavy
927 // difference in height has become greater than 1, rotate subtree left
928 self.rotate_right();
929 }
930 true
931 } else {
932 false
933 }
934 }
935 }
936 }
937 }
938 }
939
940 /// Remove leftmost node from the subtree.
941 fn take_leftmost_node(&mut self) -> Option<Box<Node<T>>> {
942 match self.node_mut() {
943 None => None, // no node here, tell caller to remove his node
944 Some(node) => {
945 match node.left.take_leftmost_node() {
946 None => {
947 // there is no left node, we are the node to remove!
948 let mut removed_node = self.node.take().unwrap();
949 *self = removed_node.right.take();
950 Some(removed_node)
951 },
952 Some(removed_node) => {
953 self.count -= 1; // one node has been removed
954 if self.balance() > 1 { // if we are too right leaning now, restore balance
955 self.rotate_left();
956 }
957 Some(removed_node)
958 }
959 }
960 }
961 }
962 }
963
964 /// Remove rightmost node from the subtree.
965 fn take_rightmost_node(&mut self) -> Option<Box<Node<T>>> {
966 match self.node_mut() {
967 None => None, // no node here, tell caller to remove his node
968 Some(node) => {
969 match node.right.take_leftmost_node() {
970 None => {
971 // there is no left node, we are the node to remove!
972 let mut removed_node = self.node.take().unwrap();
973 *self = removed_node.left.take();
974 Some(removed_node)
975 },
976 Some(removed_node) => {
977 self.count -= 1; // one node has been removed
978 if self.balance() < -1 { // if we are too right leaning now, restore balance
979 self.rotate_right();
980 }
981 Some(removed_node)
982 }
983 }
984 }
985 }
986 }
987
988 fn take_node_by_key<KEY>(&mut self, key: &KEY) -> Option<Box<Node<T>>>
989 where KEY: ?Sized + Ord,
990 T: Borrow<KEY>
991 {
992 let res = if let Some(node) = self.node_mut() {
993 match Ord::cmp(node.value.borrow(), key) {
994 Ordering::Equal => { // this is the node to remove
995 match (node.left.is_empty(), node.right.is_empty()) {
996 (true, true) => { // leaf node, can be removed directly without consequences
997 self.node.take()
998 },
999 (false, true) => { // there is a left subtree, move it up
1000 let mut removed_node = self.node.take().unwrap();
1001 *self = removed_node.left.take();
1002 Some(removed_node)
1003 },
1004 (true, false) => { // there is a right subtree, move it up
1005 let mut removed_node = self.node.take().unwrap();
1006 *self = removed_node.right.take();
1007 Some(removed_node)
1008 }
1009 (false, false) => { // there are two subtrees, take the closest node from the one with more nodes and replace the removed node with it
1010 let mut removed_node = self.node.take().unwrap();
1011 let mut left_subtree = removed_node.left.take();
1012 let mut right_subtree = removed_node.right.take();
1013 let mut new_subtree_root_node = if left_subtree.count > right_subtree.count {
1014 left_subtree.take_rightmost_node().unwrap()
1015 } else {
1016 right_subtree.take_leftmost_node().unwrap()
1017 };
1018 new_subtree_root_node.left = left_subtree;
1019 new_subtree_root_node.right = right_subtree;
1020 let new_count = new_subtree_root_node.count();
1021 self.node = Some(new_subtree_root_node);
1022 self.count = new_count;
1023 // balance should not be an issue, we took from the bigger one
1024 Some(removed_node)
1025 }
1026 }
1027 },
1028 Ordering::Less => { // node must be in the right subtree
1029 let removed_node_maybe = node.right.take_node_by_key(key);
1030 match removed_node_maybe {
1031 None => None, // not found
1032 Some(removed_node) => {
1033 Some(removed_node)
1034 }
1035 }
1036 },
1037 Ordering::Greater => { // node must be in the left subtree
1038 match node.left.take_node_by_key(key) {
1039 None => None,
1040 Some(removed_node) => {
1041 Some(removed_node)
1042 }
1043 }
1044 }
1045 }
1046 } else {
1047 None
1048 };
1049 if res.is_some() {
1050 self.rebalance();
1051 }
1052 res
1053 }
1054
1055 fn take_node_by_index(&mut self, index_to_take: usize) -> Option<Box<Node<T>>> {
1056 let res = if let Some(node) = self.node_mut() {
1057 let index_of_this_node = node.left.count;
1058 match Ord::cmp(&index_of_this_node, &index_to_take) {
1059 Ordering::Equal => { // this is the node to remove
1060 match (node.left.is_empty(), node.right.is_empty()) {
1061 (true, true) => { // leaf node, can be removed directly without consequences
1062 self.node.take()
1063 },
1064 (false, true) => { // there is a left subtree, move it up
1065 let mut removed_node = self.node.take().unwrap();
1066 *self = removed_node.left.take();
1067 Some(removed_node)
1068 },
1069 (true, false) => { // there is a right subtree, move it up
1070 let mut removed_node = self.node.take().unwrap();
1071 *self = removed_node.right.take();
1072 Some(removed_node)
1073 }
1074 (false, false) => { // there are two subtrees, take the closest node from the one with more nodes and replace the removed node with it
1075 let mut removed_node = self.node.take().unwrap();
1076 let mut left_subtree = removed_node.left.take();
1077 let mut right_subtree = removed_node.right.take();
1078 let mut new_subtree_root_node = if left_subtree.count > right_subtree.count {
1079 left_subtree.take_rightmost_node().unwrap()
1080 } else {
1081 right_subtree.take_leftmost_node().unwrap()
1082 };
1083 new_subtree_root_node.left = left_subtree;
1084 new_subtree_root_node.right = right_subtree;
1085 let new_count = new_subtree_root_node.count();
1086 self.node = Some(new_subtree_root_node);
1087 self.count = new_count;
1088 // balance should not be an issue, we took from the bigger one
1089 Some(removed_node)
1090 }
1091 }
1092 },
1093 Ordering::Less => { // node must be in the right subtree
1094 let removed_node_maybe = node.right.take_node_by_index(index_to_take - index_of_this_node - 1);
1095 match removed_node_maybe {
1096 None => None, // not found
1097 Some(removed_node) => {
1098 Some(removed_node)
1099 }
1100 }
1101 },
1102 Ordering::Greater => { // node must be in the left subtree
1103 match node.left.take_node_by_index(index_to_take) {
1104 None => None,
1105 Some(removed_node) => {
1106 Some(removed_node)
1107 }
1108 }
1109 }
1110 }
1111 } else {
1112 None
1113 };
1114 if res.is_some() {
1115 self.rebalance();
1116 }
1117 res
1118 }
1119
1120 fn rebalance(&mut self) {
1121 if let Some(node) = self.node() {
1122 self.count = node.count();
1123 let balance = self.balance();
1124 if balance < -1 {
1125 self.rotate_right();
1126 } else if balance > 1 {
1127 self.rotate_left();
1128 }
1129 } else {
1130 self.count = 0;
1131 };
1132 }
1133
1134
1135 /// Take fist value without bothering to re-balance or maintain node counts. For use within owned iterator.
1136 fn consume_next(&mut self) -> Option<T> {
1137 // Take node from left subtree if any, or
1138 // Take node from yourself, replacing it with the right subtree root node if any
1139 match &mut self.node {
1140 None => None, // no nodes left, end of iteration
1141 Some(node) => {
1142 if let Some(from_left) = node.left.consume_next() {
1143 Some(from_left)
1144 } else {
1145 let right_node = node.right.node.take();
1146 let my_node = unsafe { std::mem::replace(&mut self.node, right_node ).unwrap_unchecked() };
1147 Some(my_node.value)
1148 }
1149 }
1150 }
1151 }
1152}
1153
1154impl <T> Default for Ref<T>
1155 where T: Ord
1156{
1157 /// Empty reference
1158 fn default() -> Self {
1159 Self { count: 0, node: None }
1160 }
1161}
1162
1163
1164//--------------- Node ------------------------------------------------------------
1165
1166
1167
1168
1169impl <T> Node<T>
1170 where T: Ord
1171{
1172 /// Creates a new Node with given value and empty left & right refs
1173 fn new(value: impl Into<T>) -> Box<Node<T>> {
1174 Box::new( Node { value: value.into(), left: Ref::default(), right: Ref::default() } )
1175 }
1176
1177 /// Calculate number of nodes including this node and any subtrees pointed to by left & right
1178 fn count(&self) -> usize {
1179 self.left.count + self.right.count + 1
1180 }
1181
1182 /// returns difference in height between right and left subtrees. >0 right is bigger, <0 left is bigger.
1183 #[inline]
1184 fn balance(&self) -> isize {
1185 tree_height(self.right.count) - tree_height(self.left.count)
1186 }
1187
1188 /// Borrow value of this node immutably
1189 pub fn borrow_value<B>(&self) -> &B
1190 where T: Borrow<B>,
1191 B: ?Sized
1192 {
1193 self.value.borrow()
1194 }
1195
1196 /// Borrow value of this node mutably
1197 pub fn borrow_value_mut<B>(&mut self) -> &mut B
1198 where T: BorrowMut<B>,
1199 B: ?Sized
1200 {
1201 self.value.borrow_mut()
1202 }
1203
1204}
1205
1206
1207
1208
1209
1210
1211
1212