1
use streaming_iterator::StreamingIterator;
2

            
3
use crate::Data;
4
use crate::DataRef;
5
use crate::Ldd;
6
use crate::Storage;
7
use crate::Value;
8

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

            
17
// Returns an iterator over all vectors contained in the given LDD.
18
4329
pub fn iter<'a>(storage: &'a Storage, ldd: &Ldd) -> Iter<'a> {
19
4329
    if ldd == storage.empty_vector() || ldd == storage.empty_set() {
20
        Iter {
21
            storage,
22
            vector: Vec::new(),
23
            current: Vec::new(),
24
            stack: Vec::new(),
25
            finished: false,
26
        }
27
    } else {
28
4329
        Iter {
29
4329
            storage,
30
4329
            vector: Vec::new(),
31
4329
            current: Vec::new(),
32
4329
            stack: vec![ldd.clone()],
33
4329
            finished: false,
34
4329
        }
35
    }
36
4329
}
37

            
38
// Returns an iterator over all nodes in the given LDD. Visits each node only if the predicate holds.
39
2000
pub fn iter_nodes<'a, P>(storage: &'a Storage, ldd: &Ldd, filter: P) -> IterNode<'a, P>
40
2000
where
41
2000
    P: Fn(&Ldd) -> bool,
42
{
43
2000
    let mut stack = Vec::new();
44

            
45
2000
    if ldd != storage.empty_set() {
46
2000
        stack.push((ldd.clone(), false));
47
2000
    }
48

            
49
2000
    IterNode {
50
2000
        storage,
51
2000
        stack,
52
2000
        predicate: filter,
53
2000
    }
54
2000
}
55

            
56
pub struct IterNode<'a, P>
57
where
58
    P: Fn(&Ldd) -> bool,
59
{
60
    storage: &'a Storage,
61
    stack: Vec<(Ldd, bool)>,
62
    predicate: P,
63
}
64

            
65
impl<P> Iterator for IterNode<'_, P>
66
where
67
    P: Fn(&Ldd) -> bool,
68
{
69
    type Item = (Ldd, Data);
70

            
71
451313
    fn next(&mut self) -> Option<Self::Item> {
72
900626
        while let Some((current, visited)) = self.stack.pop() {
73
898626
            let data = self.storage.get(&current);
74

            
75
898626
            if visited {
76
449313
                return Some((current, data));
77
449313
            }
78

            
79
449313
            let Data(_, down, right) = &data;
80

            
81
            // Next time we can actually process the current node.
82
449313
            self.stack.push((current.clone(), true));
83

            
84
            // Add unvisited children to stack
85
449313
            if (self.predicate)(down) {
86
385313
                self.stack.push((down.clone(), false));
87
385313
            }
88

            
89
449313
            if (self.predicate)(right) {
90
62000
                self.stack.push((right.clone(), false));
91
387313
            }
92
        }
93

            
94
2000
        None
95
451313
    }
96
}
97

            
98
pub struct IterRight<'a> {
99
    storage: &'a Storage,
100
    current: Ldd,
101
}
102

            
103
impl Iterator for IterRight<'_> {
104
    type Item = Data;
105

            
106
193154
    fn next(&mut self) -> Option<Self::Item> {
107
193154
        if self.current == *self.storage.empty_set() {
108
964
            None
109
        } else {
110
            // Progress to the right LDD.
111
192190
            let Data(value, down, right) = self.storage.get(&self.current);
112
192190
            self.current = right.clone();
113
192190
            Some(Data(value, down, right))
114
        }
115
193154
    }
116
}
117

            
118
pub struct Iter<'a> {
119
    storage: &'a Storage,
120
    /// Stores the values of the returned vector.
121
    vector: Vec<Value>,
122
    /// Stores the current path in the LDD.
123
    current: Vec<Value>,
124
    /// Stores the stack for the depth-first search (only non 'true' or 'false' nodes)
125
    stack: Vec<Ldd>,   
126
    /// Indicates whether the iteration is finished.
127
    finished: bool, 
128
}
129

            
130
impl StreamingIterator for Iter<'_> {
131
    type Item = Vec<Value>;
132
    
133
64783
    fn advance(&mut self) {
134
        // Find the next vector by going down the chain.
135
        loop {
136
319770
            if let Some(current) = self.stack.last() {
137
315444
                let DataRef(value, down, _) = self.storage.get_ref(current);
138
315444
                self.vector.push(value);
139
315444
                if down == *self.storage.empty_vector() {
140
                    // Reserve sufficient space in current
141
60457
                    if self.current.len() < self.vector.len() {
142
4329
                        self.current.resize(self.vector.len(), Value::default());
143
56128
                    }
144

            
145
60457
                    self.current.copy_from_slice(&self.vector);
146
60457
                    break; // Stop iteration.
147
254987
                } else {
148
254987
                    self.stack.push(self.storage.protect(&down));
149
254987
                }
150
            } else {
151
                // No more elements to iterate.
152
4326
                self.finished = true;
153
4326
                return;
154
            }
155
        }
156

            
157
        // Go up the chain to find the next right sibling that is not 'false'.
158
319773
        while let Some(current) = self.stack.pop() {
159
315444
            self.vector.pop();
160
315444
            let DataRef(_, _, right) = self.storage.get_ref(&current);
161

            
162
315444
            if right != *self.storage.empty_set() {
163
56128
                self.stack.push(self.storage.protect(&right)); // This is the first right sibling.
164
56128
                break;
165
259316
            }
166
        }
167
64783
    }
168
    
169
64783
    fn get(&self) -> Option<&Self::Item> {
170
64783
        if self.finished {
171
4326
            None
172
        } else {
173
60457
            Some(&self.current)
174
        }
175
64783
    }
176
}
177

            
178
#[cfg(test)]
179
mod tests {
180

            
181
    use super::*;
182

            
183
    use merc_utilities::random_test;
184

            
185
    use crate::test_utility::from_iter;
186
    use crate::test_utility::random_vector_set;
187

            
188
    // Test the iterator implementation.
189
    #[test]
190
    #[cfg_attr(miri, ignore)]
191
1
    fn test_random_iter() {
192
100
        random_test(100, |rng| {
193
100
            let mut storage = Storage::new();
194

            
195
100
            let set = random_vector_set(rng, 32, 10, 10);
196
100
            let ldd = from_iter(&mut storage, set.iter());
197

            
198
100
            assert!(
199
100
                iter(&storage, &ldd).count() == set.len(),
200
                "Number of iterations does not match the number of elements in the set."
201
            );
202

            
203
100
            let mut iter = iter(&storage, &ldd);
204
3300
            while let Some(vector) = iter.next() {
205
3200
                assert!(set.contains(vector), "Found element not in the set.");
206
            }
207
100
        })
208
1
    }
209
}