1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
use alloc::{collections::BTreeMap, vec::Vec};
#[cfg(test)]
mod tests;
pub trait Node: Sized {
type ChildrenIter: Iterator<Item = Self>;
fn take_children(&mut self) -> Self::ChildrenIter;
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct CollapsedTree<T> {
pub elems: Vec<T>,
pub mapping: BTreeMap<usize, usize>,
}
#[derive(Copy, Clone, Debug, PartialEq, Eq)]
pub enum CollapseTreeOrd {
TopDown,
BottomUp,
}
pub fn collapse_tree<T, InIter>(input: InIter, order: CollapseTreeOrd) -> CollapsedTree<T>
where
T: Node,
InIter: Iterator<Item = T>,
{
let mut ret: CollapsedTree<T> = CollapsedTree {
elems: Vec::with_capacity({
let s_h = input.size_hint();
s_h.1.unwrap_or(s_h.0)
}),
mapping: BTreeMap::new(),
};
let is_topdown = order == CollapseTreeOrd::TopDown;
for mut i in input {
let CollapsedTree {
elems: subs_elems,
mapping: subs_mapping,
} = collapse_tree(i.take_children(), order);
let mut i = Some(i);
if is_topdown {
ret.elems.push(i.take().unwrap());
}
let (subs_cnt, subs_start_id) = (subs_elems.len(), ret.elems.len());
if subs_cnt != 0 {
ret.elems.extend(subs_elems.into_iter());
let cur_id = if is_topdown {
subs_start_id - 1
} else {
ret.elems.len()
};
for i in 0..subs_cnt {
ret.mapping.insert(
i + subs_start_id,
subs_mapping
.get(&i)
.map(|parent| *parent + subs_start_id)
.unwrap_or(cur_id),
);
}
}
if !is_topdown {
ret.elems.push(i.take().unwrap());
}
}
ret.elems.shrink_to_fit();
ret
}