1use std::error::Error;
2use std::fmt::{self, Display, Formatter};
3use std::hash::Hash;
4
5use smallvec::SmallVec;
6
7use crate::model::{TreeChildren, TreeModel, TreeRevision};
8
9#[derive(Clone, Copy, Debug, PartialEq, Eq)]
11pub enum IndexedTreeError {
12 InvalidRoot(usize),
13 DuplicateRoot(usize),
14 MissingRoot(usize),
15 InvalidChild { parent: usize, child: usize },
16 MultipleParents(usize),
17 RootHasParent(usize),
18 Cycle,
19}
20
21impl Display for IndexedTreeError {
22 fn fmt(&self, formatter: &mut Formatter<'_>) -> fmt::Result {
23 match self {
24 Self::InvalidRoot(root) => write!(formatter, "invalid root index: {root}"),
25 Self::DuplicateRoot(root) => write!(formatter, "duplicate root index: {root}"),
26 Self::MissingRoot(root) => write!(formatter, "top-level node is not a root: {root}"),
27 Self::InvalidChild { parent, child } => {
28 write!(
29 formatter,
30 "invalid child index {child} under parent {parent}"
31 )
32 }
33 Self::MultipleParents(node) => write!(formatter, "node has multiple parents: {node}"),
34 Self::RootHasParent(root) => write!(formatter, "root has a parent: {root}"),
35 Self::Cycle => formatter.write_str("indexed tree contains a cycle"),
36 }
37 }
38}
39
40impl Error for IndexedTreeError {}
41
42pub struct TreeModelRef<'a, Id, C> {
44 roots: &'a [Id],
45 children: C,
46 revision: TreeRevision,
47 size_hint: usize,
48}
49
50impl<'a, Id, C> TreeModelRef<'a, Id, C> {
51 #[must_use]
53 pub const fn new(roots: &'a [Id], children: C, revision: TreeRevision) -> Self {
54 Self {
55 roots,
56 children,
57 revision,
58 size_hint: 0,
59 }
60 }
61
62 #[must_use]
64 pub const fn with_size_hint(mut self, size_hint: usize) -> Self {
65 self.size_hint = size_hint;
66 self
67 }
68}
69
70impl<'a, Id, C> TreeModel for TreeModelRef<'a, Id, C>
71where
72 Id: Copy + Eq + Hash,
73 C: Fn(Id) -> TreeChildren<'a, Id>,
74{
75 type Id = Id;
76
77 fn roots(&self) -> impl Iterator<Item = Self::Id> + '_ {
78 self.roots.iter().copied()
79 }
80
81 fn children(&self, id: Self::Id) -> TreeChildren<'_, Self::Id> {
82 (self.children)(id)
83 }
84
85 fn revision(&self) -> TreeRevision {
86 self.revision
87 }
88
89 fn size_hint(&self) -> usize {
90 self.size_hint
91 }
92}
93
94pub struct IndexedTree<'a, C = Vec<usize>>
96where
97 C: AsRef<[usize]>,
98{
99 roots: SmallVec<[usize; 1]>,
100 children: &'a [C],
101 revision: TreeRevision,
102}
103
104impl<'a, C> IndexedTree<'a, C>
105where
106 C: AsRef<[usize]>,
107{
108 pub fn new(
115 roots: impl IntoIterator<Item = usize>,
116 children: &'a [C],
117 revision: TreeRevision,
118 ) -> Result<Self, IndexedTreeError> {
119 let roots: SmallVec<[usize; 1]> = roots.into_iter().collect();
120 let mut indegree = vec![0_usize; children.len()];
121 let mut root_seen = vec![false; children.len()];
122 for root in roots.iter().copied() {
123 let Some(seen) = root_seen.get_mut(root) else {
124 return Err(IndexedTreeError::InvalidRoot(root));
125 };
126 if *seen {
127 return Err(IndexedTreeError::DuplicateRoot(root));
128 }
129 *seen = true;
130 }
131
132 for (parent, node_children) in children.iter().enumerate() {
133 for &child in node_children.as_ref() {
134 let Some(value) = indegree.get_mut(child) else {
135 return Err(IndexedTreeError::InvalidChild { parent, child });
136 };
137 *value = value.saturating_add(1);
138 if *value > 1 {
139 return Err(IndexedTreeError::MultipleParents(child));
140 }
141 }
142 }
143 if let Some(root) = roots.iter().copied().find(|root| indegree[*root] != 0) {
144 return Err(IndexedTreeError::RootHasParent(root));
145 }
146 if let Some(root) = indegree
147 .iter()
148 .enumerate()
149 .find_map(|(id, degree)| (*degree == 0 && !root_seen[id]).then_some(id))
150 {
151 return Err(IndexedTreeError::MissingRoot(root));
152 }
153
154 let mut queue: Vec<_> = indegree
155 .iter()
156 .enumerate()
157 .filter_map(|(id, degree)| (*degree == 0).then_some(id))
158 .collect();
159 let mut processed = 0;
160 while let Some(id) = queue.pop() {
161 processed += 1;
162 for &child in children[id].as_ref() {
163 indegree[child] -= 1;
164 if indegree[child] == 0 {
165 queue.push(child);
166 }
167 }
168 }
169 if processed != children.len() {
170 return Err(IndexedTreeError::Cycle);
171 }
172
173 Ok(Self {
174 roots,
175 children,
176 revision,
177 })
178 }
179}
180
181impl<C> TreeModel for IndexedTree<'_, C>
182where
183 C: AsRef<[usize]>,
184{
185 type Id = usize;
186
187 fn roots(&self) -> impl Iterator<Item = Self::Id> + '_ {
188 self.roots.iter().copied()
189 }
190
191 fn children(&self, id: Self::Id) -> TreeChildren<'_, Self::Id> {
192 TreeChildren::loaded(self.children[id].as_ref())
193 }
194
195 fn revision(&self) -> TreeRevision {
196 self.revision
197 }
198
199 fn size_hint(&self) -> usize {
200 self.children.len()
201 }
202}
203
204#[cfg(test)]
205mod tests {
206 use super::*;
207
208 #[test]
209 fn indexed_tree_rejects_shared_nodes_and_cycles() {
210 let shared = vec![vec![2], vec![2], vec![]];
211 assert!(matches!(
212 IndexedTree::new([0, 1], &shared, TreeRevision::INITIAL),
213 Err(IndexedTreeError::MultipleParents(2))
214 ));
215
216 let cycle = vec![vec![1], vec![0]];
217 assert!(matches!(
218 IndexedTree::new([], &cycle, TreeRevision::INITIAL),
219 Err(IndexedTreeError::Cycle)
220 ));
221
222 let generic: [&[usize]; 2] = [&[1], &[]];
223 let tree = IndexedTree::new([0], &generic, TreeRevision::INITIAL)
224 .expect("slice-backed adjacency list is valid");
225 assert_eq!(tree.children(0).loaded_slice(), &[1]);
226 }
227}