tenant.openagents/omega

1use std::fmt::Debug;
2
3use ::sum_tree::SumTree;
4use collections::FxHashMap;
5use sum_tree::Bias;
6
7use crate::{FocusHandle, FocusId};
8
9/// Represents a collection of focus handles using the tab-index APIs.
10#[derive(Debug)]
11pub(crate) struct TabStopMap {
12    current_path: TabStopPath,
13    pub(crate) insertion_history: Vec<TabStopOperation>,
14    by_id: FxHashMap<FocusId, TabStopNode>,
15    order: SumTree<TabStopNode>,
16}
17
18#[derive(Debug, Clone)]
19pub enum TabStopOperation {
20    Insert(FocusHandle),
21    Group(TabIndex),
22    GroupEnd,
23}
24
25impl TabStopOperation {
26    fn focus_handle(&self) -> Option<&FocusHandle> {
27        match self {
28            TabStopOperation::Insert(focus_handle) => Some(focus_handle),
29            _ => None,
30        }
31    }
32}
33
34type TabIndex = isize;
35
36#[derive(Debug, Default, PartialEq, Eq, Clone, Ord, PartialOrd)]
37struct TabStopPath(smallvec::SmallVec<[TabIndex; 6]>);
38
39#[derive(Clone, Debug, Default, Eq, PartialEq)]
40struct TabStopNode {
41    /// Path to access the node in the tree
42    /// The final node in the list is a leaf node corresponding to an actual focus handle,
43    /// all other nodes are group nodes
44    path: TabStopPath,
45    /// index into the backing array of nodes. Corresponds to insertion order
46    node_insertion_index: usize,
47
48    /// Whether this node is a tab stop
49    tab_stop: bool,
50}
51
52impl Ord for TabStopNode {
53    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
54        self.path
55            .cmp(&other.path)
56            .then(self.node_insertion_index.cmp(&other.node_insertion_index))
57    }
58}
59
60impl PartialOrd for TabStopNode {
61    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
62        Some(self.cmp(&other))
63    }
64}
65
66impl Default for TabStopMap {
67    fn default() -> Self {
68        Self {
69            current_path: TabStopPath::default(),
70            insertion_history: Vec::new(),
71            by_id: FxHashMap::default(),
72            order: SumTree::new(()),
73        }
74    }
75}
76
77impl TabStopMap {
78    pub fn insert(&mut self, focus_handle: &FocusHandle) {
79        self.insertion_history
80            .push(TabStopOperation::Insert(focus_handle.clone()));
81        let mut path = self.current_path.clone();
82        path.0.push(focus_handle.tab_index);
83        let order = TabStopNode {
84            node_insertion_index: self.insertion_history.len() - 1,
85            tab_stop: focus_handle.tab_stop,
86            path,
87        };
88        self.by_id.insert(focus_handle.id, order.clone());
89        self.order.insert_or_replace(order, ());
90    }
91
92    pub fn begin_group(&mut self, tab_index: isize) {
93        self.insertion_history
94            .push(TabStopOperation::Group(tab_index));
95        self.current_path.0.push(tab_index);
96    }
97
98    pub fn end_group(&mut self) {
99        self.insertion_history.push(TabStopOperation::GroupEnd);
100        self.current_path.0.pop();
101    }
102
103    pub fn clear(&mut self) {
104        *self = Self::default();
105        self.current_path.0.clear();
106        self.insertion_history.clear();
107        self.by_id.clear();
108        self.order = SumTree::new(());
109    }
110
111    pub fn next(&self, focused_id: Option<&FocusId>) -> Option<FocusHandle> {
112        let Some(focused_id) = focused_id else {
113            let first = self.order.first()?;
114            if first.tab_stop {
115                return self.focus_handle_for_order(first);
116            } else {
117                return self
118                    .next_inner(first)
119                    .and_then(|order| self.focus_handle_for_order(order));
120            }
121        };
122
123        let Some(node) = self.tab_node_for_focus_id(focused_id) else {
124            return self.next(None);
125        };
126        let item = self.next_inner(node);
127
128        if let Some(item) = item {
129            self.focus_handle_for_order(&item)
130        } else {
131            self.next(None)
132        }
133    }
134
135    fn next_inner(&self, node: &TabStopNode) -> Option<&TabStopNode> {
136        let mut cursor = self.order.cursor::<TabStopNode>(());
137        cursor.seek(&node, Bias::Left);
138        cursor.next();
139        while let Some(item) = cursor.item()
140            && !item.tab_stop
141        {
142            cursor.next();
143        }
144
145        cursor.item()
146    }
147
148    pub fn prev(&self, focused_id: Option<&FocusId>) -> Option<FocusHandle> {
149        let Some(focused_id) = focused_id else {
150            let last = self.order.last()?;
151            if last.tab_stop {
152                return self.focus_handle_for_order(last);
153            } else {
154                return self
155                    .prev_inner(last)
156                    .and_then(|order| self.focus_handle_for_order(order));
157            }
158        };
159
160        let Some(node) = self.tab_node_for_focus_id(focused_id) else {
161            return self.prev(None);
162        };
163        let item = self.prev_inner(node);
164
165        if let Some(item) = item {
166            self.focus_handle_for_order(&item)
167        } else {
168            self.prev(None)
169        }
170    }
171
172    fn prev_inner(&self, node: &TabStopNode) -> Option<&TabStopNode> {
173        let mut cursor = self.order.cursor::<TabStopNode>(());
174        cursor.seek(&node, Bias::Left);
175        cursor.prev();
176        while let Some(item) = cursor.item()
177            && !item.tab_stop
178        {
179            cursor.prev();
180        }
181
182        cursor.item()
183    }
184
185    pub fn replay(&mut self, nodes: &[TabStopOperation]) {
186        for node in nodes {
187            match node {
188                TabStopOperation::Insert(focus_handle) => self.insert(focus_handle),
189                TabStopOperation::Group(tab_index) => self.begin_group(*tab_index),
190                TabStopOperation::GroupEnd => self.end_group(),
191            }
192        }
193    }
194
195    pub fn paint_index(&self) -> usize {
196        self.insertion_history.len()
197    }
198
199    pub(crate) fn tab_stop_count(&self) -> usize {
200        self.by_id.values().filter(|node| node.tab_stop).count()
201    }
202
203    fn focus_handle_for_order(&self, order: &TabStopNode) -> Option<FocusHandle> {
204        let handle = self.insertion_history[order.node_insertion_index].focus_handle();
205        debug_assert!(
206            handle.is_some(),
207            "The order node did not correspond to an element, this is a GPUI bug"
208        );
209        handle.cloned()
210    }
211
212    fn tab_node_for_focus_id(&self, focused_id: &FocusId) -> Option<&TabStopNode> {
213        let Some(order) = self.by_id.get(focused_id) else {
214            return None;
215        };
216        Some(order)
217    }
218}
219
220mod sum_tree_impl {
221    use sum_tree::SeekTarget;
222
223    use crate::tab_stop::{TabStopNode, TabStopPath};
224
225    #[derive(Clone, Debug)]
226    pub struct TabStopOrderNodeSummary {
227        max_index: usize,
228        max_path: TabStopPath,
229        pub tab_stops: usize,
230    }
231
232    pub type TabStopCount = usize;
233
234    impl sum_tree::ContextLessSummary for TabStopOrderNodeSummary {
235        fn zero() -> Self {
236            TabStopOrderNodeSummary {
237                max_index: 0,
238                max_path: TabStopPath::default(),
239                tab_stops: 0,
240            }
241        }
242
243        fn add_summary(&mut self, summary: &Self) {
244            self.max_index = summary.max_index;
245            self.max_path = summary.max_path.clone();
246            self.tab_stops += summary.tab_stops;
247        }
248    }
249
250    impl sum_tree::KeyedItem for TabStopNode {
251        type Key = Self;
252
253        fn key(&self) -> Self::Key {
254            self.clone()
255        }
256    }
257
258    impl sum_tree::Item for TabStopNode {
259        type Summary = TabStopOrderNodeSummary;
260
261        fn summary(&self, _cx: <Self::Summary as sum_tree::Summary>::Context<'_>) -> Self::Summary {
262            TabStopOrderNodeSummary {
263                max_index: self.node_insertion_index,
264                max_path: self.path.clone(),
265                tab_stops: if self.tab_stop { 1 } else { 0 },
266            }
267        }
268    }
269
270    impl<'a> sum_tree::Dimension<'a, TabStopOrderNodeSummary> for TabStopCount {
271        fn zero(_: <TabStopOrderNodeSummary as sum_tree::Summary>::Context<'_>) -> Self {
272            0
273        }
274
275        fn add_summary(
276            &mut self,
277            summary: &'a TabStopOrderNodeSummary,
278            _: <TabStopOrderNodeSummary as sum_tree::Summary>::Context<'_>,
279        ) {
280            *self += summary.tab_stops;
281        }
282    }
283
284    impl<'a> sum_tree::Dimension<'a, TabStopOrderNodeSummary> for TabStopNode {
285        fn zero(_: <TabStopOrderNodeSummary as sum_tree::Summary>::Context<'_>) -> Self {
286            TabStopNode::default()
287        }
288
289        fn add_summary(
290            &mut self,
291            summary: &'a TabStopOrderNodeSummary,
292            _: <TabStopOrderNodeSummary as sum_tree::Summary>::Context<'_>,
293        ) {
294            self.node_insertion_index = summary.max_index;
295            self.path = summary.max_path.clone();
296        }
297    }
298
299    impl<'a, 'b> SeekTarget<'a, TabStopOrderNodeSummary, TabStopNode> for &'b TabStopNode {
300        fn cmp(
301            &self,
302            cursor_location: &TabStopNode,
303            _: <TabStopOrderNodeSummary as sum_tree::Summary>::Context<'_>,
304        ) -> std::cmp::Ordering {
305            Iterator::cmp(self.path.0.iter(), cursor_location.path.0.iter()).then(
306                <usize as Ord>::cmp(
307                    &self.node_insertion_index,
308                    &cursor_location.node_insertion_index,
309                ),
310            )
311        }
312    }
313}
314
315#[cfg(test)]
316mod tests {
317    use itertools::Itertools as _;
318
319    use crate::{FocusHandle, FocusId, FocusMap, TabStopMap};
320    use std::sync::Arc;
321
322    #[test]
323    fn test_tab_handles() {
324        let focus_map = Arc::new(FocusMap::default());
325        let mut tab_index_map = TabStopMap::default();
326
327        let focus_handles = [
328            FocusHandle::new(&focus_map).tab_stop(true).tab_index(0),
329            FocusHandle::new(&focus_map).tab_stop(true).tab_index(1),
330            FocusHandle::new(&focus_map).tab_stop(true).tab_index(1),
331            FocusHandle::new(&focus_map),
332            FocusHandle::new(&focus_map).tab_index(2),
333            FocusHandle::new(&focus_map).tab_stop(true).tab_index(0),
334            FocusHandle::new(&focus_map).tab_stop(true).tab_index(2),
335        ];
336
337        for handle in focus_handles.iter() {
338            tab_index_map.insert(handle);
339        }
340        let expected = [
341            focus_handles[0].clone(),
342            focus_handles[5].clone(),
343            focus_handles[1].clone(),
344            focus_handles[2].clone(),
345            focus_handles[6].clone(),
346        ];
347
348        let mut prev = None;
349        let mut found = vec![];
350        for _ in 0..expected.len() {
351            let handle = tab_index_map.next(prev.as_ref()).unwrap();
352            prev = Some(handle.id);
353            found.push(handle.id);
354        }
355
356        assert_eq!(
357            found,
358            expected.iter().map(|handle| handle.id).collect::<Vec<_>>()
359        );
360
361        // Select first tab index if no handle is currently focused.
362        assert_eq!(tab_index_map.next(None), Some(expected[0].clone()));
363        // Select last tab index if no handle is currently focused.
364        assert_eq!(tab_index_map.prev(None), expected.last().cloned(),);
365
366        assert_eq!(
367            tab_index_map.next(Some(&expected[0].id)),
368            Some(expected[1].clone())
369        );
370        assert_eq!(
371            tab_index_map.next(Some(&expected[1].id)),
372            Some(expected[2].clone())
373        );
374        assert_eq!(
375            tab_index_map.next(Some(&expected[2].id)),
376            Some(expected[3].clone())
377        );
378        assert_eq!(
379            tab_index_map.next(Some(&expected[3].id)),
380            Some(expected[4].clone())
381        );
382        assert_eq!(
383            tab_index_map.next(Some(&expected[4].id)),
384            Some(expected[0].clone())
385        );
386
387        // prev
388        assert_eq!(tab_index_map.prev(None), Some(expected[4].clone()));
389        assert_eq!(
390            tab_index_map.prev(Some(&expected[0].id)),
391            Some(expected[4].clone())
392        );
393        assert_eq!(
394            tab_index_map.prev(Some(&expected[1].id)),
395            Some(expected[0].clone())
396        );
397        assert_eq!(
398            tab_index_map.prev(Some(&expected[2].id)),
399            Some(expected[1].clone())
400        );
401        assert_eq!(
402            tab_index_map.prev(Some(&expected[3].id)),
403            Some(expected[2].clone())
404        );
405        assert_eq!(
406            tab_index_map.prev(Some(&expected[4].id)),
407            Some(expected[3].clone())
408        );
409    }
410
411    #[test]
412    fn test_tab_non_stop_filtering() {
413        let focus_map = Arc::new(FocusMap::default());
414        let mut tab_index_map = TabStopMap::default();
415
416        // Check that we can query next from a non-stop tab
417        let tab_non_stop_1 = FocusHandle::new(&focus_map).tab_stop(false).tab_index(1);
418        let tab_stop_2 = FocusHandle::new(&focus_map).tab_stop(true).tab_index(2);
419        tab_index_map.insert(&tab_non_stop_1);
420        tab_index_map.insert(&tab_stop_2);
421        let result = tab_index_map.next(Some(&tab_non_stop_1.id)).unwrap();
422        assert_eq!(result.id, tab_stop_2.id);
423
424        // Check that we skip over non-stop tabs
425        let tab_stop_0 = FocusHandle::new(&focus_map).tab_stop(true).tab_index(0);
426        let tab_non_stop_0 = FocusHandle::new(&focus_map).tab_stop(false).tab_index(0);
427        tab_index_map.insert(&tab_stop_0);
428        tab_index_map.insert(&tab_non_stop_0);
429        let result = tab_index_map.next(Some(&tab_stop_0.id)).unwrap();
430        assert_eq!(result.id, tab_stop_2.id);
431    }
432
433    #[must_use]
434    struct TabStopMapTest {
435        tab_map: TabStopMap,
436        focus_map: Arc<FocusMap>,
437        expected: Vec<(usize, FocusId)>,
438    }
439
440    impl TabStopMapTest {
441        #[must_use]
442        fn new() -> Self {
443            Self {
444                tab_map: TabStopMap::default(),
445                focus_map: Arc::new(FocusMap::default()),
446                expected: Vec::default(),
447            }
448        }
449
450        #[must_use]
451        fn tab_non_stop(mut self, index: isize) -> Self {
452            let handle = FocusHandle::new(&self.focus_map)
453                .tab_stop(false)
454                .tab_index(index);
455            self.tab_map.insert(&handle);
456            self
457        }
458
459        #[must_use]
460        fn tab_stop(mut self, index: isize, expected: usize) -> Self {
461            let handle = FocusHandle::new(&self.focus_map)
462                .tab_stop(true)
463                .tab_index(index);
464            self.tab_map.insert(&handle);
465            self.expected.push((expected, handle.id));
466            self.expected.sort_by_key(|(expected, _)| *expected);
467            self
468        }
469
470        #[must_use]
471        fn tab_group(mut self, tab_index: isize, children: impl FnOnce(Self) -> Self) -> Self {
472            self.tab_map.begin_group(tab_index);
473            self = children(self);
474            self.tab_map.end_group();
475            self
476        }
477
478        fn traverse_tab_map(
479            &self,
480            traverse: impl Fn(&TabStopMap, Option<&FocusId>) -> Option<FocusHandle>,
481        ) -> Vec<FocusId> {
482            let mut last_focus_id = None;
483            let mut found = vec![];
484            for _ in 0..self.expected.len() {
485                let handle = traverse(&self.tab_map, last_focus_id.as_ref()).unwrap();
486                last_focus_id = Some(handle.id);
487                found.push(handle.id);
488            }
489            found
490        }
491
492        fn assert(self) {
493            let mut expected = self.expected.iter().map(|(_, id)| *id).collect_vec();
494
495            // Check next order
496            let forward_found = self.traverse_tab_map(|tab_map, prev| tab_map.next(prev));
497            assert_eq!(forward_found, expected);
498
499            // Test overflow. Last to first
500            assert_eq!(
501                self.tab_map
502                    .next(forward_found.last())
503                    .map(|handle| handle.id),
504                expected.first().cloned()
505            );
506
507            // Check previous order
508            let reversed_found = self.traverse_tab_map(|tab_map, prev| tab_map.prev(prev));
509            expected.reverse();
510            assert_eq!(reversed_found, expected);
511
512            // Test overflow. First to last
513            assert_eq!(
514                self.tab_map
515                    .prev(reversed_found.last())
516                    .map(|handle| handle.id),
517                expected.first().cloned(),
518            );
519        }
520    }
521
522    #[test]
523    fn test_with_disabled_tab_stop() {
524        TabStopMapTest::new()
525            .tab_stop(0, 0)
526            .tab_non_stop(1)
527            .tab_stop(2, 1)
528            .tab_stop(3, 2)
529            .assert();
530    }
531
532    #[test]
533    fn test_with_multiple_disabled_tab_stops() {
534        TabStopMapTest::new()
535            .tab_non_stop(0)
536            .tab_stop(1, 0)
537            .tab_non_stop(3)
538            .tab_stop(3, 1)
539            .tab_non_stop(4)
540            .assert();
541    }
542
543    #[test]
544    fn test_tab_group_functionality() {
545        TabStopMapTest::new()
546            .tab_stop(0, 0)
547            .tab_stop(0, 1)
548            .tab_group(2, |t| t.tab_stop(0, 2).tab_stop(1, 3))
549            .tab_stop(3, 4)
550            .tab_stop(4, 5)
551            .assert()
552    }
553
554    #[test]
555    fn test_sibling_groups() {
556        TabStopMapTest::new()
557            .tab_stop(0, 0)
558            .tab_stop(1, 1)
559            .tab_group(2, |test| test.tab_stop(0, 2).tab_stop(1, 3))
560            .tab_stop(3, 4)
561            .tab_stop(4, 5)
562            .tab_group(6, |test| test.tab_stop(0, 6).tab_stop(1, 7))
563            .tab_stop(7, 8)
564            .tab_stop(8, 9)
565            .assert();
566    }
567
568    #[test]
569    fn test_nested_group() {
570        TabStopMapTest::new()
571            .tab_stop(0, 0)
572            .tab_stop(1, 1)
573            .tab_group(2, |t| {
574                t.tab_group(0, |t| t.tab_stop(0, 2).tab_stop(1, 3))
575                    .tab_stop(1, 4)
576            })
577            .tab_stop(3, 5)
578            .tab_stop(4, 6)
579            .assert();
580    }
581
582    #[test]
583    fn test_sibling_nested_groups() {
584        TabStopMapTest::new()
585            .tab_stop(0, 0)
586            .tab_stop(1, 1)
587            .tab_group(2, |builder| {
588                builder
589                    .tab_stop(0, 2)
590                    .tab_stop(2, 5)
591                    .tab_group(1, |builder| builder.tab_stop(0, 3).tab_stop(1, 4))
592                    .tab_group(3, |builder| builder.tab_stop(0, 6).tab_stop(1, 7))
593            })
594            .tab_stop(3, 8)
595            .tab_stop(4, 9)
596            .assert();
597    }
598
599    #[test]
600    fn test_sibling_nested_groups_out_of_order() {
601        TabStopMapTest::new()
602            .tab_stop(9, 9)
603            .tab_stop(8, 8)
604            .tab_group(7, |builder| {
605                builder
606                    .tab_stop(0, 2)
607                    .tab_stop(2, 5)
608                    .tab_group(3, |builder| builder.tab_stop(1, 7).tab_stop(0, 6))
609                    .tab_group(1, |builder| builder.tab_stop(0, 3).tab_stop(1, 4))
610            })
611            .tab_stop(3, 0)
612            .tab_stop(4, 1)
613            .assert();
614    }
615}
616