immutable_seq/
node.rs

1use std::ops;
2
3use lazy::{Lazy, strict, value};
4use self::Node::{Leaf,Node2,Node3};
5use measure::Measure;
6use digit::Digit;
7use digit::Digit::{One,Two};
8
9/// A node in a 2-3 tree.
10///
11/// The children are stared as lazy references
12#[derive(Debug)]
13pub enum Node<T, M>
14{
15    Leaf(T),
16    Node2(M, Lazy<Node<T,M>>, Lazy<Node<T,M>>),
17    Node3(M, Lazy<Node<T,M>>, Lazy<Node<T,M>>, Lazy<Node<T,M>>),
18}
19
20/// Construct a lazy reference to a leaf
21pub fn leaf<T,M>(v: T) -> Lazy<Node<T,M>> {
22    strict(Leaf(v))
23}
24
25/// Construct a lazy reference to a node with two children
26pub fn node2<T,M>(left: Lazy<Node<T,M>>, right: Lazy<Node<T,M>>) -> Lazy<Node<T,M>>
27    where T: Measure<M> + 'static,
28          M: ops::Add<Output=M> + Copy + 'static
29{
30    lazy!{
31        let m = left.measure() + right.measure();
32        let left: Lazy<Node<T,M>> = left;
33        value(Node2(m, left, right))
34    }
35}
36
37/// Construct a lazy reference to a node with three children
38pub fn node3<T,M>(left: Lazy<Node<T,M>>, middle: Lazy<Node<T,M>>, right: Lazy<Node<T,M>>) -> Lazy<Node<T,M>>
39    where T: Measure<M> + 'static,
40          M: ops::Add<Output=M> + Copy + 'static
41{
42    lazy!{
43        let m = left.measure() + middle.measure() + right.measure();
44        value(Node3(m, left, middle, right))
45    }
46}
47
48#[doc(hidden)]
49#[macro_export]
50macro_rules! node {
51    ($left: expr, $right: expr) => {
52        Node::node2($left, $right)
53    };
54    ($left: expr, $middle: expr, $right: expr) => {
55        Node::node3($left, $middle, $right)
56    }
57}
58
59pub fn lookup<T,M,P>(pred: P, i: M, node: &Node<T,M>) -> (&T,M)
60    where T: Measure<M> + 'static,
61          M: ops::Add<Output=M> + Copy + 'static,
62          P: Fn(M) -> bool
63{
64    match *node {
65        Leaf(ref x) => (x, i),
66        Node2(_, ref left, ref right) => {
67            let i1 = i + left.measure();
68            if pred(i1) {
69                lookup(pred, i, left)
70            } else {
71                lookup(pred, i1, right)
72            }
73        },
74        Node3(_, ref left, ref middle, ref right) => {
75            let i1 = i + left.measure();
76            if pred(i1) {
77                lookup(pred, i, left)
78            } else {
79                let i2 = i1 + middle.measure();
80                if pred(i2) {
81                    lookup(pred, i1, middle)
82                } else {
83                    lookup(pred, i2, right)
84                }
85            }
86        }
87    }
88}
89
90pub fn adjust<T,M,P,F>(func: F, pred: P, i: M, node: &Node<T,M>) -> Lazy<Node<T,M>>
91    where T: Measure<M> + 'static,
92          M: ops::Add<Output=M> + Copy + 'static,
93          P: Fn(M) -> bool,
94          F: FnOnce(&T) -> T
95{
96    match *node {
97        Leaf(ref x) => leaf(func(x)),
98        Node2(_, ref left, ref right) => {
99            let i1 = i + left.measure();
100            if pred(i1) {
101                node2(adjust(func, pred, i, left), right.clone())
102            } else {
103                node2(left.clone(), adjust(func, pred, i1, right))
104            }
105        },
106        Node3(_, ref left, ref middle, ref right) => {
107            let i1 = i + left.measure();
108            if pred(i1) {
109                node3(adjust(func, pred, i, left), middle.clone(), right.clone())
110            } else {
111                let i2 = i1 + middle.measure();
112                if pred(i2) {
113                    node3(left.clone(), adjust(func, pred, i1, middle), right.clone())
114                } else {
115                    node3(left.clone(), middle.clone(), adjust(func, pred, i2, right))
116                }
117            }
118        }
119    }
120}
121
122pub fn split_once<'a,T,M,P>(pred: &P, i: M, node: &'a Node<T,M>)
123                    -> (Option<Digit<T,M>>, &'a Lazy<Node<T,M>>, Option<Digit<T,M>>)
124    where T: Measure<M> + 'static,
125          M: ops::Add<Output=M> + Copy + 'static,
126          P: Fn(M) -> bool
127{
128    match *node {
129        Leaf(_) => panic!("split_once on Leaf"),
130        Node2(_, ref left, ref right) => {
131            let i1 = i + left.measure();
132            if pred(i1) {
133                (None, left, Some(One(right.clone())))
134            } else {
135                (Some(One(left.clone())), right, None)
136            }
137        },
138        Node3(_, ref left, ref middle, ref right) => {
139            let i1 = i + left.measure();
140            if pred(i1) {
141                return (None, left, Some(Two(middle.clone(), right.clone())))
142            }
143            let i2 = i1 + middle.measure();
144            if pred(i2) {
145                (Some(One(left.clone())), middle, Some(One(right.clone())))
146            } else {
147                (Some(Two(left.clone(), middle.clone())), right, None)
148            }
149        }
150    }
151}
152
153impl<T,M> Node<T,M>
154{
155    /// Iterates over the values in the leaves
156    pub fn iter<'a>(&'a self) -> Iter<'a,T,M> {
157        Iter::new(self)
158    }
159}
160
161impl<'a,T,M> IntoIterator for &'a Node<T,M> {
162    type Item = &'a T;
163
164    type IntoIter = Iter<'a, T, M>;
165
166    fn into_iter(self) -> Iter<'a,T,M> {
167        self.iter()
168    }
169}
170
171impl<T,M> Measure<M> for Node<T,M>
172    where T: Measure<M>,
173          M: Copy
174{
175    fn measure(&self) -> M {
176        match self {
177            &Leaf(ref value) => value.measure(),
178            &Node2(measure, _, _) => measure,
179            &Node3(measure, _, _, _) => measure,
180        }
181    }
182}
183
184/// Iterator over the values in the leaves of a 2-3 tree
185#[derive(Debug)]
186pub struct Iter<'a, T:'a, M:'a> {
187    stack: Vec<&'a Node<T,M>>,
188}
189
190impl<'a, T, M> Iter<'a, T, M> {
191    fn new(node: &'a Node<T,M>) -> Iter<'a, T, M> {
192        Iter {
193            stack: vec![node],
194        }
195    }
196    /// An `iter` that is empty (yields no values).
197    ///
198    /// This is helpful in  `finger_tree::Iter`.
199    pub fn empty() -> Iter<'a, T, M> {
200        Iter {
201            stack: vec![],
202        }
203    }
204}
205
206impl<'a, T:'a, M> Iterator for Iter<'a,T,M> {
207    type Item = &'a T;
208    fn next(&mut self) -> Option<&'a T> {
209        let mut node: &'a Node<T,M> = {
210            match self.stack.pop() {
211                None => return None,
212                Some(node) => node
213            }
214        };
215        loop {
216            match *node {
217                Leaf(ref x) => return Some(x),
218                Node2(_, ref left, ref right) => {
219                    self.stack.push(&*right);
220                    node = &*left;
221                },
222                Node3(_, ref left, ref middle, ref right) => {
223                    self.stack.push(&*right);
224                    self.stack.push(&*middle);
225                    node = &*left;
226                }
227            }
228        }
229    }
230}
231
232#[cfg(test)]
233mod test {
234    use super::*;
235    use measure::Measure;
236
237    struct Item<T>(T);
238
239    impl<T> Measure<usize> for Item<T> {
240        fn measure(&self) -> usize {
241            1
242        }
243    }
244
245    #[test]
246    fn test_node_iter() {
247        let tree: Lazy<Node<Item<u32>, usize>> =
248            node2(
249                node3(
250                    leaf(Item(0)),
251                    leaf(Item(1)),
252                    leaf(Item(2))),
253                node2(
254                    leaf(Item(3)),
255                    leaf(Item(4))));
256        let result:Vec<u32> = tree.iter().map(|&Item(x)| x).collect();
257        let expected:Vec<u32> = vec![0,1,2,3,4];
258        assert_eq!(result, expected);
259    }
260
261    #[test]
262    fn test_tree23_measure() {
263        let tree: Lazy<Node<Item<u32>, usize>> =
264            node2(
265                node3(
266                    leaf(Item(0)),
267                    leaf(Item(1)),
268                    leaf(Item(2))),
269                node2(
270                    leaf(Item(3)),
271                    leaf(Item(4))));
272        assert_eq!(tree.measure(), 5);
273    }
274}