binary_tree/
count.rs

1//! Counting tree implementation.
2//!
3//! ## When should you use `CountTree`?
4//!
5//! - You want to maintain a possibly large unsorted list.
6//! - You want to access, modify, insert, and delete elements at arbitrary
7//!   position with O(log(n)) time complexity.
8//! - You can tolerate O(n log(n)) time-complexity for (not implemented yet):
9//!   - splitting at arbitrary position
10//!   - truncating the length (complexity unclear)
11//!   - appending another list (complexity unclear)
12//! - You have less than 4.29 billion (`u32::MAX`) elements!
13//!
14//! ## Benchmarks
15//!
16//! The constants in the complexity bounds are not small enough to make an
17//! immediate switch from `Vec` to `CountTree`. In the benchmarks below, `*_ct`
18//! indicate `CountTree`, `*_ll` indicate `LinkedList` and `*_vec` indicate
19//! `Vec`. `from_iter_*` indicate the performance of using `collect()`,
20//! `insert_at_random_*` indicate that of inserting N elements at random
21//! positions, and `remove_at_random_*` indicate that of first creating an N
22//! sized object using `collect()` and then removing all N elements at random
23//! one-by-one. See `benches` directory for more details.
24//!
25//! ### Bench: N=2048
26//!
27//! ```sh
28//!      Running target/release/from_iter-86b0c3c534a106e8
29//!
30//! running 3 tests
31//! test from_iter_ct  ... bench:     316,440 ns/iter (+/- 92,914)
32//! test from_iter_ll  ... bench:     299,755 ns/iter (+/- 2,096)
33//! test from_iter_vec ... bench:       2,839 ns/iter (+/- 24)
34//!
35//!      Running target/release/insert-892f694b6341f60d
36//!
37//! running 3 tests
38//! test insert_at_random_ct  ... bench:   1,569,694 ns/iter (+/- 239,724)
39//! test insert_at_random_ll  ... bench:   2,882,318 ns/iter (+/- 20,555)
40//! test insert_at_random_vec ... bench:     570,018 ns/iter (+/- 3,710)
41//!
42//!      Running target/release/remove-12ee8f5093c08f36
43//!
44//! running 3 tests
45//! test remove_at_random_ct  ... bench:   1,800,295 ns/iter (+/- 5,761)
46//! test remove_at_random_ll  ... bench:   2,702,035 ns/iter (+/- 22,632)
47//! test remove_at_random_vec ... bench:     568,502 ns/iter (+/- 3,626)
48//! ```
49//!
50//! ### Bench: N=4096
51//!
52//! ```sh
53//!      Running target/release/from_iter-86b0c3c534a106e8
54//!
55//! running 3 tests
56//! test from_iter_ct  ... bench:     698,944 ns/iter (+/- 25,819)
57//! test from_iter_ll  ... bench:     579,370 ns/iter (+/- 12,161)
58//! test from_iter_vec ... bench:       5,582 ns/iter (+/- 26)
59//!
60//!      Running target/release/insert-892f694b6341f60d
61//!
62//! running 3 tests
63//! test insert_at_random_ct  ... bench:   3,495,019 ns/iter (+/- 123,200)
64//! test insert_at_random_ll  ... bench:  14,470,666 ns/iter (+/- 39,605)
65//! test insert_at_random_vec ... bench:   1,896,108 ns/iter (+/- 3,925)
66//!
67//!      Running target/release/remove-12ee8f5093c08f36
68//!
69//! running 3 tests
70//! test remove_at_random_ct  ... bench:   3,966,049 ns/iter (+/- 25,852)
71//! test remove_at_random_ll  ... bench:  11,981,076 ns/iter (+/- 77,152)
72//! test remove_at_random_vec ... bench:   1,909,054 ns/iter (+/- 5,475)
73//! ```
74//!
75//! ### Bench: N=8192
76//!
77//! ```sh
78//!      Running target/release/from_iter-86b0c3c534a106e8
79//!
80//! running 3 tests
81//! test from_iter_ct  ... bench:   1,422,694 ns/iter (+/- 224,526)
82//! test from_iter_ll  ... bench:   1,190,954 ns/iter (+/- 17,328)
83//! test from_iter_vec ... bench:      11,487 ns/iter (+/- 52)
84//!
85//!      Running target/release/insert-892f694b6341f60d
86//!
87//! running 3 tests
88//! test insert_at_random_ct  ... bench:   7,651,408 ns/iter (+/- 232,136)
89//! test insert_at_random_ll  ... bench:  67,522,968 ns/iter (+/- 508,089)
90//! test insert_at_random_vec ... bench:   8,062,135 ns/iter (+/- 46,386)
91//!
92//!      Running target/release/remove-12ee8f5093c08f36
93//!
94//! running 3 tests
95//! test remove_at_random_ct  ... bench:   8,972,611 ns/iter (+/- 99,882)
96//! test remove_at_random_ll  ... bench:  50,513,436 ns/iter (+/- 161,649)
97//! test remove_at_random_vec ... bench:   8,166,272 ns/iter (+/- 35,268)
98//! ```
99//!
100//! ### Conclusion
101//!
102//! In short, if you want to maintiain a list of type `T` such that:
103//!
104//! ```sh
105//! [number of elements] * [size of T] > 256 KB
106//! ```
107//!
108//! then `CountTree` might be a good choice, otherwise you are better off using
109//! `Vec`.
110
111use std::mem;
112use std::iter::FromIterator;
113use std::fmt::{self, Debug};
114
115#[cfg(feature="quickcheck")]
116use quickcheck::{Arbitrary, Gen};
117
118use Node;
119use NodeMut;
120use BinaryTree;
121use WalkAction;
122use iter::Iter as GenIter;
123use iter::IntoIter as GenIntoIter;
124
125pub type NodePtr<T> = Box<CountNode<T>>;
126
127macro_rules! index_walker {
128    ($index:ident, $node:ident, $up_count:ident, $stop:block) => {
129        {
130            let cur_index = $node.lcount() as usize + $up_count;
131            if $index < cur_index {
132                Left
133            } else if $index == cur_index {
134                $stop
135                Stop
136            } else {
137                $up_count = cur_index + 1;
138                Right
139            }
140        }
141    }
142}
143
144/// Counting tree.
145///
146/// A balanced binary tree which keeps track of total number of child nodes in
147/// each node, so that elements can be inserted and deleted using its in-order
148/// index. The algorithm used internally is a variation of [AVL Tree][avlwiki].
149/// Time complexities mentioned below are that of worst case scenario (and are
150/// of the same order as that of an AVL tree).
151///
152/// [avlwiki]: https://en.wikipedia.org/wiki/AVL_tree
153///
154/// # Examples
155///
156/// ```rust
157/// # extern crate binary_tree;
158/// # use binary_tree::count::CountTree;
159/// # fn main() {
160/// let mut ct: CountTree<i32> = CountTree::new();
161/// ct.push_front(20);
162/// ct.push_front(10);
163/// assert_eq!(ct.pop_back().unwrap(), 20);
164/// # }
165/// ```
166///
167/// You can also use `collect` to create one from an iterator. This has a time
168/// complexity of O(n), which is much more efficient than inserting iteratively.
169///
170/// ```rust
171/// # extern crate binary_tree;
172/// # use binary_tree::count::CountTree;
173/// # fn main() {
174/// let mut ct: CountTree<i32> = (0..100).collect();
175/// assert_eq!(ct.remove(32), 32);
176/// # }
177/// ```
178pub struct CountTree<T>(Option<NodePtr<T>>);
179
180impl<T> CountTree<T> {
181    fn root_must(&mut self) -> &mut CountNode<T> {
182        &mut **self.0.as_mut().unwrap()
183    }
184
185    /// Returns an empty `CountTree`
186    pub fn new() -> CountTree<T> {
187        CountTree(None)
188    }
189
190    /// Returns `true` if the tree contains no elements.
191    pub fn is_empty(&self) -> bool {
192        self.0.is_none()
193    }
194
195    /// Returns the number elements in the tree. Time complexity: O(1)
196    pub fn len(&self) -> usize {
197        self.root().map_or(0, |node| node.count as usize)
198    }
199
200    /// Clears the tree, dropping all elements iteratively.
201    pub fn clear(&mut self) {
202        let mut inner = None;
203        mem::swap(&mut self.0, &mut inner);
204        let _: GenIntoIter<CountNode<T>> = GenIntoIter::new(inner);
205    }
206
207    /// Returns the element at the given index, or `None` if index is out of
208    /// bounds. Time complexity: O(log(n))
209    pub fn get(&self, index: usize) -> Option<&T> {
210        use WalkAction::*;
211
212        if index >= self.len() {
213            None
214        } else {
215            let mut val = None;
216            let mut up_count = 0;
217            self.root().unwrap().walk(|node| {
218                index_walker!(index, node, up_count, {
219                    val = Some(node.value());
220                })
221            });
222            debug_assert!(val.is_some());
223            val
224        }
225    }
226
227    /// Returns a mutable reference to the element at the given index, or `None`
228    /// if out of bounds. Time complexity: O(log(n))
229    pub fn get_mut(&mut self, index: usize) -> Option<&mut T> {
230        use WalkAction::*;
231
232        if index >= self.len() {
233            None
234        } else {
235            let mut val = None;
236            let mut up_count = 0;
237            let root = self.root_must();
238            root.walk_mut(|node| index_walker!(index, node, up_count, {}),
239                          |node| val = Some(node.value_mut()));
240            debug_assert!(val.is_some());
241            val
242        }
243    }
244
245    /// Inserts an element at the given index. Time complexity: O(log(n))
246    ///
247    /// ## Panics
248    ///
249    /// Panics if index is greater than `self.len()`
250    pub fn insert(&mut self, index: usize, value: T) {
251        use WalkAction::*;
252
253        let len = self.len();
254        if index == 0 {
255            self.push_front(value);
256        } else if index < len {
257            let new_node = Box::new(CountNode::new(value));
258            let mut up_count = 0;
259            let root = self.root_must();
260            root.walk_reshape(|node| index_walker!(index, node, up_count, {}),
261                              move |node| {
262                                  node.insert_before(new_node,
263                                                     |node, _| node.rebalance());
264                              },
265                              |node, _| node.rebalance());
266        } else if index == len {
267            self.push_back(value);
268        } else {
269            panic!("index out of bounds!");
270        }
271    }
272
273    /// Prepends an element at the beginning.
274    pub fn push_front(&mut self, value: T) {
275        let new_node = Box::new(CountNode::new(value));
276        if self.is_empty() {
277            self.0 = Some(new_node);
278        } else {
279            self.root_must().walk_reshape(|_| WalkAction::Left,
280                                          move |node| {
281                                              node.insert_left(Some(new_node));
282                                          },
283                                          |node, _| node.rebalance());
284        }
285    }
286
287    /// Appends an element at the end.
288    pub fn push_back(&mut self, value: T) {
289        let new_node = Box::new(CountNode::new(value));
290        if self.is_empty() {
291            self.0 = Some(new_node);
292        } else {
293            self.root_must().walk_reshape(|_| WalkAction::Right,
294                                          move |node| {
295                                              node.insert_right(Some(new_node));
296                                          },
297                                          |node, _| node.rebalance());
298        }
299    }
300
301    /// Removes the element at the given index. Time complexity: O(log(n))
302    ///
303    /// ## Panics
304    ///
305    /// Panics if index is out of bounds.
306    pub fn remove(&mut self, index: usize) -> T {
307        use WalkAction::*;
308
309        let len = self.len();
310        if index == 0 {
311            self.pop_front().expect("Tree is empty!")
312        } else if index + 1 < len {
313            let mut up_count = 0;
314            let root = self.root_must();
315            root.walk_extract(|node| index_walker!(index, node, up_count, {}),
316                              |node, ret| {
317                                  *ret = node.try_remove(|node, _| node.rebalance());
318                              },
319                              |node, _| node.rebalance())
320                .unwrap()
321                .into_value()
322        } else if index + 1 == len {
323            self.pop_back().unwrap()
324        } else {
325            panic!("index out of bounds!");
326        }
327    }
328
329    /// Removes and returns the first element, or `None` if empty.
330    pub fn pop_front(&mut self) -> Option<T> {
331        if self.is_empty() {
332            None
333        } else if self.len() == 1 {
334            Some(self.0.take().unwrap().into_value())
335        } else {
336            let root = self.root_must();
337            Some(root.walk_extract(|_| WalkAction::Left,
338                                   |node, ret| {
339                                       if let Some(mut right) = node.detach_right() {
340                                           mem::swap(&mut *right, node);
341                                           *ret = Some(right);
342                                       }
343                                   },
344                                   |node, _| node.rebalance())
345                     .unwrap()
346                     .into_value())
347        }
348    }
349
350    /// Removes and returns the last element, or `None` if empty.
351    pub fn pop_back(&mut self) -> Option<T> {
352        // FIXME Ewww! Code duplication!
353        if self.is_empty() {
354            None
355        } else if self.len() == 1 {
356            Some(self.0.take().unwrap().into_value())
357        } else {
358            let root = self.root_must();
359            Some(root.walk_extract(|_| WalkAction::Right,
360                                   |node, ret| {
361                                       if let Some(mut left) = node.detach_left() {
362                                           mem::swap(&mut *left, node);
363                                           *ret = Some(left);
364                                       }
365                                   },
366                                   |node, _| node.rebalance())
367                     .unwrap()
368                     .into_value())
369        }
370    }
371
372    // TODO ? iter_mut
373    // TODO { O(n) } truncate, append, split_off, retain
374}
375
376impl<T> BinaryTree for CountTree<T> {
377    type Node = CountNode<T>;
378
379    fn root(&self) -> Option<&Self::Node> {
380        self.0.as_ref().map(|nodeptr| &**nodeptr)
381    }
382}
383
384impl<T> Debug for CountTree<T>
385    where T: Debug
386{
387    fn fmt(&self, f: &mut fmt::Formatter) -> Result<(), fmt::Error> {
388        let mut ds = f.debug_struct("CountTree");
389        if let Some(ref root) = self.0 {
390            ds.field("_count", &root.count);
391            ds.field("_height", &root.height);
392            ds.field("_inner", &DebugPrefix("^", root));
393        } else {
394            ds.field("_count", &0);
395            ds.field("_height", &0);
396            ds.field("_inner", &DebugPrefix("^", &()));
397        }
398        ds.finish()
399    }
400}
401
402impl<T> Drop for CountTree<T> {
403    fn drop(&mut self) {
404        self.clear();
405    }
406}
407
408fn is_power(v: u32) -> bool {
409    if v == 0 {
410        false
411    } else {
412        v & (v - 1) == 0
413    }
414}
415
416fn exp_floor_log(v: u32) -> u32 {
417    if v == 0 || is_power(v) {
418        v
419    } else {
420        let mut efl = v - 1;
421        efl |= efl >> 1;
422        efl |= efl >> 2;
423        efl |= efl >> 4;
424        efl |= efl >> 8;
425        efl |= efl >> 16;
426        efl += 1;
427        efl >> 1
428    }
429}
430
431impl<T> FromIterator<T> for CountTree<T> {
432    /// Time complexity: &Theta;(n + log<sup>2</sup>(n))
433    fn from_iter<I>(iterable: I) -> Self
434        where I: IntoIterator<Item = T>
435    {
436        use WalkAction::*;
437
438        let mut iter = iterable.into_iter();
439        if let Some(item) = iter.next() {
440            let mut node = Box::new(CountNode::new(item));
441            let mut count = 1;
442            for item in iter {
443                let mut new_node = Box::new(CountNode::new(item));
444                new_node.insert_left(Some(node));
445                node = new_node;
446                count += 1;
447                let rcount = if is_power(count + 1) {
448                    count >> 1
449                } else {
450                    count
451                };
452                let mut rotate_points = 1;
453                while rcount & rotate_points == rotate_points {
454                    node.rotate_right().unwrap();
455                    rotate_points <<= 1;
456                    rotate_points |= 1;
457                }
458            }
459            let balanced_till = exp_floor_log(count + 1) - 1;
460            count = node.lcount() + 1; // not needed
461            while count > balanced_till {
462                node.rotate_right().unwrap();
463                node.right
464                    .as_mut()
465                    .unwrap()
466                    .walk_reshape(|node| {
467                                      if node.balance_factor() > 1 {
468                                          node.rotate_right().unwrap();
469                                          Right
470                                      } else {
471                                          Stop
472                                      }
473                                  },
474                                  |_| (),
475                                  |_, _| ());
476                count = node.lcount() + 1;
477            }
478            CountTree(Some(node))
479        } else {
480            CountTree::new()
481        }
482    }
483}
484
485impl<'a, T> IntoIterator for &'a CountTree<T> {
486    type Item = &'a T;
487    type IntoIter = Iter<'a, T>;
488
489    fn into_iter(self) -> Self::IntoIter {
490        Iter {
491            inner: GenIter::new(self.root()),
492            remaining: self.len(),
493        }
494    }
495}
496
497pub struct Iter<'a, T: 'a> {
498    inner: GenIter<'a, CountNode<T>>,
499    remaining: usize,
500}
501
502impl<'a, T> Iterator for Iter<'a, T> {
503    type Item = &'a T;
504
505    fn next(&mut self) -> Option<&'a T> {
506        if self.remaining > 0 {
507            self.remaining -= 1;
508        }
509        self.inner.next()
510    }
511
512    fn size_hint(&self) -> (usize, Option<usize>) {
513        (self.remaining, Some(self.remaining))
514    }
515}
516
517impl<'a, T> ExactSizeIterator for Iter<'a, T> {}
518
519impl<T> IntoIterator for CountTree<T> {
520    type Item = T;
521    type IntoIter = IntoIter<T>;
522
523    fn into_iter(mut self) -> Self::IntoIter {
524        let len = self.len();
525        let mut inner = None;
526        mem::swap(&mut self.0, &mut inner);
527        IntoIter {
528            inner: GenIntoIter::new(inner),
529            remaining: len,
530        }
531    }
532}
533
534pub struct IntoIter<T> {
535    inner: GenIntoIter<CountNode<T>>,
536    remaining: usize,
537}
538
539impl<T> Iterator for IntoIter<T> {
540    type Item = T;
541
542    fn next(&mut self) -> Option<T> {
543        if self.remaining > 0 {
544            self.remaining -= 1;
545        }
546        self.inner.next()
547    }
548
549    fn size_hint(&self) -> (usize, Option<usize>) {
550        (self.remaining, Some(self.remaining))
551    }
552}
553
554impl<T> ExactSizeIterator for IntoIter<T> {}
555
556/// Node of a `CountTree`.
557///
558/// The only way of getting your hands on a `CountNode` is through
559/// [`CountTree::root()`](struct.CountTree.html#method.root) method which
560/// returns a shared reference to its root.  Thus `NodeMut` methods are not
561/// accessible to users.
562pub struct CountNode<T> {
563    val: T,
564    left: Option<NodePtr<T>>,
565    right: Option<NodePtr<T>>,
566    count: u32,
567    height: u16,
568}
569
570impl<T> CountNode<T> {
571    fn new(val: T) -> CountNode<T> {
572        CountNode {
573            val: val,
574            left: None,
575            right: None,
576            count: 1,
577            height: 0,
578        }
579    }
580
581    fn lcount(&self) -> u32 {
582        self.left.as_ref().map_or(0, |tree| tree.count)
583    }
584
585    fn rcount(&self) -> u32 {
586        self.right.as_ref().map_or(0, |tree| tree.count)
587    }
588
589    // generalized version of AVL tree balance factor: h(left) - h(right)
590    fn balance_factor(&self) -> i32 {
591        self.left.as_ref().map_or(-1, |node| node.height as i32) -
592            self.right.as_ref().map_or(-1, |node| node.height as i32)
593    }
594
595    // AVL tree algorithm
596    fn rebalance(&mut self) {
597        if self.balance_factor() > 1 {
598            self.left.as_mut().map(|node| {
599                if node.balance_factor() < 0 {
600                    node.rotate_left().unwrap();
601                }
602            });
603            self.rotate_right().unwrap();
604        } else if self.balance_factor() < -1 {
605            self.right.as_mut().map(|node| {
606                if node.balance_factor() > 0 {
607                    node.rotate_right().unwrap();
608                }
609            });
610            self.rotate_left().unwrap();
611        }
612    }
613
614    fn update_stats(&mut self) {
615        use std::cmp::max;
616        self.count = self.lcount() + self.rcount() + 1;
617        self.height = max(self.left.as_ref().map_or(0, |tree| tree.height),
618                          self.right.as_ref().map_or(0, |tree| tree.height));
619        if self.count > 1 {
620            self.height += 1;
621        }
622    }
623
624    fn into_value(self) -> T {
625        debug_assert!(self.count == 1, "count = {}", self.count);
626        self.val
627    }
628}
629
630impl<T> Node for CountNode<T> {
631    type Value = T;
632
633    fn left(&self) -> Option<&Self> {
634        self.left.as_ref().map(|st| &**st)
635    }
636
637    fn right(&self) -> Option<&Self> {
638        self.right.as_ref().map(|st| &**st)
639    }
640
641    fn value(&self) -> &T {
642        &self.val
643    }
644}
645
646impl<T> NodeMut for CountNode<T> {
647    type NodePtr = NodePtr<T>;
648
649    fn detach_left(&mut self) -> Option<Self::NodePtr> {
650        let tree = self.left.take();
651        self.update_stats();
652        tree
653    }
654
655    fn detach_right(&mut self) -> Option<Self::NodePtr> {
656        let tree = self.right.take();
657        self.update_stats();
658        tree
659    }
660
661    fn insert_left(&mut self, mut tree: Option<Self::NodePtr>) -> Option<Self::NodePtr> {
662        mem::swap(&mut self.left, &mut tree);
663        self.update_stats();
664        tree
665    }
666
667    fn insert_right(&mut self, mut tree: Option<Self::NodePtr>) -> Option<Self::NodePtr> {
668        mem::swap(&mut self.right, &mut tree);
669        self.update_stats();
670        tree
671    }
672
673    fn value_mut(&mut self) -> &mut T {
674        &mut self.val
675    }
676
677    fn into_parts(self) -> (T, Option<Self::NodePtr>, Option<Self::NodePtr>) {
678        (self.val, self.left, self.right)
679    }
680
681    fn left_mut(&mut self) -> Option<&mut Self> {
682        self.left.as_mut().map(|l| &mut **l)
683    }
684
685    fn right_mut(&mut self) -> Option<&mut Self> {
686        self.right.as_mut().map(|r| &mut **r)
687    }
688}
689
690struct DebugPrefix<'a, 'b, T: 'b>(&'a str, &'b T);
691
692impl<'a, 'b, T> Debug for DebugPrefix<'a, 'b, T>
693    where T: Debug
694{
695    fn fmt(&self, f: &mut fmt::Formatter) -> Result<(), fmt::Error> {
696        try!(f.write_str(self.0));
697        self.1.fmt(f)
698    }
699}
700
701impl<T> Debug for CountNode<T>
702    where T: Debug
703{
704    fn fmt(&self, f: &mut fmt::Formatter) -> Result<(), fmt::Error> {
705        let mut dt = f.debug_tuple("");
706        dt.field(&self.val);
707        if let Some(ref left) = self.left {
708            dt.field(&DebugPrefix("L", left));
709        }
710        if let Some(ref right) = self.right {
711            dt.field(&DebugPrefix("R", right));
712        }
713        dt.finish()
714    }
715}
716
717#[cfg(feature="quickcheck")]
718impl Arbitrary for CountTree<usize> {
719    fn arbitrary<G: Gen>(g: &mut G) -> CountTree<usize> {
720        let size = { let s = g.size(); g.gen_range(0, s) };
721        let mut ct = CountTree::new();
722        for i in 0..size {
723            ct.insert(g.gen_range(0, i + 1), i);
724        }
725        ct
726    }
727
728    fn shrink(&self) -> Box<Iterator<Item=CountTree<usize>>> {
729        Box::new(quickcheck::Shrinker::new(self))
730    }
731}
732
733#[cfg(feature="quickcheck")]
734impl<T> Clone for CountTree<T>
735    where T: Clone
736{
737    fn clone(&self) -> Self {
738        CountTree(self.0.clone())
739    }
740}
741
742#[cfg(feature="quickcheck")]
743impl<T> Clone for CountNode<T>
744    where T: Clone
745{
746    fn clone(&self) -> Self {
747        CountNode {
748            val: self.val.clone(),
749            left: self.left.clone(),
750            right: self.right.clone(),
751            count: self.count,
752            height: self.height,
753        }
754    }
755}
756
757#[cfg(feature="quickcheck")]
758pub mod quickcheck {
759    use super::CountTree;
760    use BinaryTree;
761
762    #[derive(Clone, Copy)]
763    enum ShrinkerState {
764        Value,
765        Left,
766        Right,
767        End,
768    }
769
770    pub struct Shrinker {
771        inner: CountTree<usize>,
772        state: ShrinkerState,
773    }
774
775    impl Shrinker {
776        pub fn new(inner: &CountTree<usize>) -> Shrinker {
777            Shrinker {
778                inner: inner.clone(),
779                state: ShrinkerState::Value,
780            }
781        }
782    }
783
784    impl Iterator for Shrinker {
785        type Item = CountTree<usize>;
786
787        fn next(&mut self) -> Option<Self::Item> {
788            if self.inner.0.is_none() {
789                None
790            } else {
791                use self::ShrinkerState::*;
792                let root = self.inner.root().unwrap();
793                match self.state {
794                    Value => {
795                        let mut ct = CountTree::new();
796                        if root.count > 1 {
797                            ct.push_back(root.val);
798                            self.state = Left;
799                        } else {
800                            self.state = End;
801                        }
802                        Some(ct)
803                    }
804                    Left => {
805                        self.state = Right;
806                        Some(CountTree(root.left.clone()))
807                    }
808                    Right => {
809                        self.state = End;
810                        Some(CountTree(root.right.clone()))
811                    }
812                    End => {
813                        None
814                    }
815                }
816            }
817        }
818    }
819}
820
821#[cfg(test)]
822mod tests {
823    use BinaryTree;
824    use NodeMut;
825    use super::CountNode;
826    use super::CountTree;
827    use test::compute_level;
828    use test::Level;
829
830    fn test_nodes() -> Box<CountNode<u32>> {
831        let mut cn = Box::new(CountNode::new(7));
832        cn.insert_before(Box::new(CountNode::new(8)), |_, _| ());
833        cn.insert_before(Box::new(CountNode::new(12)), |_, _| ());
834        cn.insert_right(Some(Box::new(CountNode::new(5))));
835        cn
836    }
837
838    #[test]
839    fn custom() {
840        let ct = CountTree(Some(test_nodes()));
841        assert_eq!(ct.get(0), Some(&8));
842        assert_eq!(ct.get(1), Some(&12));
843        assert_eq!(ct.get(2), Some(&7));
844        assert_eq!(ct.get(3), Some(&5));
845        assert_eq!(ct.get(4), None);
846        let mut ct = ct;
847        ct.get_mut(3).map(|v| *v = 100);
848        assert_eq!(ct.get(3), Some(&100));
849    }
850
851    #[test]
852    fn counting() {
853        let cn = test_nodes();
854        assert_eq!(cn.lcount(), 2);
855        assert_eq!(cn.rcount(), 1);
856        assert_eq!(cn.count, 4);
857        assert_eq!(cn.height, 2);
858    }
859
860    #[test]
861    fn rebalance() {
862        let mut cn = test_nodes();
863        assert_eq!(cn.balance_factor(), 1);
864        cn.detach_right();
865        cn.rebalance();
866        assert_eq!(cn.balance_factor(), 0);
867        assert_eq!(compute_level(&*cn, 1), Level::Balanced(2));
868        let ct = CountTree(Some(cn));
869        assert_eq!(ct.get(0), Some(&8));
870        assert_eq!(ct.get(1), Some(&12));
871        assert_eq!(ct.get(2), Some(&7));
872        assert_eq!(ct.get(3), None);
873    }
874
875    #[test]
876    fn insert() {
877        let mut ct = CountTree::new();
878        assert_eq!(ct.get(0), None);
879        ct.insert(0, 2);
880        ct.insert(0, 3);
881        ct.insert(0, 4);
882        ct.insert(0, 5);
883        ct.insert(0, 6);
884        assert_eq!(ct.get(0), Some(&6));
885        assert_eq!(ct.get(1), Some(&5));
886        assert_eq!(ct.get(2), Some(&4));
887        ct.insert(0, 7);
888        assert_eq!(ct.get(4), Some(&3));
889        assert_eq!(ct.get(5), Some(&2));
890        assert_eq!(ct.root().unwrap().height, 2);
891        assert_eq!(compute_level(ct.root().unwrap(), 1), Level::Balanced(3));
892        ct.insert(6, 1);
893        assert_eq!(ct.get(6), Some(&1));
894        assert_eq!(ct.root().unwrap().height, 3);
895        assert_eq!(compute_level(ct.root().unwrap(), 1), Level::Balanced(4));
896    }
897
898    #[test]
899    fn from_iter() {
900        let ct: CountTree<_> = (0..63).collect();
901        let root = ct.root().unwrap();
902        assert_eq!(root.height, 5);
903        assert_eq!(compute_level(root, 0), Level::Balanced(6));
904
905        let ct: CountTree<_> = (0..94).collect();
906        let root = ct.root().unwrap();
907        assert_eq!(root.balance_factor(), -1);
908        assert_eq!(root.height, 6);
909        assert_eq!(compute_level(root, 1), Level::Balanced(7));
910    }
911
912    #[test]
913    fn remove() {
914        let mut ct: CountTree<_> = (0..94).collect();
915        for i in 0..20 {
916            assert_eq!(ct.remove(64), 64 + i);
917            assert!(compute_level(ct.root().unwrap(), 1).is_balanced());
918        }
919    }
920}