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#[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
20pub fn leaf<T,M>(v: T) -> Lazy<Node<T,M>> {
22 strict(Leaf(v))
23}
24
25pub 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
37pub 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 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#[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 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}