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