1use 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
144pub 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 pub fn new() -> CountTree<T> {
187 CountTree(None)
188 }
189
190 pub fn is_empty(&self) -> bool {
192 self.0.is_none()
193 }
194
195 pub fn len(&self) -> usize {
197 self.root().map_or(0, |node| node.count as usize)
198 }
199
200 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 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 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 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 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 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 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 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 pub fn pop_back(&mut self) -> Option<T> {
352 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 }
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 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; 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
556pub 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 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 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}