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