Skip to main content

gpui_rhai/
virtual_list.rs

1use std::cell::RefCell;
2use std::collections::{BTreeMap, BTreeSet};
3use std::ops::Range;
4use std::rc::Rc;
5use std::sync::Arc;
6
7use thiserror::Error;
8
9use crate::{ComponentInstancePath, UiNode};
10
11#[derive(Clone, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
12pub struct VirtualCollectionId {
13    pub component: ComponentInstancePath,
14    pub key: String,
15}
16
17#[derive(Clone, Debug, PartialEq)]
18pub struct VirtualCollectionNodeSpec {
19    pub id: VirtualCollectionId,
20    pub label: String,
21    pub data: crate::VirtualCollectionData,
22    pub realized: BTreeMap<usize, UiNode>,
23    pub estimated_height: f64,
24    /// Fixed logical-pixel viewport height, or `None` to fill the resolved
25    /// height offered by the parent flex layout.
26    pub height: Option<f64>,
27    pub overdraw_pixels: f64,
28    pub bottom_align: bool,
29    pub follow_tail: bool,
30    /// Stable data key to reveal when this controlled target changes.
31    pub reveal_key: Option<String>,
32    /// Item indices whose realized nodes act as top-pinned section headers.
33    pub sticky_headers: Arc<BTreeSet<usize>>,
34    /// Presentation-only motion group inherited from a surrounding
35    /// `motion_group`; applied to items realized after the initial render.
36    pub inherited_motion_group: Option<String>,
37    /// The component instance that declared the inherited group.
38    pub inherited_motion_scope: Option<String>,
39}
40
41#[derive(Clone, Debug, Default)]
42pub struct VirtualRequestRegistry {
43    requests: Rc<RefCell<BTreeMap<VirtualCollectionId, BTreeSet<usize>>>>,
44    metrics: Rc<RefCell<BTreeMap<VirtualCollectionId, VirtualCollectionMetrics>>>,
45    item_bounds: Rc<RefCell<BTreeMap<VirtualCollectionId, ItemBoundsReader>>>,
46    wake: crate::async_runtime::AsyncWake,
47}
48
49/// Reads a mounted list's item bounds, in window coordinates, at call time.
50#[derive(Clone)]
51pub(crate) struct ItemBoundsReader(Rc<dyn Fn(usize) -> Option<crate::GeometryBounds>>);
52
53impl std::fmt::Debug for ItemBoundsReader {
54    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
55        formatter.write_str("ItemBoundsReader")
56    }
57}
58
59#[derive(Clone, Debug, PartialEq)]
60pub struct VirtualCollectionMetrics {
61    pub id: VirtualCollectionId,
62    pub item_count: usize,
63    pub realized_count: usize,
64    pub realized_range: Range<usize>,
65    pub requested_count: usize,
66    pub requested_range: Range<usize>,
67    pub visible_range: Range<usize>,
68    pub viewport_height: f64,
69    pub scroll_item: usize,
70    pub scroll_offset: f64,
71    pub is_scrolled: bool,
72    pub sticky_header: Option<usize>,
73    pub bottom_align: bool,
74    pub follow_tail: bool,
75}
76
77impl VirtualCollectionMetrics {
78    fn new(id: VirtualCollectionId) -> Self {
79        Self {
80            id,
81            item_count: 0,
82            realized_count: 0,
83            realized_range: 0..0,
84            requested_count: 0,
85            requested_range: 0..0,
86            visible_range: 0..0,
87            viewport_height: 0.0,
88            scroll_item: 0,
89            scroll_offset: 0.0,
90            is_scrolled: false,
91            sticky_header: None,
92            bottom_align: false,
93            follow_tail: false,
94        }
95    }
96}
97
98#[derive(Clone, Debug)]
99pub(crate) struct VirtualRequestSnapshot {
100    requests: BTreeMap<VirtualCollectionId, BTreeSet<usize>>,
101    metrics: BTreeMap<VirtualCollectionId, VirtualCollectionMetrics>,
102}
103
104impl VirtualRequestRegistry {
105    #[must_use]
106    pub fn new() -> Self {
107        Self::default()
108    }
109
110    pub fn request(&self, id: VirtualCollectionId, indices: impl IntoIterator<Item = usize>) {
111        let indices = indices.into_iter().collect::<Vec<_>>();
112        let requested_count = {
113            let mut requests = self.requests.borrow_mut();
114            let requested = requests.entry(id.clone()).or_default();
115            requested.extend(indices.iter().copied());
116            requested.len()
117        };
118        let new_range = index_range(indices);
119        let mut metrics = self.metrics.borrow_mut();
120        let metrics = metrics
121            .entry(id.clone())
122            .or_insert_with(|| VirtualCollectionMetrics::new(id));
123        metrics.requested_count = requested_count;
124        if !new_range.is_empty() {
125            metrics.requested_range = if metrics.requested_range.is_empty() {
126                new_range
127            } else {
128                metrics.requested_range.start.min(new_range.start)
129                    ..metrics.requested_range.end.max(new_range.end)
130            };
131        }
132        self.wake.notify();
133    }
134
135    /// Replace one collection's pending request with a complete atomic target.
136    ///
137    /// Unlike [`Self::request`], this does not union partial layout callbacks.
138    /// The virtual element calls it only after collecting a complete prepaint.
139    pub(crate) fn request_target(
140        &self,
141        id: VirtualCollectionId,
142        indices: impl IntoIterator<Item = usize>,
143    ) {
144        let indices = indices.into_iter().collect::<BTreeSet<_>>();
145        let requested_count = indices.len();
146        self.requests
147            .borrow_mut()
148            .insert(id.clone(), indices.clone());
149        let requested_range = index_range(indices);
150        let mut metrics = self.metrics.borrow_mut();
151        let metrics = metrics
152            .entry(id.clone())
153            .or_insert_with(|| VirtualCollectionMetrics::new(id));
154        metrics.requested_count = requested_count;
155        metrics.requested_range = requested_range;
156        self.wake.notify();
157    }
158
159    pub(crate) fn clear_target(&self, id: &VirtualCollectionId) {
160        self.requests.borrow_mut().remove(id);
161        if let Some(metrics) = self.metrics.borrow_mut().get_mut(id) {
162            metrics.requested_count = 0;
163            metrics.requested_range = 0..0;
164        }
165    }
166
167    #[must_use]
168    pub fn is_empty(&self) -> bool {
169        self.requests.borrow().is_empty()
170    }
171
172    #[must_use]
173    pub(crate) fn has_scope(&self, root: &ComponentInstancePath) -> bool {
174        self.requests
175            .borrow()
176            .keys()
177            .any(|id| id.component.is_within(root))
178    }
179
180    pub(crate) fn drain(&self) -> BTreeMap<VirtualCollectionId, BTreeSet<usize>> {
181        let requests = std::mem::take(&mut *self.requests.borrow_mut());
182        for metrics in self.metrics.borrow_mut().values_mut() {
183            metrics.requested_count = 0;
184            metrics.requested_range = 0..0;
185        }
186        requests
187    }
188
189    pub(crate) fn wake(&self) -> crate::async_runtime::AsyncWake {
190        self.wake.clone()
191    }
192
193    pub(crate) fn snapshot(&self) -> VirtualRequestSnapshot {
194        VirtualRequestSnapshot {
195            requests: self.requests.borrow().clone(),
196            metrics: self.metrics.borrow().clone(),
197        }
198    }
199
200    pub(crate) fn restore(&self, snapshot: VirtualRequestSnapshot) {
201        *self.requests.borrow_mut() = snapshot.requests;
202        *self.metrics.borrow_mut() = snapshot.metrics;
203    }
204
205    pub(crate) fn retain(&self, active: &BTreeSet<VirtualCollectionId>) {
206        self.requests
207            .borrow_mut()
208            .retain(|id, _| active.contains(id));
209        self.metrics
210            .borrow_mut()
211            .retain(|id, _| active.contains(id));
212        self.item_bounds
213            .borrow_mut()
214            .retain(|id, _| active.contains(id));
215    }
216
217    /// Install the reader of a mounted list; the list removes it when it is
218    /// dropped (the reader holds the list's state).
219    pub(crate) fn set_item_bounds_reader(
220        &self,
221        id: VirtualCollectionId,
222        read: impl Fn(usize) -> Option<crate::GeometryBounds> + 'static,
223    ) -> ItemBoundsReader {
224        let reader = ItemBoundsReader(Rc::new(read));
225        self.item_bounds.borrow_mut().insert(id, reader.clone());
226        reader
227    }
228
229    pub(crate) fn remove_item_bounds_reader(
230        &self,
231        id: &VirtualCollectionId,
232        reader: &ItemBoundsReader,
233    ) {
234        let mut readers = self.item_bounds.borrow_mut();
235        if readers
236            .get(id)
237            .is_some_and(|current| Rc::ptr_eq(&current.0, &reader.0))
238        {
239            readers.remove(id);
240        }
241    }
242
243    /// The window bounds of a laid-out item of a mounted collection.
244    pub(crate) fn item_bounds(
245        &self,
246        id: &VirtualCollectionId,
247        index: usize,
248    ) -> Option<crate::GeometryBounds> {
249        let reader = self.item_bounds.borrow().get(id).cloned()?;
250        (reader.0)(index)
251    }
252
253    pub(crate) fn report_frame(
254        &self,
255        spec: &VirtualCollectionNodeSpec,
256        viewport_height: f64,
257        scroll_item: usize,
258        scroll_offset: f64,
259        measured_visible: Option<Range<usize>>,
260    ) {
261        let mut metrics = self.metrics.borrow_mut();
262        let metrics = metrics
263            .entry(spec.id.clone())
264            .or_insert_with(|| VirtualCollectionMetrics::new(spec.id.clone()));
265        metrics.item_count = spec.data.len();
266        metrics.realized_count = spec.realized.len();
267        metrics.realized_range = index_range(spec.realized.keys().copied());
268        if let Some(visible_range) = measured_visible {
269            metrics.visible_range = visible_range;
270        } else if metrics.visible_range.is_empty() {
271            metrics.visible_range = metrics.realized_range.clone();
272        } else {
273            metrics.visible_range = metrics.visible_range.start.min(spec.data.len())
274                ..metrics.visible_range.end.min(spec.data.len());
275        }
276        metrics.viewport_height = viewport_height;
277        metrics.scroll_item = scroll_item;
278        metrics.scroll_offset = scroll_offset;
279        metrics.sticky_header = spec
280            .sticky_headers
281            .range(..=scroll_item)
282            .next_back()
283            .copied();
284        metrics.bottom_align = spec.bottom_align;
285        metrics.follow_tail = spec.follow_tail;
286    }
287
288    pub(crate) fn report_scroll(
289        &self,
290        id: &VirtualCollectionId,
291        visible_range: Range<usize>,
292        is_scrolled: bool,
293    ) {
294        let mut metrics = self.metrics.borrow_mut();
295        let metrics = metrics
296            .entry(id.clone())
297            .or_insert_with(|| VirtualCollectionMetrics::new(id.clone()));
298        metrics.visible_range = visible_range;
299        metrics.is_scrolled = is_scrolled;
300    }
301
302    #[must_use]
303    pub fn inspect(&self) -> Vec<VirtualCollectionMetrics> {
304        self.metrics.borrow().values().cloned().collect()
305    }
306
307    pub(crate) fn metrics(&self, id: &VirtualCollectionId) -> Option<VirtualCollectionMetrics> {
308        self.metrics.borrow().get(id).cloned()
309    }
310}
311
312fn index_range(indices: impl IntoIterator<Item = usize>) -> Range<usize> {
313    let mut indices = indices.into_iter();
314    let Some(first) = indices.next() else {
315        return 0..0;
316    };
317    let (min, max) = indices.fold((first, first), |(min, max), index| {
318        (min.min(index), max.max(index))
319    });
320    min..max.saturating_add(1)
321}
322
323fn validate_viewport(scroll_offset: f64, viewport_height: f64) -> Result<(), VirtualListError> {
324    if !scroll_offset.is_finite() || scroll_offset < 0.0 {
325        return Err(VirtualListError::InvalidScrollOffset(scroll_offset));
326    }
327    if !viewport_height.is_finite() || viewport_height < 0.0 {
328        return Err(VirtualListError::InvalidViewport(viewport_height));
329    }
330    Ok(())
331}
332
333#[derive(Clone, Copy, Debug, PartialEq)]
334pub struct VariableListSpec {
335    pub estimated_height: f64,
336    pub overscan_pixels: f64,
337}
338
339impl VariableListSpec {
340    /// Construct a variable-height one-dimensional virtualization policy.
341    ///
342    /// # Errors
343    ///
344    /// Returns for non-finite/non-positive estimates or negative overscan.
345    pub fn new(estimated_height: f64, overscan_pixels: f64) -> Result<Self, VirtualListError> {
346        if !estimated_height.is_finite() || estimated_height <= 0.0 {
347            return Err(VirtualListError::InvalidEstimatedHeight(estimated_height));
348        }
349        if !overscan_pixels.is_finite() || overscan_pixels < 0.0 {
350            return Err(VirtualListError::InvalidOverscan(overscan_pixels));
351        }
352        Ok(Self {
353            estimated_height,
354            overscan_pixels,
355        })
356    }
357}
358
359#[derive(Clone, Debug, Default)]
360pub struct VariableListState {
361    keys: Vec<String>,
362    indices: BTreeMap<String, usize>,
363    measured: BTreeMap<String, f64>,
364    heights: Vec<f64>,
365    prefix: FenwickTree,
366}
367
368impl VariableListState {
369    /// Install ordered keys while preserving measurements by identity.
370    ///
371    /// # Errors
372    ///
373    /// Returns for invalid policy or duplicate keys.
374    pub fn set_keys(
375        &mut self,
376        keys: Vec<String>,
377        spec: VariableListSpec,
378    ) -> Result<(), VirtualListError> {
379        VariableListSpec::new(spec.estimated_height, spec.overscan_pixels)?;
380        let mut unique = BTreeSet::new();
381        if let Some(key) = keys.iter().find(|key| !unique.insert((*key).clone())) {
382            return Err(VirtualListError::DuplicateKey(key.clone()));
383        }
384        self.measured.retain(|key, _| unique.contains(key));
385        self.indices = keys
386            .iter()
387            .enumerate()
388            .map(|(index, key)| (key.clone(), index))
389            .collect();
390        self.heights = keys
391            .iter()
392            .map(|key| {
393                self.measured
394                    .get(key)
395                    .copied()
396                    .unwrap_or(spec.estimated_height)
397            })
398            .collect();
399        self.prefix = FenwickTree::from_values(&self.heights);
400        self.keys = keys;
401        Ok(())
402    }
403
404    /// Record a measured item height and return the scroll correction required
405    /// to keep an anchor key at the same viewport position.
406    ///
407    /// # Errors
408    ///
409    /// Returns for unknown keys or invalid measurements.
410    pub fn measure(
411        &mut self,
412        key: &str,
413        height: f64,
414        anchor: Option<&str>,
415    ) -> Result<f64, VirtualListError> {
416        if !height.is_finite() || height <= 0.0 {
417            return Err(VirtualListError::InvalidMeasurement(height));
418        }
419        let index = *self
420            .indices
421            .get(key)
422            .ok_or_else(|| VirtualListError::UnknownKey(key.to_owned()))?;
423        let anchor_index = anchor
424            .map(|anchor| {
425                self.indices
426                    .get(anchor)
427                    .copied()
428                    .ok_or_else(|| VirtualListError::UnknownKey(anchor.to_owned()))
429            })
430            .transpose()?;
431        let before = anchor_index.map_or(0.0, |anchor| self.prefix.sum(anchor));
432        let delta = height - self.heights[index];
433        self.heights[index] = height;
434        self.measured.insert(key.to_owned(), height);
435        self.prefix.add(index, delta);
436        let after = anchor_index.map_or(0.0, |anchor| self.prefix.sum(anchor));
437        Ok(after - before)
438    }
439
440    /// Compute the bounded realization window and spacer geometry.
441    ///
442    /// # Errors
443    ///
444    /// Returns viewport validation errors.
445    pub fn window(
446        &self,
447        spec: VariableListSpec,
448        scroll_offset: f64,
449        viewport_height: f64,
450    ) -> Result<VariableListWindow, VirtualListError> {
451        validate_viewport(scroll_offset, viewport_height)?;
452        VariableListSpec::new(spec.estimated_height, spec.overscan_pixels)?;
453        if self.keys.is_empty() {
454            return Ok(VariableListWindow::default());
455        }
456        let start_offset = (scroll_offset - spec.overscan_pixels).max(0.0);
457        let end_offset = scroll_offset + viewport_height + spec.overscan_pixels;
458        let start = self.prefix.lower_bound(start_offset).min(self.keys.len());
459        let end = self
460            .prefix
461            .lower_bound(end_offset)
462            .saturating_add(1)
463            .min(self.keys.len());
464        Ok(VariableListWindow {
465            range: start..end,
466            leading: self.prefix.sum(start),
467            trailing: (self.prefix.total() - self.prefix.sum(end)).max(0.0),
468            total: self.prefix.total(),
469        })
470    }
471
472    /// Return the positive visible offset that keeps the bottom aligned.
473    #[must_use]
474    pub fn tail_offset(&self, viewport_height: f64) -> f64 {
475        (self.prefix.total() - viewport_height.max(0.0)).max(0.0)
476    }
477
478    #[must_use]
479    pub fn measured_count(&self) -> usize {
480        self.measured.len()
481    }
482}
483
484#[derive(Clone, Debug, Default, PartialEq)]
485pub struct VariableListWindow {
486    pub range: Range<usize>,
487    pub leading: f64,
488    pub trailing: f64,
489    pub total: f64,
490}
491
492#[derive(Clone, Debug, Default)]
493struct FenwickTree {
494    tree: Vec<f64>,
495}
496
497impl FenwickTree {
498    fn from_values(values: &[f64]) -> Self {
499        let mut tree = Self {
500            tree: vec![0.0; values.len() + 1],
501        };
502        for (index, value) in values.iter().copied().enumerate() {
503            tree.add(index, value);
504        }
505        tree
506    }
507
508    fn add(&mut self, index: usize, delta: f64) {
509        let mut cursor = index + 1;
510        while cursor < self.tree.len() {
511            self.tree[cursor] += delta;
512            cursor += cursor & cursor.wrapping_neg();
513        }
514    }
515
516    fn sum(&self, end: usize) -> f64 {
517        let mut cursor = end.min(self.tree.len().saturating_sub(1));
518        let mut total = 0.0;
519        while cursor > 0 {
520            total += self.tree[cursor];
521            cursor &= cursor - 1;
522        }
523        total
524    }
525
526    fn total(&self) -> f64 {
527        self.sum(self.tree.len().saturating_sub(1))
528    }
529
530    fn lower_bound(&self, target: f64) -> usize {
531        let mut index = 0usize;
532        let mut accumulated = 0.0;
533        let mut step = self.tree.len().next_power_of_two() / 2;
534        while step > 0 {
535            let next = index + step;
536            if next < self.tree.len() && accumulated + self.tree[next] <= target {
537                index = next;
538                accumulated += self.tree[next];
539            }
540            step /= 2;
541        }
542        index.min(self.tree.len().saturating_sub(1))
543    }
544}
545
546#[derive(Clone, Debug, Error, PartialEq)]
547pub enum VirtualListError {
548    #[error("scroll offset must be finite and non-negative, got {0}")]
549    InvalidScrollOffset(f64),
550    #[error("viewport height must be finite and non-negative, got {0}")]
551    InvalidViewport(f64),
552    #[error("virtual list key `{0}` is duplicated")]
553    DuplicateKey(String),
554    #[error("virtual list key `{0}` is unknown")]
555    UnknownKey(String),
556    #[error("estimated item height must be finite and positive, got {0}")]
557    InvalidEstimatedHeight(f64),
558    #[error("virtual-list overscan pixels must be finite and non-negative, got {0}")]
559    InvalidOverscan(f64),
560    #[error("measured item height must be finite and positive, got {0}")]
561    InvalidMeasurement(f64),
562}
563
564#[cfg(test)]
565mod tests {
566    use super::*;
567    use crate::UiValue;
568
569    #[test]
570    fn variable_height_window_is_bounded_and_measurement_preserves_anchor() {
571        let spec = VariableListSpec::new(20.0, 80.0).unwrap();
572        let mut list = VariableListState::default();
573        list.set_keys(
574            (0..10_000).map(|index| format!("row-{index}")).collect(),
575            spec,
576        )
577        .unwrap();
578        let window = list.window(spec, 50_000.0, 600.0).unwrap();
579        assert!(window.range.len() <= 40, "{window:?}");
580        let correction = list.measure("row-0", 50.0, Some("row-2500")).unwrap();
581        assert!((correction - 30.0).abs() < f64::EPSILON);
582        let corrected = list.window(spec, 50_030.0, 600.0).unwrap();
583        assert_eq!(corrected.range, window.range);
584        assert_eq!(list.measured_count(), 1);
585    }
586
587    #[test]
588    fn variable_height_measurements_survive_reorder_and_tail_alignment() {
589        let spec = VariableListSpec::new(20.0, 0.0).unwrap();
590        let mut list = VariableListState::default();
591        list.set_keys(vec!["a".into(), "b".into(), "c".into()], spec)
592            .unwrap();
593        list.measure("b", 60.0, None).unwrap();
594        list.set_keys(vec!["c".into(), "b".into(), "a".into()], spec)
595            .unwrap();
596        assert_eq!(list.measured_count(), 1);
597        assert!((list.tail_offset(50.0) - 50.0).abs() < f64::EPSILON);
598    }
599
600    #[test]
601    fn virtual_collection_metrics_track_frame_scroll_requests_and_rollback() {
602        let id = VirtualCollectionId {
603            component: ComponentInstancePath::root("Chat", "main"),
604            key: "messages".to_owned(),
605        };
606        let item = |key: &str| {
607            UiValue::Map(BTreeMap::from([(
608                "key".to_owned(),
609                UiValue::String(key.to_owned()),
610            )]))
611        };
612        let spec = VirtualCollectionNodeSpec {
613            id: id.clone(),
614            label: "Messages".to_owned(),
615            data: (0..5).map(|index| item(&format!("row-{index}"))).collect(),
616            realized: BTreeMap::from([(1, UiNode::text("one")), (2, UiNode::text("two"))]),
617            estimated_height: 24.0,
618            height: Some(120.0),
619            overdraw_pixels: 48.0,
620            bottom_align: true,
621            follow_tail: true,
622            reveal_key: None,
623            sticky_headers: Arc::new(BTreeSet::new()),
624            inherited_motion_group: None,
625            inherited_motion_scope: None,
626        };
627        let registry = VirtualRequestRegistry::new();
628        registry.report_frame(&spec, 120.0, 1, 3.5, Some(1..3));
629        registry.request(id.clone(), [3, 4]);
630        registry.report_scroll(&id, 1..4, true);
631        let metrics = &registry.inspect()[0];
632        assert_eq!(metrics.realized_range, 1..3);
633        assert_eq!(metrics.requested_range, 3..5);
634        assert_eq!(metrics.visible_range, 1..4);
635        assert!(metrics.bottom_align && metrics.follow_tail && metrics.is_scrolled);
636        assert_eq!(metrics.sticky_header, None);
637
638        let snapshot = registry.snapshot();
639        let _ = registry.drain();
640        assert_eq!(registry.inspect()[0].requested_count, 0);
641        registry.restore(snapshot);
642        assert_eq!(registry.inspect()[0].requested_count, 2);
643        registry.retain(&BTreeSet::new());
644        assert_eq!(
645            registry.inspect(),
646            [] as [crate::virtual_list::VirtualCollectionMetrics; 0]
647        );
648    }
649}