summaryrefslogtreecommitdiff
path: root/src/iterators.rs
blob: 2d7a673634336e2cf51c54290e712e4225fc277f (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
use ::tree;
use ::node;
use std::collections::Bound;

pub struct RangePairIter<'a, K:'a+Ord+Copy,D:'a> {
    tree: &'a tree::AVLTree<K, D>,
    from: Bound<K>,
    to: Bound<K>,
    prev: Option<&'a K>,
}

impl<'a, K:'a+Ord+Copy,D:'a> RangePairIter<'a, K, D> {

    pub fn new(tree: &'a tree::AVLTree<K,D>, lower: Bound<K>, upper: Bound<K>) -> RangePairIter<'a,K,D>{
        RangePairIter{tree: tree, from: lower, to: upper, prev:None}
    }

    fn get_next_key_under(&mut self, root: &'a Box<node::Node<K,D>>) -> Option<(&'a K,&'a D)>{
        let res = self.get_next_pair(root).and_then(|p| self.check_upper_bound(p));
        if let Some((key,_)) = res { self.prev = Some(key) }
        return res
    }

    fn get_next_pair(&mut self, root: &'a Box<node::Node<K,D>>) -> Option<(&'a K, &'a D)>{
        match self.prev{
            None => self.get_lower_bound_pair(root),
            Some(key) => node::min_after::<K,D>(key, root)
        }
    }

    fn get_lower_bound_pair(&self, root: &'a Box<node::Node<K,D>>) -> Option<(&'a K, &'a D)>{
        match self.from {
            Bound::Included(ref key) => node::search_pair(key, root).or_else(|| node::min_after(key, root)),
            Bound::Excluded(ref key) => node::min_after(key, root),
            Bound::Unbounded => Some(node::min_pair(root))
        }
    }

    fn check_upper_bound(&self, current: (&'a K, &'a D)) -> Option<(&'a K, &'a D)> {
        let ok = match self.to {
            Bound::Included(ref key) => current.0 <= key,
            Bound::Excluded(ref key) => current.0 < key,
            Bound::Unbounded => true
        };
        return if ok { Some(current) } else { None };
    }
}

impl<'a, K:'a+Ord+Copy,D:'a> Iterator for RangePairIter<'a, K, D> {

    type Item = (&'a K,&'a D);

    fn next(&mut self) -> Option<(&'a K,&'a D)> {
        self.tree.root.as_ref().map_or(None,|node| self.get_next_key_under(node))
    }


}

#[test]
fn test_iterators(){
    let mut tree = tree::AVLTree::<u64,i32>::new();
    tree.insert(18, 1337);
    tree.insert(13, 1338);
    tree.insert(17, 1339);
    tree.insert(10, 1321);
    tree.insert(1, 1321);
    tree.insert(3, 1322);
    let init_key = 0;
    let mut iter = RangePairIter::<u64,i32>{tree: &tree, prev: Some(&init_key), from: Bound::Unbounded, to: Bound::Unbounded};
    assert_eq!(iter.next().expect("should have a few values").0, &1);
    assert_eq!(iter.next().expect("should have a few values").0, &3);
    assert_eq!(iter.next().expect("should have a few values").0, &10);
    assert_eq!(iter.next().expect("should have a few values").0, &13);
    assert_eq!(iter.next().expect("should have a few values").0, &17);
    assert_eq!(iter.next().expect("should have a few values").0, &18);
    assert!(iter.next().is_none());

    let mut iter = RangePairIter::new(&tree, Bound::Unbounded, Bound::Unbounded);
    assert_eq!(iter.next().expect("should have a few values").0, &1);
    assert_eq!(iter.next().expect("should have a few values").0, &3);
    assert_eq!(iter.next().expect("should have a few values").0, &10);
    assert_eq!(iter.next().expect("should have a few values").0, &13);
    assert_eq!(iter.next().expect("should have a few values").0, &17);
    assert_eq!(iter.next().expect("should have a few values").0, &18);
    assert!(iter.next().is_none());

    let mut iter = RangePairIter::new(&tree, Bound::Included(3), Bound::Included(17));
    assert_eq!(iter.next().expect("should have a few values").0, &3);
    assert_eq!(iter.next().expect("should have a few values").0, &10);
    assert_eq!(iter.next().expect("should have a few values").0, &13);
    assert_eq!(iter.next().expect("should have a few values").0, &17);
    assert!(iter.next().is_none());

    let mut iter = RangePairIter::new(&tree, Bound::Excluded(1), Bound::Excluded(18));
    assert_eq!(iter.next().expect("should have a few values").0, &3);
    assert_eq!(iter.next().expect("should have a few values").0, &10);
    assert_eq!(iter.next().expect("should have a few values").0, &13);
    assert_eq!(iter.next().expect("should have a few values").0, &17);
    assert!(iter.next().is_none());
}