tenant.openagents/omega

1use crate::{Bounds, Half};
2use std::{
3    cmp,
4    fmt::Debug,
5    ops::{Add, Sub},
6    ptr::NonNull,
7};
8
9/// Maximum children per internal node (R-tree style branching factor).
10/// Higher values = shorter tree = fewer cache misses, but more work per node.
11const MAX_CHILDREN: usize = 12;
12
13/// A spatial tree optimized for finding maximum ordering among intersecting bounds.
14///
15/// This is an R-tree variant specifically designed for the use case of assigning
16/// z-order to overlapping UI elements. Key optimizations:
17/// - Tracks the leaf with global max ordering for O(1) fast-path queries
18/// - Uses higher branching factor (4) for lower tree height
19/// - Aggressive pruning during search based on max_order metadata
20#[derive(Debug)]
21pub(crate) struct BoundsTree<U>
22where
23    U: Clone + Debug + Default + PartialEq,
24{
25    /// All nodes stored contiguously for cache efficiency.
26    nodes: Vec<Node<U>>,
27    /// Index of the root node, if any.
28    root: Option<usize>,
29    /// Index of the leaf with the highest ordering (for fast-path lookups).
30    max_leaf: Option<usize>,
31    /// Reusable stack for tree traversal during insertion.
32    insert_path: Vec<usize>,
33    /// Reusable stack for search operations.
34    search_stack: Vec<NonNull<Node<U>>>,
35}
36
37/// A node in the bounds tree.
38#[derive(Debug, Clone)]
39struct Node<U>
40where
41    U: Clone + Debug + Default + PartialEq,
42{
43    /// Bounding box containing this node and all descendants.
44    bounds: Bounds<U>,
45    /// Maximum ordering value in this subtree.
46    max_order: u32,
47    /// Node-specific data.
48    kind: NodeKind,
49}
50
51#[derive(Debug, Clone)]
52enum NodeKind {
53    /// Leaf node containing actual bounds data.
54    Leaf {
55        /// The ordering assigned to this bounds.
56        order: u32,
57    },
58    /// Internal node with children.
59    Internal {
60        /// Indices of child nodes (2 to MAX_CHILDREN).
61        children: NodeChildren,
62    },
63}
64
65/// Fixed-size array for child indices, avoiding heap allocation.
66#[derive(Debug, Clone)]
67struct NodeChildren {
68    // Keeps an invariant where the max order child is always at the end
69    indices: [usize; MAX_CHILDREN],
70    len: u8,
71}
72
73impl NodeChildren {
74    fn new() -> Self {
75        Self {
76            indices: [0; MAX_CHILDREN],
77            len: 0,
78        }
79    }
80
81    fn push(&mut self, index: usize) {
82        debug_assert!((self.len as usize) < MAX_CHILDREN);
83        self.indices[self.len as usize] = index;
84        self.len += 1;
85    }
86
87    fn len(&self) -> usize {
88        self.len as usize
89    }
90
91    fn as_slice(&self) -> &[usize] {
92        &self.indices[..self.len as usize]
93    }
94}
95
96impl<U> BoundsTree<U>
97where
98    U: Clone
99        + Debug
100        + PartialEq
101        + PartialOrd
102        + Add<U, Output = U>
103        + Sub<Output = U>
104        + Half
105        + Default,
106{
107    /// Clears all nodes from the tree.
108    pub fn clear(&mut self) {
109        self.nodes.clear();
110        self.root = None;
111        self.max_leaf = None;
112        self.insert_path.clear();
113        self.search_stack.clear();
114    }
115
116    /// Inserts bounds into the tree and returns its assigned ordering.
117    ///
118    /// The ordering is one greater than the maximum ordering of any
119    /// existing bounds that intersect with the new bounds.
120    pub fn insert(&mut self, new_bounds: Bounds<U>) -> u32 {
121        // Find maximum ordering among intersecting bounds
122        let max_intersecting = self.find_max_ordering(&new_bounds);
123        let ordering = max_intersecting + 1;
124
125        // Insert the new leaf
126        let new_leaf_idx = self.insert_leaf(new_bounds, ordering);
127
128        // Update max_leaf tracking
129        self.max_leaf = match self.max_leaf {
130            None => Some(new_leaf_idx),
131            Some(old_idx) if self.nodes[old_idx].max_order < ordering => Some(new_leaf_idx),
132            some => some,
133        };
134
135        ordering
136    }
137
138    /// Finds the maximum ordering among all bounds that intersect with the query.
139    fn find_max_ordering(&mut self, query: &Bounds<U>) -> u32 {
140        let Some(root_idx) = self.root else {
141            return 0;
142        };
143
144        // Fast path: check if the max-ordering leaf intersects
145        if let Some(max_idx) = self.max_leaf {
146            let max_node = &self.nodes[max_idx];
147            if query.intersects(&max_node.bounds) {
148                return max_node.max_order;
149            }
150        }
151
152        // Slow path: search the tree
153        self.search_stack.clear();
154        self.search_stack.push(NonNull::from(&self.nodes[root_idx]));
155
156        let mut max_found = 0u32;
157
158        while let Some(node) = self.search_stack.pop() {
159            // SAFETY: `node` is guaranteed to be valid as the `nodes` stack is unmodified in this function
160            // and the `search_stack` only contains pointers from this function call.
161            let node = unsafe { node.as_ref() };
162
163            // Pruning: skip if this subtree can't improve our result
164            if node.max_order <= max_found {
165                continue;
166            }
167
168            // Spatial pruning: skip if bounds don't intersect
169            if !query.intersects(&node.bounds) {
170                continue;
171            }
172
173            match &node.kind {
174                NodeKind::Leaf { order } => {
175                    max_found = cmp::max(max_found, *order);
176                }
177                NodeKind::Internal { children } => {
178                    // Children are maintained with highest max_order at the end.
179                    // Push in forward order to highest (last) is popped first.
180                    self.search_stack.extend(
181                        children
182                            .as_slice()
183                            .iter()
184                            .map(|&child_idx| &self.nodes[child_idx])
185                            .filter(|node| node.max_order > max_found)
186                            .map(NonNull::from),
187                    );
188                }
189            }
190        }
191
192        max_found
193    }
194
195    /// Inserts a leaf node with the given bounds and ordering.
196    /// Returns the index of the new leaf.
197    fn insert_leaf(&mut self, bounds: Bounds<U>, order: u32) -> usize {
198        let new_leaf_idx = self.nodes.len();
199        self.nodes.push(Node {
200            bounds: bounds.clone(),
201            max_order: order,
202            kind: NodeKind::Leaf { order },
203        });
204
205        let Some(root_idx) = self.root else {
206            // Tree is empty, new leaf becomes root
207            self.root = Some(new_leaf_idx);
208            return new_leaf_idx;
209        };
210
211        // If root is a leaf, create internal node with both
212        if matches!(self.nodes[root_idx].kind, NodeKind::Leaf { .. }) {
213            let root_bounds = self.nodes[root_idx].bounds.clone();
214            let root_order = self.nodes[root_idx].max_order;
215
216            let mut children = NodeChildren::new();
217            // Max end invariant
218            if order > root_order {
219                children.push(root_idx);
220                children.push(new_leaf_idx);
221            } else {
222                children.push(new_leaf_idx);
223                children.push(root_idx);
224            }
225
226            let new_root_idx = self.nodes.len();
227            self.nodes.push(Node {
228                bounds: root_bounds.union(&bounds),
229                max_order: cmp::max(root_order, order),
230                kind: NodeKind::Internal { children },
231            });
232            self.root = Some(new_root_idx);
233            return new_leaf_idx;
234        }
235
236        // Descend to find the best internal node to insert into
237        self.insert_path.clear();
238        let mut current_idx = root_idx;
239
240        loop {
241            let current = &self.nodes[current_idx];
242            let NodeKind::Internal { children } = &current.kind else {
243                unreachable!("Should only traverse internal nodes");
244            };
245
246            self.insert_path.push(current_idx);
247
248            // Find the best child to descend into
249            let mut best_child_idx = children.as_slice()[0];
250            let mut best_child_pos = 0;
251            let mut best_cost = bounds
252                .union(&self.nodes[best_child_idx].bounds)
253                .half_perimeter();
254
255            for (pos, &child_idx) in children.as_slice().iter().enumerate().skip(1) {
256                let cost = bounds.union(&self.nodes[child_idx].bounds).half_perimeter();
257                if cost < best_cost {
258                    best_cost = cost;
259                    best_child_idx = child_idx;
260                    best_child_pos = pos;
261                }
262            }
263
264            // Check if best child is a leaf or internal
265            if matches!(self.nodes[best_child_idx].kind, NodeKind::Leaf { .. }) {
266                // Best child is a leaf. Check if current node has room for another child.
267                if children.len() < MAX_CHILDREN {
268                    // Add new leaf directly to this node
269                    let node = &mut self.nodes[current_idx];
270
271                    if let NodeKind::Internal { children } = &mut node.kind {
272                        children.push(new_leaf_idx);
273                        // Swap new leaf only if it has the highest max_order
274                        if order <= node.max_order {
275                            let last = children.len() - 1;
276                            children.indices.swap(last - 1, last);
277                        }
278                    }
279
280                    node.bounds = node.bounds.union(&bounds);
281                    node.max_order = cmp::max(node.max_order, order);
282                    break;
283                } else {
284                    // Node is full, create new internal with [best_leaf, new_leaf]
285                    let sibling_bounds = self.nodes[best_child_idx].bounds.clone();
286                    let sibling_order = self.nodes[best_child_idx].max_order;
287
288                    let mut new_children = NodeChildren::new();
289                    // Max end invariant
290                    if order > sibling_order {
291                        new_children.push(best_child_idx);
292                        new_children.push(new_leaf_idx);
293                    } else {
294                        new_children.push(new_leaf_idx);
295                        new_children.push(best_child_idx);
296                    }
297
298                    let new_internal_idx = self.nodes.len();
299                    let new_internal_max = cmp::max(sibling_order, order);
300                    self.nodes.push(Node {
301                        bounds: sibling_bounds.union(&bounds),
302                        max_order: new_internal_max,
303                        kind: NodeKind::Internal {
304                            children: new_children,
305                        },
306                    });
307
308                    // Replace the leaf with the new internal in parent
309                    let parent = &mut self.nodes[current_idx];
310                    if let NodeKind::Internal { children } = &mut parent.kind {
311                        let children_len = children.len();
312
313                        children.indices[best_child_pos] = new_internal_idx;
314
315                        // If new internal has highest max_order, swap it to the end
316                        // to maintain sorting invariant
317                        if new_internal_max > parent.max_order {
318                            children.indices.swap(best_child_pos, children_len - 1);
319                        }
320                    }
321                    break;
322                }
323            } else {
324                // Best child is internal, continue descent
325                current_idx = best_child_idx;
326            }
327        }
328
329        // Propagate bounds and max_order updates up the tree
330        let mut updated_child_idx = None;
331        for &node_idx in self.insert_path.iter().rev() {
332            let node = &mut self.nodes[node_idx];
333            node.bounds = node.bounds.union(&bounds);
334
335            if node.max_order < order {
336                node.max_order = order;
337
338                // Swap updated child to end (skip first iteration since the invariant is already handled by previous cases)
339                if let Some(child_idx) = updated_child_idx {
340                    if let NodeKind::Internal { children } = &mut node.kind {
341                        if let Some(pos) = children.as_slice().iter().position(|&c| c == child_idx)
342                        {
343                            let last = children.len() - 1;
344                            if pos != last {
345                                children.indices.swap(pos, last);
346                            }
347                        }
348                    }
349                }
350            }
351
352            updated_child_idx = Some(node_idx);
353        }
354
355        new_leaf_idx
356    }
357}
358
359impl<U> Default for BoundsTree<U>
360where
361    U: Clone + Debug + Default + PartialEq,
362{
363    fn default() -> Self {
364        BoundsTree {
365            nodes: Vec::new(),
366            root: None,
367            max_leaf: None,
368            insert_path: Vec::new(),
369            search_stack: Vec::new(),
370        }
371    }
372}
373
374#[cfg(test)]
375mod tests {
376    use super::*;
377    use crate::{Bounds, Point, Size};
378    use rand::{Rng, SeedableRng};
379
380    #[test]
381    fn test_insert() {
382        let mut tree = BoundsTree::<f32>::default();
383        let bounds1 = Bounds {
384            origin: Point { x: 0.0, y: 0.0 },
385            size: Size {
386                width: 10.0,
387                height: 10.0,
388            },
389        };
390        let bounds2 = Bounds {
391            origin: Point { x: 5.0, y: 5.0 },
392            size: Size {
393                width: 10.0,
394                height: 10.0,
395            },
396        };
397        let bounds3 = Bounds {
398            origin: Point { x: 10.0, y: 10.0 },
399            size: Size {
400                width: 10.0,
401                height: 10.0,
402            },
403        };
404
405        // Insert the bounds into the tree and verify the order is correct
406        assert_eq!(tree.insert(bounds1), 1);
407        assert_eq!(tree.insert(bounds2), 2);
408        assert_eq!(tree.insert(bounds3), 3);
409
410        // Insert non-overlapping bounds and verify they can reuse orders
411        let bounds4 = Bounds {
412            origin: Point { x: 20.0, y: 20.0 },
413            size: Size {
414                width: 10.0,
415                height: 10.0,
416            },
417        };
418        let bounds5 = Bounds {
419            origin: Point { x: 40.0, y: 40.0 },
420            size: Size {
421                width: 10.0,
422                height: 10.0,
423            },
424        };
425        let bounds6 = Bounds {
426            origin: Point { x: 25.0, y: 25.0 },
427            size: Size {
428                width: 10.0,
429                height: 10.0,
430            },
431        };
432        assert_eq!(tree.insert(bounds4), 1); // bounds4 does not overlap with bounds1, bounds2, or bounds3
433        assert_eq!(tree.insert(bounds5), 1); // bounds5 does not overlap with any other bounds
434        assert_eq!(tree.insert(bounds6), 2); // bounds6 overlaps with bounds4, so it should have a different order
435    }
436
437    #[test]
438    fn test_random_iterations() {
439        let max_bounds = 100;
440        for seed in 1..=1000 {
441            // let seed = 44;
442            let mut tree = BoundsTree::default();
443            let mut rng = rand::rngs::StdRng::seed_from_u64(seed as u64);
444            let mut expected_quads: Vec<(Bounds<f32>, u32)> = Vec::new();
445
446            // Insert a random number of random AABBs into the tree.
447            let num_bounds = rng.random_range(1..=max_bounds);
448            for _ in 0..num_bounds {
449                let min_x: f32 = rng.random_range(-100.0..100.0);
450                let min_y: f32 = rng.random_range(-100.0..100.0);
451                let width: f32 = rng.random_range(0.0..50.0);
452                let height: f32 = rng.random_range(0.0..50.0);
453                let bounds = Bounds {
454                    origin: Point { x: min_x, y: min_y },
455                    size: Size { width, height },
456                };
457
458                let expected_ordering = expected_quads
459                    .iter()
460                    .filter_map(|quad| quad.0.intersects(&bounds).then_some(quad.1))
461                    .max()
462                    .unwrap_or(0)
463                    + 1;
464                expected_quads.push((bounds, expected_ordering));
465
466                // Insert the AABB into the tree and collect intersections.
467                let actual_ordering = tree.insert(bounds);
468                assert_eq!(actual_ordering, expected_ordering);
469            }
470        }
471    }
472}
473