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
119368
pub fn iter_right<'a>(storage: &'a Storage, ldd: &Ldd) -> IterRight<'a> {
11
119368
    IterRight {
12
119368
        storage,
13
119368
        current: ldd.clone(),
14
119368
    }
15
119368
}
16

            
17
// Returns an iterator over all vectors contained in the given LDD.
18
532084
pub fn iter<'a>(storage: &'a Storage, ldd: &Ldd) -> Iter<'a> {
19
532084
    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
532084
        Iter {
29
532084
            storage,
30
532084
            vector: Vec::new(),
31
532084
            current: Vec::new(),
32
532084
            stack: vec![ldd.clone()],
33
532084
            finished: false,
34
532084
        }
35
    }
36
532084
}
37

            
38
// Calls `f` with `&mut storage` and each vector in the LDD.  Unlike `iter`, the closure receives a
39
// mutable storage reference so it can insert nodes without a separate collection step.
40
100
pub fn for_each_mut<F>(storage: &mut Storage, ldd: &Ldd, mut f: F)
41
100
where
42
100
    F: FnMut(&mut Storage, &[Value]),
43
{
44
100
    if ldd == storage.empty_vector() || ldd == storage.empty_set() {
45
        return;
46
100
    }
47

            
48
100
    let mut vector: Vec<Value> = Vec::new();
49
100
    let mut stack: Vec<Ldd> = vec![ldd.clone()];
50

            
51
    loop {
52
        // Descend to the next leaf (node whose down == empty_vector).
53
        loop {
54
29249
            let (value, down) = {
55
29249
                let DataRef(value, down, _) = storage.get_ref(stack.last().unwrap());
56
29249
                let is_leaf = down == *storage.empty_vector();
57
29249
                let down_protected = (!is_leaf).then(|| storage.protect(&down));
58
29249
                (value, down_protected)
59
            };
60
29249
            vector.push(value);
61
29249
            match down {
62
26049
                Some(down_ldd) => stack.push(down_ldd),
63
3200
                None => break, // reached leaf
64
            }
65
        }
66

            
67
        // Storage is no longer borrowed — the closure may mutate it freely.
68
3200
        f(storage, &vector);
69

            
70
        // Ascend to find the next right sibling.
71
29349
        while let Some(current) = stack.pop() {
72
29249
            vector.pop();
73
29249
            let right = {
74
29249
                let DataRef(_, _, right) = storage.get_ref(&current);
75
29249
                if right != *storage.empty_set() {
76
3100
                    Some(storage.protect(&right))
77
                } else {
78
26149
                    None
79
                }
80
            };
81
29249
            if let Some(right_ldd) = right {
82
3100
                stack.push(right_ldd);
83
3100
                break;
84
26149
            }
85
        }
86

            
87
3200
        if stack.is_empty() {
88
100
            break;
89
3100
        }
90
    }
91
100
}
92

            
93
// Returns an iterator over all nodes in the given LDD. Visits each node only if the predicate holds.
94
2000
pub fn iter_nodes<'a, P>(storage: &'a Storage, ldd: &Ldd, filter: P) -> IterNode<'a, P>
95
2000
where
96
2000
    P: Fn(&Ldd) -> bool,
97
{
98
2000
    let mut stack = Vec::new();
99

            
100
2000
    if ldd != storage.empty_set() {
101
2000
        stack.push((ldd.clone(), false));
102
2000
    }
103

            
104
2000
    IterNode {
105
2000
        storage,
106
2000
        stack,
107
2000
        predicate: filter,
108
2000
    }
109
2000
}
110

            
111
pub struct IterNode<'a, P>
112
where
113
    P: Fn(&Ldd) -> bool,
114
{
115
    storage: &'a Storage,
116
    stack: Vec<(Ldd, bool)>,
117
    predicate: P,
118
}
119

            
120
impl<P> Iterator for IterNode<'_, P>
121
where
122
    P: Fn(&Ldd) -> bool,
123
{
124
    type Item = (Ldd, Data);
125

            
126
451322
    fn next(&mut self) -> Option<Self::Item> {
127
900644
        while let Some((current, visited)) = self.stack.pop() {
128
898644
            let data = self.storage.get(&current);
129

            
130
898644
            if visited {
131
449322
                return Some((current, data));
132
449322
            }
133

            
134
449322
            let Data(_, down, right) = &data;
135

            
136
            // Next time we can actually process the current node.
137
449322
            self.stack.push((current.clone(), true));
138

            
139
            // Add unvisited children to stack
140
449322
            if (self.predicate)(down) {
141
385322
                self.stack.push((down.clone(), false));
142
385322
            }
143

            
144
449322
            if (self.predicate)(right) {
145
62000
                self.stack.push((right.clone(), false));
146
387322
            }
147
        }
148

            
149
2000
        None
150
451322
    }
151
}
152

            
153
pub struct IterRight<'a> {
154
    storage: &'a Storage,
155
    current: Ldd,
156
}
157

            
158
impl Iterator for IterRight<'_> {
159
    type Item = Data;
160

            
161
194373
    fn next(&mut self) -> Option<Self::Item> {
162
194373
        if self.current == *self.storage.empty_set() {
163
983
            None
164
        } else {
165
            // Progress to the right LDD.
166
193390
            let Data(value, down, right) = self.storage.get(&self.current);
167
193390
            self.current = right.clone();
168
193390
            Some(Data(value, down, right))
169
        }
170
194373
    }
