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