1use crate::VirtualScrollbar;
2use crate::gpui_compat::element_id;
3use crate::tree::TreeNode;
4use gpui::{
5 App, Context, Entity, IntoElement, ListAlignment, ListState, MouseButton, Pixels, Render,
6 SharedString, Window, div, list, prelude::*, px,
7};
8use liora_core::Config;
9use liora_icons::Icon;
10use liora_icons_lucide::IconName;
11use std::collections::HashSet;
12use std::sync::Arc;
13
14type NodeCallback = dyn Fn(SharedString, &mut Window, &mut App) + 'static;
15
16#[derive(Debug, Clone, PartialEq, Eq)]
18pub struct VirtualTreeItem {
19 pub id: SharedString,
20 pub label: SharedString,
21 pub depth: u32,
22 pub has_children: bool,
23}
24
25pub struct VirtualizedTree {
31 data: Vec<TreeNode>,
32 expanded_keys: HashSet<SharedString>,
33 selected_keys: HashSet<SharedString>,
34 flattened: Vec<VirtualTreeItem>,
35 list_state: ListState,
36 multiple: bool,
37 indent: Pixels,
38 row_height: Pixels,
39 height: Pixels,
40 overdraw: Pixels,
41 show_checkbox: bool,
42 on_node_click: Option<Arc<NodeCallback>>,
43}
44
45impl VirtualizedTree {
46 pub fn new(data: Vec<TreeNode>, _cx: &mut Context<Self>) -> Self {
47 let flattened = flatten_visible(&data, &HashSet::new());
48 let overdraw = px(640.0);
49 let list_state = ListState::new(flattened.len(), ListAlignment::Top, overdraw);
50 Self {
51 data,
52 expanded_keys: HashSet::new(),
53 selected_keys: HashSet::new(),
54 flattened,
55 list_state,
56 multiple: false,
57 indent: px(18.0),
58 row_height: px(34.0),
59 height: px(360.0),
60 overdraw,
61 show_checkbox: false,
62 on_node_click: None,
63 }
64 }
65
66 pub fn entity(data: Vec<TreeNode>, cx: &mut App) -> Entity<Self> {
67 cx.new(|cx| Self::new(data, cx))
68 }
69
70 pub fn height(mut self, height: impl Into<Pixels>) -> Self {
71 self.height = height.into();
72 self
73 }
74
75 pub fn row_height(mut self, height: impl Into<Pixels>) -> Self {
76 self.row_height = height.into();
77 self.list_state.reset(self.flattened.len());
78 self
79 }
80
81 pub fn indent(mut self, indent: impl Into<Pixels>) -> Self {
82 self.indent = indent.into();
83 self
84 }
85
86 pub fn overdraw(mut self, overdraw: impl Into<Pixels>) -> Self {
87 self.overdraw = overdraw.into();
88 self.rebuild_list_state();
89 self
90 }
91
92 pub fn multiple(mut self, multiple: bool) -> Self {
93 self.multiple = multiple;
94 self
95 }
96
97 pub fn show_checkbox(mut self, show: bool) -> Self {
98 self.show_checkbox = show;
99 self
100 }
101
102 pub fn default_expanded_keys(mut self, keys: impl IntoIterator<Item = SharedString>) -> Self {
103 self.expanded_keys = keys.into_iter().collect();
104 self.rebuild_flattened();
105 self
106 }
107
108 pub fn default_selected_keys(mut self, keys: impl IntoIterator<Item = SharedString>) -> Self {
109 self.selected_keys = keys.into_iter().collect();
110 self
111 }
112
113 pub fn expand_all(mut self) -> Self {
114 let mut keys = HashSet::new();
115 collect_parent_keys(&self.data, &mut keys);
116 self.expanded_keys = keys;
117 self.rebuild_flattened();
118 self
119 }
120
121 pub fn on_node_click(
122 mut self,
123 callback: impl Fn(SharedString, &mut Window, &mut App) + 'static,
124 ) -> Self {
125 self.on_node_click = Some(Arc::new(callback));
126 self
127 }
128
129 pub fn visible_len(&self) -> usize {
130 self.flattened.len()
131 }
132
133 pub fn is_expanded(&self, id: &SharedString) -> bool {
134 self.expanded_keys.contains(id)
135 }
136
137 pub fn is_selected(&self, id: &SharedString) -> bool {
138 self.selected_keys.contains(id)
139 }
140
141 pub fn list_state(&self) -> ListState {
142 self.list_state.clone()
143 }
144
145 fn rebuild_flattened(&mut self) {
146 let next = flatten_visible(&self.data, &self.expanded_keys);
147 let count_changed = next.len() != self.flattened.len();
148 self.flattened = next;
149 if count_changed {
150 self.rebuild_list_state();
151 } else {
152 self.list_state.reset(self.flattened.len());
153 }
154 }
155
156 fn rebuild_list_state(&mut self) {
157 self.list_state = ListState::new(self.flattened.len(), ListAlignment::Top, self.overdraw);
158 }
159
160 fn toggle_expand(&mut self, id: SharedString, cx: &mut Context<Self>) {
161 if self.expanded_keys.contains(&id) {
162 self.expanded_keys.remove(&id);
163 } else {
164 self.expanded_keys.insert(id);
165 }
166 self.rebuild_flattened();
167 cx.notify();
168 }
169
170 fn select_node(&mut self, id: SharedString, window: &mut Window, cx: &mut Context<Self>) {
171 if self.multiple {
172 if self.selected_keys.contains(&id) {
173 self.selected_keys.remove(&id);
174 } else {
175 self.selected_keys.insert(id.clone());
176 }
177 } else {
178 self.selected_keys.clear();
179 self.selected_keys.insert(id.clone());
180 }
181
182 if let Some(callback) = self.on_node_click.clone() {
183 callback(id, window, cx);
184 }
185 cx.notify();
186 }
187
188 fn click_node(
189 &mut self,
190 id: SharedString,
191 has_children: bool,
192 window: &mut Window,
193 cx: &mut Context<Self>,
194 ) {
195 if has_children {
196 if self.expanded_keys.contains(&id) {
197 self.expanded_keys.remove(&id);
198 } else {
199 self.expanded_keys.insert(id.clone());
200 }
201 self.rebuild_flattened();
202 }
203 self.select_node(id, window, cx);
204 }
205}
206
207impl Render for VirtualizedTree {
208 fn render(&mut self, _window: &mut Window, cx: &mut Context<Self>) -> impl IntoElement {
209 let theme = cx.global::<Config>().theme.clone();
210 let flattened = self.flattened.clone();
211 let expanded_keys = self.expanded_keys.clone();
212 let selected_keys = self.selected_keys.clone();
213 let indent = self.indent;
214 let row_height = self.row_height;
215 let show_checkbox = self.show_checkbox;
216 let entity = cx.entity().clone();
217 let list_state = self.list_state.clone();
218
219 div()
220 .relative()
221 .w_full()
222 .h(self.height)
223 .overflow_hidden()
224 .rounded(px(theme.radius.md))
225 .border_1()
226 .border_color(theme.neutral.border)
227 .bg(theme.neutral.card)
228 .child(
229 list(list_state.clone(), move |index, _window, _cx| {
230 let Some(item) = flattened.get(index).cloned() else {
231 return div().into_any_element();
232 };
233 let id = item.id.clone();
234 let is_expanded = expanded_keys.contains(&id);
235 let is_selected = selected_keys.contains(&id);
236 let has_children = item.has_children;
237 let padding_left = px(f32::from(indent) * item.depth as f32);
238 let expand_entity = entity.clone();
239 let click_entity = entity.clone();
240 let expand_id = id.clone();
241 let click_id = id.clone();
242
243 div()
244 .id(element_id(format!("virtual-tree-row-{}", id)))
245 .cursor_pointer()
246 .flex()
247 .flex_row()
248 .items_center()
249 .gap_1()
250 .w_full()
251 .min_h(row_height)
252 .pl(padding_left)
253 .pr_4()
254 .text_color(if is_selected {
255 theme.primary.base
256 } else {
257 theme.neutral.text_1
258 })
259 .bg(if is_selected {
260 theme.primary.base.opacity(0.1)
261 } else {
262 gpui::transparent_black()
263 })
264 .hover(|s| s.bg(theme.neutral.hover))
265 .child(
266 div()
267 .flex()
268 .items_center()
269 .justify_center()
270 .w(px(22.0))
271 .when(has_children, |s| {
272 s.on_mouse_down(MouseButton::Left, move |_, _, cx| {
273 expand_entity.update(cx, |tree, cx| {
274 tree.toggle_expand(expand_id.clone(), cx);
275 });
276 cx.stop_propagation();
277 })
278 .child(
279 Icon::new(if is_expanded {
280 IconName::ChevronDown
281 } else {
282 IconName::ChevronRight
283 })
284 .size(px(14.0))
285 .color(theme.neutral.text_3),
286 )
287 }),
288 )
289 .when(show_checkbox, |s| {
290 s.child(
291 Icon::new(if is_selected {
292 IconName::Check
293 } else {
294 IconName::Square
295 })
296 .size(px(15.0))
297 .color(if is_selected {
298 theme.primary.base
299 } else {
300 theme.neutral.text_3
301 }),
302 )
303 })
304 .child(
305 div()
306 .flex_1()
307 .text_size(px(theme.font_size.sm))
308 .child(item.label.clone()),
309 )
310 .on_click(move |_, window, cx| {
311 click_entity.update(cx, |tree, cx| {
312 tree.click_node(click_id.clone(), has_children, window, cx);
313 });
314 })
315 .into_any_element()
316 })
317 .size_full(),
318 )
319 .child(VirtualScrollbar::new(list_state))
320 }
321}
322
323pub fn flatten_visible(
324 data: &[TreeNode],
325 expanded_keys: &HashSet<SharedString>,
326) -> Vec<VirtualTreeItem> {
327 let mut output = Vec::new();
328 for node in data {
329 flatten_node(node, 0, expanded_keys, &mut output);
330 }
331 output
332}
333
334fn flatten_node(
335 node: &TreeNode,
336 depth: u32,
337 expanded_keys: &HashSet<SharedString>,
338 output: &mut Vec<VirtualTreeItem>,
339) {
340 let has_children = !node.children.is_empty();
341 output.push(VirtualTreeItem {
342 id: node.id.clone(),
343 label: node.label.clone(),
344 depth,
345 has_children,
346 });
347 if has_children && expanded_keys.contains(&node.id) {
348 for child in &node.children {
349 flatten_node(child, depth + 1, expanded_keys, output);
350 }
351 }
352}
353
354fn collect_parent_keys(data: &[TreeNode], output: &mut HashSet<SharedString>) {
355 for node in data {
356 if !node.children.is_empty() {
357 output.insert(node.id.clone());
358 collect_parent_keys(&node.children, output);
359 }
360 }
361}
362
363#[cfg(test)]
364mod tests {
365 use super::*;
366
367 fn sample_tree() -> Vec<TreeNode> {
368 vec![
369 TreeNode::new("root", "Root")
370 .child(TreeNode::new("a", "A"))
371 .child(TreeNode::new("b", "B").child(TreeNode::new("b1", "B1"))),
372 TreeNode::new("other", "Other"),
373 ]
374 }
375
376 #[test]
377 fn flatten_visible_only_includes_expanded_descendants() {
378 let tree = sample_tree();
379 let collapsed = flatten_visible(&tree, &HashSet::new());
380 assert_eq!(
381 collapsed
382 .iter()
383 .map(|item| item.id.as_ref())
384 .collect::<Vec<_>>(),
385 vec!["root", "other"]
386 );
387
388 let expanded = HashSet::from([SharedString::from("root")]);
389 let visible = flatten_visible(&tree, &expanded);
390 assert_eq!(
391 visible
392 .iter()
393 .map(|item| item.id.as_ref())
394 .collect::<Vec<_>>(),
395 vec!["root", "a", "b", "other"]
396 );
397
398 let expanded = HashSet::from([SharedString::from("root"), SharedString::from("b")]);
399 let visible = flatten_visible(&tree, &expanded);
400 assert_eq!(
401 visible
402 .iter()
403 .map(|item| item.id.as_ref())
404 .collect::<Vec<_>>(),
405 vec!["root", "a", "b", "b1", "other"]
406 );
407 }
408
409 #[test]
410 fn virtualized_tree_uses_list_state_and_visible_metadata() {
411 let source = include_str!("virtualized_tree.rs");
412
413 assert!(source.contains("pub struct VirtualizedTree"));
414 assert!(source.contains("ListState::new"));
415 assert!(source.contains("list(list_state.clone()"));
416 assert!(source.contains("VirtualScrollbar::new"));
417 assert!(source.contains("flatten_visible"));
418 assert!(source.contains("flattened: Vec<VirtualTreeItem>"));
419 }
420}