1
//#![cfg(feature = "unstable")]
2

            
3
use std::sync::Arc;
4

            
5
use im::vector;
6
use proptest::prelude::*;
7

            
8
use self::Tree::*;
9

            
10
#[derive(Clone, Debug, Eq, PartialEq)]
11
pub enum Tree<T: Sized + Clone + Eq> {
12
    Zero,
13
    One(T),
14
    Many(im::Vector<Tree<T>>),
15
}
16

            
17
// NOTE (niklasdewally): This converts the entire tree into a list. Therefore this is only really
18
// worth it when we use all the children returned. This is what we use this for inside Uniplate.
19
// Because of this, I think a .iter() / IntoIterator for Tree<&T> is a bad idea.
20

            
21
impl<T: Sized + Clone + Eq + 'static> IntoIterator for Tree<T> {
22
    type Item = T;
23

            
24
    type IntoIter = im::vector::ConsumingIter<T>;
25

            
26
    fn into_iter(self) -> Self::IntoIter {
27
        self.list().0.into_iter()
28
    }
29
}
30
impl<T: Sized + Clone + Eq + 'static> Tree<T> {
31
    /// Returns the tree as a list alongside a function to reconstruct the tree from a list.
32
    ///
33
    /// This preserves the structure of the tree.
34
    #[allow(clippy::type_complexity)]
35
    pub fn list(self) -> (im::Vector<T>, Box<dyn Fn(im::Vector<T>) -> Tree<T>>) {
36
        // inspired by the Uniplate Haskell equivalent Data.Generics.Str::strStructure
37
        // https://github.com/ndmitchell/uniplate/blob/master/Data/Generics/Str.hs#L85
38

            
39
        fn flatten<T: Sized + Clone + Eq>(t: Tree<T>, xs: im::Vector<T>) -> im::Vector<T> {
40
            match (t, xs) {
41
                (Zero, xs) => xs,
42
                (One(x), xs) => {
43
                    let mut xs1 = xs.clone();
44
                    xs1.push_back(x);
45
                    xs1
46
                }
47
                (Many(ts), xs) => ts.into_iter().fold(xs, |xs, t| flatten(t, xs)),
48
            }
49
        }
50

            
51
        // Iterate over both the old tree and the new list.
52
        // We use the node types of the old tree to know what node types to use for the new tree.
53
        fn recons<T: Sized + Clone + Eq>(
54
            old_tree: Tree<T>,
55
            xs: im::Vector<T>,
56
        ) -> (Tree<T>, im::Vector<T>) {
57
            #[allow(clippy::unwrap_used)]
58
            match (old_tree, xs) {
59
                (Zero, xs) => (Zero, xs),
60
                (One(_), xs) => {
61
                    let mut xs1 = xs.clone();
62
                    (One(xs1.pop_front().unwrap()), xs1)
63
                }
64
                (Many(ts), xs) => {
65
                    let (ts1, xs1) = ts.into_iter().fold((vector![], xs), |(ts1, xs), t| {
66
                        let (t1, xs1) = recons(t, xs);
67
                        let mut ts2 = ts1.clone();
68
                        ts2.push_back(t1);
69
                        (ts2, xs1)
70
                    });
71
                    (Many(ts1), xs1)
72
                }
73
            }
74
        }
75
        (
76
            flatten(self.clone(), vector![]),
77
            Box::new(move |xs| recons(self.clone(), xs).0),
78
        )
79
    }
80

            
81
    // Perform a map over all elements in the tree.
82
    pub fn map(self, op: Arc<dyn Fn(T) -> T>) -> Tree<T> {
83
        match self {
84
            Zero => Zero,
85
            One(t) => One(op(t)),
86
            Many(ts) => Many(ts.into_iter().map(|t| t.map(op.clone())).collect()),
87
        }
88
    }
89
}
90

            
91
#[allow(dead_code)]
92
// Used by proptest for generating test instances of Tree<i32>.
93
fn proptest_integer_trees() -> impl Strategy<Value = Tree<i32>> {
94
    // https://proptest-rs.github.io/proptest/proptest/tutorial/enums.html
95
    // https://proptest-rs.github.io/proptest/proptest/tutorial/recursive.html
96
    let leaf = prop_oneof![Just(Tree::Zero), any::<i32>().prop_map(Tree::One),];
97

            
98
    leaf.prop_recursive(
99
        10,  // levels deep
100
        512, // Shoot for maximum size of 512 nodes
101
        20,  // We put up to 20 items per collection
102
        |inner| im::proptest::vector(inner.clone(), 0..20).prop_map(Tree::Many),
103
    )
104
}
105

            
106
#[cfg(test)]
107
mod tests {
108
    use std::iter::zip;
109

            
110
    use super::*;
111

            
112
    proptest! {
113
        #[test]
114
        // Is tree.recons() isomorphic?
115
        fn list_is_isomorphic(tree in proptest_integer_trees()) {
116
            let (children,func) = tree.clone().list();
117
            let new_tree = func(children);
118
            prop_assert_eq!(new_tree,tree);
119
        }
120

            
121
        #[test]
122
        fn map_add(tree in proptest_integer_trees(), diff in -100i32..100i32) {
123
            let new_tree = tree.clone().map(Arc::new(move |a| a+diff));
124
            let (old_children,_) = tree.list();
125
            let (new_children,_) = new_tree.list();
126

            
127
            for (old,new) in zip(old_children,new_children) {
128
                prop_assert_eq!(old+diff,new);
129
            }
130
        }
131
    }
132
    #[test]
133
    fn list_preserves_ordering() {
134
        let my_tree: Tree<i32> = Many(vector![
135
            Many(vector![One(0), Zero]),
136
            Many(vector![Many(vector![Zero, One(1), One(2)])]),
137
            One(3),
138
            Zero,
139
            One(4)
140
        ]);
141

            
142
        let flat = my_tree.list().0;
143

            
144
        for i in 0..5 {
145
            assert_eq!(flat[i], i.try_into().unwrap());
146
        }
147
    }
148
}