1use std::hash::Hash;
2
3use rustc_hash::{FxBuildHasher, FxHashMap, FxHashSet};
4use smallvec::SmallVec;
5
6use crate::context::{TreeExpansionState, TreeMatchState};
7use crate::model::{
8 TreeChildren, TreeFilter, TreeFilterConfig, TreeModel, TreeQuery, TreeRevision,
9 TreeRootVisibility, TreeSort,
10};
11use crate::traversal::TreePostorder;
12
13pub struct OccurrencePath<Id> {
14 root_parent: Option<Id>,
15 ids: SmallVec<[Id; 16]>,
16}
17
18impl<Id> OccurrencePath<Id> {
19 pub fn len(&self) -> usize {
20 self.ids.len()
21 }
22}
23
24#[derive(Clone, Copy, Debug, PartialEq, Eq)]
26pub struct ProjectedNode<Id> {
27 id: Id,
28 parent: Option<Id>,
29 parent_index: Option<usize>,
30 level: usize,
31 is_last_sibling: bool,
32 visible_child_count: usize,
33 expansion: TreeExpansionState,
34 match_state: TreeMatchState,
35}
36
37impl<Id: Copy> ProjectedNode<Id> {
38 #[must_use]
39 pub const fn id(self) -> Id {
40 self.id
41 }
42
43 #[must_use]
44 pub const fn parent(self) -> Option<Id> {
45 self.parent
46 }
47
48 #[must_use]
53 pub const fn parent_index(self) -> Option<usize> {
54 self.parent_index
55 }
56
57 #[must_use]
58 pub const fn level(self) -> usize {
59 self.level
60 }
61
62 #[must_use]
63 pub const fn is_last_sibling(self) -> bool {
64 self.is_last_sibling
65 }
66
67 #[must_use]
68 pub const fn visible_child_count(self) -> usize {
69 self.visible_child_count
70 }
71
72 #[must_use]
73 pub const fn expansion(self) -> TreeExpansionState {
74 self.expansion
75 }
76
77 #[must_use]
78 pub const fn match_state(self) -> TreeMatchState {
79 self.match_state
80 }
81}
82
83#[derive(Clone, Copy, Debug, PartialEq, Eq)]
84struct ProjectionStamp {
85 model: TreeRevision,
86 filter: PolicyStamp,
87 sort: PolicyStamp,
88 expansion: TreeRevision,
89 filter_config: TreeFilterConfig,
90 root_visibility: TreeRootVisibility,
91}
92
93#[derive(Clone, Copy, Debug, PartialEq, Eq)]
94struct PolicyStamp {
95 revision: TreeRevision,
96 generation: TreeRevision,
97}
98
99impl PolicyStamp {
100 const fn new(revision: TreeRevision, generation: TreeRevision) -> Self {
101 Self {
102 revision,
103 generation,
104 }
105 }
106}
107
108pub struct TreeProjection<Id> {
110 nodes: Vec<ProjectedNode<Id>>,
111 index: FxHashMap<Id, usize>,
112 filter_memo: FxHashMap<Id, bool>,
113 direct_matches: FxHashSet<Id>,
114 stamp: Option<ProjectionStamp>,
115}
116
117impl<Id: Copy + Eq + Hash> TreeProjection<Id> {
118 pub(crate) fn with_capacity(capacity: usize) -> Self {
119 Self {
120 nodes: Vec::with_capacity(capacity),
121 index: FxHashMap::with_capacity_and_hasher(capacity, FxBuildHasher),
122 filter_memo: FxHashMap::with_capacity_and_hasher(capacity, FxBuildHasher),
123 direct_matches: FxHashSet::with_capacity_and_hasher(capacity, FxBuildHasher),
124 stamp: None,
125 }
126 }
127
128 #[must_use]
130 pub fn nodes(&self) -> &[ProjectedNode<Id>] {
131 &self.nodes
132 }
133
134 #[must_use]
136 pub const fn len(&self) -> usize {
137 self.nodes.len()
138 }
139
140 #[must_use]
142 pub const fn is_empty(&self) -> bool {
143 self.nodes.is_empty()
144 }
145
146 #[must_use]
148 pub fn index_of(&self, id: Id) -> Option<usize> {
149 self.index.get(&id).copied()
150 }
151
152 #[must_use]
154 pub fn get_by_id(&self, id: Id) -> Option<ProjectedNode<Id>> {
155 self.index_of(id)
156 .and_then(|index| self.nodes.get(index))
157 .copied()
158 }
159
160 pub(crate) fn is_current<T, F, S>(
161 &self,
162 model: &T,
163 query: &TreeQuery<F, S>,
164 expansion: TreeRevision,
165 ) -> bool
166 where
167 T: TreeModel<Id = Id>,
168 {
169 self.stamp == Some(Self::stamp(model, query, expansion))
170 }
171
172 pub(crate) fn rebuild<T, F, S, E>(
173 &mut self,
174 model: &T,
175 query: &TreeQuery<F, S>,
176 expansion_revision: TreeRevision,
177 is_expanded: E,
178 ) where
179 T: TreeModel<Id = Id>,
180 F: TreeFilter<T>,
181 S: TreeSort<T>,
182 E: Fn(Option<Id>, Id) -> bool,
183 {
184 self.nodes.clear();
185 self.index.clear();
186 self.reserve(model.size_hint());
187
188 let filtering = matches!(query.filter_config(), TreeFilterConfig::Enabled { .. });
189 if filtering {
190 self.compute_filter_matches(model, query.filter());
191 } else {
192 self.filter_memo.clear();
193 self.direct_matches.clear();
194 }
195
196 let mut roots: SmallVec<[Id; 8]> = model.roots().collect();
197 Self::sort_ids(model, query.sort(), &mut roots);
198 let mut stack = Vec::with_capacity(model.size_hint().min(1024).max(roots.len()));
199
200 match query.root_visibility() {
201 TreeRootVisibility::Visible => {
202 Self::push_children(&mut stack, &roots, None, None, 0);
203 }
204 TreeRootVisibility::Hidden => {
205 for root in roots.iter().rev().copied() {
206 let mut children =
207 self.visible_children(query, model.children(root).loaded_slice());
208 Self::sort_ids(model, query.sort(), &mut children);
209 Self::push_children(&mut stack, &children, Some(root), None, 0);
210 }
211 }
212 }
213
214 while let Some(frame) = stack.pop() {
215 if filtering && !self.filter_memo.get(&frame.id).copied().unwrap_or(false) {
216 continue;
217 }
218
219 let children_state = model.children(frame.id);
220 let mut visible_children = match children_state {
221 TreeChildren::Loaded(children) => self.visible_children(query, children),
222 TreeChildren::Leaf | TreeChildren::Unloaded | TreeChildren::Loading => {
223 SmallVec::new()
224 }
225 };
226 Self::sort_ids(model, query.sort(), &mut visible_children);
227
228 let expansion = match children_state {
229 TreeChildren::Leaf => TreeExpansionState::Leaf,
230 TreeChildren::Unloaded => TreeExpansionState::Unloaded,
231 TreeChildren::Loading => TreeExpansionState::Loading,
232 TreeChildren::Loaded(_) if visible_children.is_empty() => TreeExpansionState::Leaf,
233 TreeChildren::Loaded(_) => match query.filter_config() {
234 TreeFilterConfig::Enabled { auto_expand: true } => {
235 TreeExpansionState::ForcedByFilter
236 }
237 TreeFilterConfig::Disabled
238 | TreeFilterConfig::Enabled { auto_expand: false } => {
239 if is_expanded(frame.parent, frame.id) {
240 TreeExpansionState::Expanded
241 } else {
242 TreeExpansionState::Collapsed
243 }
244 }
245 },
246 };
247 let match_state = if !filtering {
248 TreeMatchState::Unfiltered
249 } else if self.direct_matches.contains(&frame.id) {
250 TreeMatchState::Direct
251 } else {
252 TreeMatchState::Ancestor
253 };
254
255 let index = self.nodes.len();
256 self.nodes.push(ProjectedNode {
257 id: frame.id,
258 parent: frame.parent,
259 parent_index: frame.parent_index,
260 level: frame.level,
261 is_last_sibling: frame.is_last_sibling,
262 visible_child_count: visible_children.len(),
263 expansion,
264 match_state,
265 });
266 self.index.entry(frame.id).or_insert(index);
267
268 if expansion.is_expanded() {
269 Self::push_children(
270 &mut stack,
271 &visible_children,
272 Some(frame.id),
273 Some(index),
274 frame.level.saturating_add(1),
275 );
276 }
277 }
278
279 self.stamp = Some(Self::stamp(model, query, expansion_revision));
280 }
281
282 fn stamp<T, F, S>(
283 model: &T,
284 query: &TreeQuery<F, S>,
285 expansion: TreeRevision,
286 ) -> ProjectionStamp
287 where
288 T: TreeModel<Id = Id>,
289 {
290 ProjectionStamp {
291 model: model.revision(),
292 filter: PolicyStamp::new(query.filter_revision(), query.filter_generation()),
293 sort: PolicyStamp::new(query.sort_revision(), query.sort_generation()),
294 expansion,
295 filter_config: query.filter_config(),
296 root_visibility: query.root_visibility(),
297 }
298 }
299
300 fn reserve(&mut self, hint: usize) {
301 if hint == 0 {
302 return;
303 }
304 reserve_to(&mut self.nodes, hint);
305 reserve_map_to(&mut self.index, hint);
306 reserve_map_to(&mut self.filter_memo, hint);
307 let extra = hint.saturating_sub(self.direct_matches.len());
308 self.direct_matches.reserve(extra);
309 }
310
311 fn visible_children<F, S>(
312 &self,
313 query: &TreeQuery<F, S>,
314 children: &[Id],
315 ) -> SmallVec<[Id; 8]> {
316 match query.filter_config() {
317 TreeFilterConfig::Disabled => children.iter().copied().collect(),
318 TreeFilterConfig::Enabled { .. } => children
319 .iter()
320 .copied()
321 .filter(|child| self.filter_memo.get(child).copied().unwrap_or(false))
322 .collect(),
323 }
324 }
325
326 fn compute_filter_matches<T, F>(&mut self, model: &T, filter: &F)
327 where
328 T: TreeModel<Id = Id>,
329 F: TreeFilter<T>,
330 {
331 self.filter_memo.clear();
332 self.direct_matches.clear();
333 for node in TreePostorder::forest(model) {
334 let direct = filter.is_match(model, node.id);
335 if direct {
336 self.direct_matches.insert(node.id);
337 }
338 let descendant = node
339 .children
340 .iter()
341 .any(|child| self.filter_memo.get(child).copied().unwrap_or(false));
342 self.filter_memo.insert(node.id, direct || descendant);
343 }
344 }
345
346 fn sort_ids<T, S>(model: &T, sort: &S, ids: &mut [Id])
347 where
348 T: TreeModel<Id = Id>,
349 S: TreeSort<T>,
350 {
351 if sort.is_enabled() {
352 ids.sort_by(|left, right| sort.compare(model, *left, *right));
353 }
354 }
355
356 pub(crate) fn occurrence_path(&self, index: usize) -> Option<OccurrencePath<Id>> {
357 let mut ids = SmallVec::new();
358 let mut cursor = Some(index);
359 let mut root_parent = None;
360 while let Some(index) = cursor {
361 let node = self.nodes.get(index)?;
362 ids.push(node.id);
363 cursor = node.parent_index;
364 if cursor.is_none() {
365 root_parent = node.parent;
366 }
367 }
368 ids.reverse();
369 Some(OccurrencePath { root_parent, ids })
370 }
371
372 pub(crate) fn index_of_path(&self, path: &OccurrencePath<Id>) -> Option<usize> {
373 self.index_of_path_prefix(path, path.len())
374 }
375
376 pub(crate) fn index_of_path_prefix(
377 &self,
378 path: &OccurrencePath<Id>,
379 end: usize,
380 ) -> Option<usize> {
381 let ids = path.ids.get(..end)?;
382 let (&id, _) = ids.split_last()?;
383 let first = self.index_of(id)?;
384 if self.path_matches(first, path.root_parent, ids) {
385 return Some(first);
386 }
387 self.nodes[first + 1..]
388 .iter()
389 .enumerate()
390 .filter(|(_, node)| node.id == id)
391 .find_map(|(offset, _)| {
392 let index = first + 1 + offset;
393 self.path_matches(index, path.root_parent, ids)
394 .then_some(index)
395 })
396 }
397
398 fn path_matches(&self, index: usize, root_parent: Option<Id>, ids: &[Id]) -> bool {
399 let mut cursor = Some(index);
400 let mut actual_root_parent = None;
401 for &expected_id in ids.iter().rev() {
402 let Some(node) = cursor.and_then(|index| self.nodes.get(index)) else {
403 return false;
404 };
405 if node.id != expected_id {
406 return false;
407 }
408 cursor = node.parent_index;
409 actual_root_parent = node.parent;
410 }
411 cursor.is_none() && actual_root_parent == root_parent
412 }
413
414 fn push_children(
415 stack: &mut Vec<ProjectionFrame<Id>>,
416 children: &[Id],
417 parent: Option<Id>,
418 parent_index: Option<usize>,
419 level: usize,
420 ) {
421 let last = children.len().saturating_sub(1);
422 stack.extend(
423 children
424 .iter()
425 .copied()
426 .enumerate()
427 .rev()
428 .map(|(index, id)| ProjectionFrame {
429 id,
430 parent,
431 parent_index,
432 level,
433 is_last_sibling: index == last,
434 }),
435 );
436 }
437}
438
439struct ProjectionFrame<Id> {
440 id: Id,
441 parent: Option<Id>,
442 parent_index: Option<usize>,
443 level: usize,
444 is_last_sibling: bool,
445}
446
447fn reserve_to<T>(values: &mut Vec<T>, capacity: usize) {
448 values.reserve(capacity.saturating_sub(values.len()));
449}
450
451fn reserve_map_to<K: Eq + Hash, V>(values: &mut FxHashMap<K, V>, capacity: usize) {
452 values.reserve(capacity.saturating_sub(values.len()));
453}