171
}
172

            
173
pub struct Iter<'a> {
174
    storage: &'a Storage,
175
    /// Stores the values of the returned vector.
176
    vector: Vec<Value>,
177
    /// Stores the current path in the LDD.
178
    current: Vec<Value>,
179
    /// Stores the stack for the depth-first search (only non 'true' or 'false' nodes)
180
    stack: Vec<Ldd>,
181
    /// Indicates whether the iteration is finished.
182
    finished: bool,
183
}
184

            
185
impl StreamingIterator for Iter<'_> {
186
    type Item = Vec<Value>;
187

            
188
5961686
    fn advance(&mut self) {
189
        // Find the next vector by going down the chain.
190
        loop {
191
42979782
            if let Some(current) = self.stack.last() {
192
42448301
                let DataRef(value, down, _) = self.storage.get_ref(current);
193
42448301
                self.vector.push(value);
194
42448301
                if down == *self.storage.empty_vector() {
195
                    // Reserve sufficient space in current
196
5430205
                    if self.current.len() < self.vector.len() {
197
532084
                        self.current.resize(self.vector.len(), Value::default());
198
4898121
                    }
199

            
200
5430205
                    self.current.copy_from_slice(&self.vector);
201
5430205
                    break; // Stop iteration.
202
37018096
                } else {
203
37018096
                    self.stack.push(self.storage.protect(&down));
204
37018096
                }
205
            } else {
206
                // No more elements to iterate.
207
531481
                self.finished = true;
208
531481
                return;
209
            }
210
        }
211

            
212
        // Go up the chain to find the next right sibling that is not 'false'.
213
42980385
        while let Some(current) = self.stack.pop() {
214
42448301
            self.vector.pop();
215
42448301
            let DataRef(_, _, right) = self.storage.get_ref(&current);
216

            
217
42448301
            if right != *self.storage.empty_set() {
218
4898121
                self.stack.push(self.storage.protect(&right)); // This is the first right sibling.
219
4898121
                break;
220
37550180
            }
221
        }
222
5961686
    }
223

            
224
5961686
    fn get(&self) -> Option<&Self::Item> {
225
5961686
        if self.finished { None } else { Some(&self.current) }
226
5961686
    }
227
}
228

            
229
#[cfg(test)]
230
mod tests {
231

            
232
    use super::*;
233

            
234
    use merc_utilities::random_test;
235

            
236
    use crate::test_utility::from_iter;
237
    use crate::test_utility::random_vector_set;
238

            
239
    // Test that iter and for_each_mut yield the same set of cubes.
240
    #[test]
241
    #[cfg_attr(miri, ignore)]
242
1
    fn test_random_iter_for_each_mut() {
243
100
        random_test(100, |rng| {
244
100
            let mut storage = Storage::new();
245

            
246
100
            let set = random_vector_set(rng, 32, 10, 10);
247
100
            let ldd = from_iter(&mut storage, set.iter());
248

            
249
            // Collect via iter.
250
100
            let mut iter_result: Vec<Vec<Value>> = Vec::new();
251
100
            let mut it = iter(&storage, &ldd);
252
3300
            while let Some(vector) = it.next() {
253
3200
                iter_result.push(vector.to_vec());
254
3200
            }
255

            
256
            // Collect via for_each_mut.
257
100
            let mut for_each_result: Vec<Vec<Value>> = Vec::new();
258
3200
            for_each_mut(&mut storage, &ldd, |_storage, vector| {
259
3200
                for_each_result.push(vector.to_vec());
260
3200
            });
261

            
262
100
            for_each_result.sort();
263
100
            let original_len = for_each_result.len();
264
100
            for_each_result.dedup();
265
100
            assert_eq!(
266
                original_len,
267
100
                for_each_result.len(),
268
                "for_each_mut returned duplicate vectors."
269
            );
270

            
271
100
            assert_eq!(
272
                iter_result, for_each_result,
273
                "iter and for_each_mut yielded different sequences of cubes."
274
            );
275
100
        })
276
1
    }
277

            
278
    // Test the iterator implementation.
279
    #[test]
280
    #[cfg_attr(miri, ignore)]
281
1
    fn test_random_iter() {
282
100
        random_test(100, |rng| {
283
100
            let mut storage = Storage::new();
284

            
285
100
            let set = random_vector_set(rng, 32, 10, 10);
286
100
            let ldd = from_iter(&mut storage, set.iter());
287

            
288
100
            assert!(
289
100
                iter(&storage, &ldd).count() == set.len(),
290
                "Number of iterations does not match the number of elements in the set."
291
            );
292

            
293
100
            let mut results: Vec<Vec<Value>> = Vec::new();
294
100
            let mut it = iter(&storage, &ldd);
295
3300
            while let Some(vector) = it.next() {
296
3200
                assert!(set.contains(vector), "Found element not in the set.");
297
3200
                results.push(vector.to_vec());
298
            }
299

            
300
100
            results.sort();
301
100
            let original_len = results.len();
302
100
            results.dedup();
303
100
            assert_eq!(original_len, results.len(), "iter returned duplicate vectors.");
304
100
        })
305
1
    }
306
}