1
use crate::Data;
2
use crate::Ldd;
3
use crate::Storage;
4
use crate::Value;
5

            
6
// Returns an iterator over all right siblings of the given LDD.
7
119210
pub fn iter_right<'a>(storage: &'a Storage, ldd: &Ldd) -> IterRight<'a> {
8
119210
    IterRight {
9
119210
        storage,
10
119210
        current: ldd.clone(),
11
119210
    }
12
119210
}
13

            
14
// Returns an iterator over all vectors contained in the given LDD.
15
1600
pub fn iter<'a>(storage: &'a Storage, ldd: &Ldd) -> Iter<'a> {
16
1600
    if ldd == storage.empty_set() {
17
        Iter {
18
            storage,
19
            vector: Vec::new(),
20
            stack: Vec::new(),
21
        }
22
    } else {
23
1600
        Iter {
24
1600
            storage,
25
1600
            vector: Vec::new(),
26
1600
            stack: vec![ldd.clone()],
27
1600
        }
28
    }
29
1600
}
30

            
31
// Returns an iterator over all nodes in the given LDD. Visits each node only if the predicate holds.
32
2000
pub fn iter_nodes<'a, P>(storage: &'a Storage, ldd: &Ldd, filter: P) -> IterNode<'a, P>
33
2000
where
34
2000
    P: Fn(&Ldd) -> bool,
35
{
36
2000
    let mut stack = Vec::new();
37

            
38
2000
    if ldd != storage.empty_set() {
39
2000
        stack.push((ldd.clone(), false));
40
2000
    }
41

            
42
2000
    IterNode {
43
2000
        storage,
44
2000
        stack,
45
2000
        predicate: filter,
46
2000
    }
47
2000
}
48

            
49
pub struct IterNode<'a, P>
50
where
51
    P: Fn(&Ldd) -> bool,
52
{
53
    storage: &'a Storage,
54
    stack: Vec<(Ldd, bool)>,
55
    predicate: P,
56
}
57

            
58
impl<P> Iterator for IterNode<'_, P>
59
where
60
    P: Fn(&Ldd) -> bool,
61
{
62
    type Item = (Ldd, Data);
63

            
64
451338
    fn next(&mut self) -> Option<Self::Item> {
65
900676
        while let Some((current, visited)) = self.stack.pop() {
66
898676
            let data = self.storage.get(&current);
67

            
68
898676
            if visited {
69
449338
                return Some((current, data));
70
449338
            }
71

            
72
449338
            let Data(_, down, right) = &data;
73

            
74
            // Next time we can actually process the current node.
75
449338
            self.stack.push((current.clone(), true));
76

            
77
            // Add unvisited children to stack
78
449338
            if (self.predicate)(down) {
79
385338
                self.stack.push((down.clone(), false));
80
385338
            }
81

            
82
449338
            if (self.predicate)(right) {
83
62000
                self.stack.push((right.clone(), false));
84
387338
            }
85
        }
86

            
87
2000
        None
88
451338
    }
89
}
90

            
91
pub struct IterRight<'a> {
92
    storage: &'a Storage,
93
    current: Ldd,
94
}
95

            
96
impl Iterator for IterRight<'_> {
97
    type Item = Data;
98

            
99
193282
    fn next(&mut self) -> Option<Self::Item> {
100
193282
        if self.current == *self.storage.empty_set() {
101
934
            None
102
        } else {
103
            // Progress to the right LDD.
104
192348
            let Data(value, down, right) = self.storage.get(&self.current);
105
192348
            self.current = right.clone();
106
192348
            Some(Data(value, down, right))
107
        }
108
193282
    }
109
}
110

            
111
pub struct Iter<'a> {
112
    storage: &'a Storage,
113
    vector: Vec<Value>, // Stores the values of the returned vector.
114
    stack: Vec<Ldd>,    // Stores the stack for the depth-first search (only non 'true' or 'false' nodes)
115
}
116

            
117
impl Iterator for Iter<'_> {
118
    type Item = Vec<Value>;
119

            
120
27838
    fn next(&mut self) -> Option<Self::Item> {
121
        // Find the next vector by going down the chain.
122
        let vector: Vec<Value>;
123
        loop {
124
219593
            let current = self.stack.last()?;
125

            
126
217993
            let Data(value, down, _) = self.storage.get(current);
127
217993
            self.vector.push(value);
128
217993
            if down == *self.storage.empty_vector() {
129
26238
                vector = self.vector.clone();
130
26238
                break; // Stop iteration.
131
191755
            } else {
132
191755
                self.stack.push(down.clone());
133
191755
            }
134
        }
135

            
136
        // Go up the chain to find the next right sibling that is not 'false'.
137
219593
        while let Some(current) = self.stack.pop() {
138
217993
            self.vector.pop();
139
217993
            let Data(_, _, right) = self.storage.get(&current);
140

            
141
217993
            if right != *self.storage.empty_set() {
142
24638
                self.stack.push(right); // This is the first right sibling.
143
24638
                break;
144
193355
            }
145
        }
146

            
147
26238
        Some(vector)
148
27838
    }
149
}
150

            
151
#[cfg(test)]
152
mod tests {
153

            
154
    use super::*;
155

            
156
    use merc_utilities::random_test;
157

            
158
    use crate::test_utility::from_iter;
159
    use crate::test_utility::random_vector_set;
160

            
161
    // Test the iterator implementation.
162
    #[test]
163
    #[cfg_attr(miri, ignore)]
164
1
    fn test_random_iter() {
165
100
        random_test(100, |rng| {
166
100
            let mut storage = Storage::new();
167

            
168
100
            let set = random_vector_set(rng, 32, 10, 10);
169
100
            let ldd = from_iter(&mut storage, set.iter());
170

            
171
100
            assert!(
172
100
                iter(&storage, &ldd).count() == set.len(),
173
                "Number of iterations does not match the number of elements in the set."
174
            );
175

            
176
3200
            for vector in iter(&storage, &ldd) {
177
3200
                assert!(set.contains(&vector), "Found element not in the set.");
178
            }
179
100
        })
180
1
    }
181
}