Skip to main content

lattice_host/
per_buffer_cache.rs

1//! `PerBufferCache<T>` — per-buffer cache primitive for the
2//! 5.8.AF.5 Slice 3b migration.
3//!
4//! Every LSP feature that produces a per-buffer cache (inlay
5//! hints, folding ranges, semantic tokens, code lenses, document
6//! links, document colors, pull diagnostics, …) needs:
7//!
8//! - **Wait-free reads** at render time (renderer paints every
9//!   frame; can't take a lock).
10//! - **Concurrent writes** from background tasks on the LSP
11//!   runtime (the spawned request task `.store()`s the result
12//!   when the response arrives — paramount goal #4).
13//! - **Per-buffer keying** so closing buffer A doesn't affect
14//!   buffer B's cache.
15//! - **Shareability** so the renderer's `RenderState` snapshot
16//!   can hold a clone that observes writes by the task.
17//!
18//! `Arc<ArcSwap<HashMap<BufferId, Arc<T>>>>` satisfies all four:
19//!
20//! - The outer `Arc` lets the spawned task clone the slot into
21//!   itself.
22//! - The `ArcSwap` makes the inner `HashMap` swappable atomically.
23//! - The inner `Arc<T>` per entry lets readers detach a value
24//!   cheaply (one Arc bump) without holding a lock.
25//! - The `ArcSwap` snapshot is wait-free for readers.
26//!
27//! Writes are copy-on-write: a writer loads the current
28//! `Arc<HashMap>`, clones the underlying `HashMap`, mutates the
29//! clone, and stores. For typical sessions (~5–50 open buffers)
30//! the clone cost is microseconds and writes are rare
31//! (~1 per LSP response per cache type), so the cost stays well
32//! below the 100µs publication budget.
33
34use std::collections::HashMap;
35use std::sync::Arc;
36
37use arc_swap::ArcSwap;
38use lattice_core::BufferId;
39
40/// Per-buffer cache slot — the canonical Slice 3b shape for
41/// per-buffer LSP feature data.
42///
43/// Use [`PerBufferCacheExt`] for the read / write / remove
44/// helpers; the raw type exposes only what `ArcSwap` exposes.
45pub type PerBufferCache<T> = Arc<ArcSwap<HashMap<BufferId, Arc<T>>>>;
46
47/// Construct an empty `PerBufferCache<T>`. The Slice 3b boot
48/// path uses this; downstream code clones the resulting Arc.
49pub fn empty<T>() -> PerBufferCache<T> {
50    Arc::new(ArcSwap::from_pointee(HashMap::new()))
51}
52
53/// Convenience trait carrying the standard read / insert /
54/// remove operations for a [`PerBufferCache`].
55///
56/// All operations are non-blocking. Writes use a copy-on-write
57/// pattern: load the current `Arc<HashMap>`, clone the inner
58/// `HashMap`, mutate, store. Concurrent writers may race; the
59/// last-writer-wins outcome is acceptable for LSP feature caches
60/// (each buffer's writer is itself single-flight via
61/// cancellation tokens; cross-buffer races are independent).
62pub trait PerBufferCacheExt<T> {
63    /// Wait-free read: returns a detached `Arc<T>` snapshot of
64    /// the cache entry for `id`, or `None` if no entry exists.
65    /// The returned `Arc` is independent of any subsequent
66    /// store; the renderer can hold it across the frame.
67    fn get_for(&self, id: BufferId) -> Option<Arc<T>>;
68
69    /// Store (or replace) the cache entry for `id`. Copy-on-
70    /// write: clones the current `HashMap`, inserts, stores.
71    fn insert_for(&self, id: BufferId, value: T);
72
73    /// Remove the cache entry for `id` if present. No-op when
74    /// the entry doesn't exist. Copy-on-write semantics.
75    fn remove_for(&self, id: BufferId);
76
77    /// Retain only entries matching the predicate. Mirrors
78    /// `HashMap::retain`'s semantics: the predicate sees each
79    /// `(BufferId, &T)`; entries returning `false` are dropped.
80    ///
81    /// Copy-on-write: if no entries need to be dropped, the
82    /// underlying Arc is unchanged. Used by LSP `*/refresh`
83    /// drains to evict per-server caches when a server's
84    /// invalidation notification arrives.
85    fn retain<F: FnMut(BufferId, &T) -> bool>(&self, predicate: F);
86
87    /// Returns `true` when no entries exist. Wait-free.
88    fn is_empty_snapshot(&self) -> bool;
89}
90
91impl<T> PerBufferCacheExt<T> for PerBufferCache<T> {
92    fn get_for(&self, id: BufferId) -> Option<Arc<T>> {
93        self.load().get(&id).cloned()
94    }
95
96    fn insert_for(&self, id: BufferId, value: T) {
97        let current = self.load();
98        let mut next = (**current).clone();
99        next.insert(id, Arc::new(value));
100        self.store(Arc::new(next));
101    }
102
103    fn remove_for(&self, id: BufferId) {
104        let current = self.load();
105        if !current.contains_key(&id) {
106            // Avoid the clone when the key is absent — common
107            // when buffer-close fires for buffers the cache
108            // never observed.
109            return;
110        }
111        let mut next = (**current).clone();
112        next.remove(&id);
113        self.store(Arc::new(next));
114    }
115
116    fn retain<F: FnMut(BufferId, &T) -> bool>(&self, mut predicate: F) {
117        let current = self.load();
118        // First pass: identify keys to drop without cloning the
119        // map. Skip the rebuild entirely when nothing matches.
120        let drop_keys: Vec<BufferId> = current
121            .iter()
122            .filter_map(|(id, value)| {
123                if predicate(*id, value) {
124                    None
125                } else {
126                    Some(*id)
127                }
128            })
129            .collect();
130        if drop_keys.is_empty() {
131            return;
132        }
133        let mut next = (**current).clone();
134        for id in drop_keys {
135            next.remove(&id);
136        }
137        self.store(Arc::new(next));
138    }
139
140    fn is_empty_snapshot(&self) -> bool {
141        self.load().is_empty()
142    }
143}
144
145#[cfg(test)]
146mod tests {
147    use super::*;
148
149    /// Basic insert / get round-trip. Verifies the Arc-bump
150    /// read contract.
151    #[test]
152    fn insert_and_get_roundtrip() {
153        let cache: PerBufferCache<u32> = empty();
154        let id = BufferId(7);
155        cache.insert_for(id, 42);
156        let v = cache.get_for(id).expect("entry must be present");
157        assert_eq!(*v, 42);
158    }
159
160    /// Remove on a non-existent key is a no-op (does NOT clone
161    /// the map, which matters for the buffer-close hot path
162    /// that runs on every close regardless of whether the cache
163    /// holds an entry for that buffer).
164    #[test]
165    fn remove_missing_skips_clone() {
166        let cache: PerBufferCache<u32> = empty();
167        let before = Arc::as_ptr(&cache.load_full());
168        cache.remove_for(BufferId(99));
169        let after = Arc::as_ptr(&cache.load_full());
170        assert_eq!(
171            before, after,
172            "remove of absent key must not allocate a new Arc"
173        );
174    }
175
176    /// A clone of the outer `Arc` observes writes made through
177    /// the original — this is the property that lets
178    /// `RenderState.lsp.<cache>` see writes made by the spawned
179    /// task on the LSP runtime without re-publishing
180    /// `RenderState`.
181    #[test]
182    fn clone_observes_writes_through_arcswap() {
183        let cache: PerBufferCache<u32> = empty();
184        let shared = cache.clone();
185        cache.insert_for(BufferId(1), 100);
186        let v = shared.get_for(BufferId(1)).expect("clone sees write");
187        assert_eq!(*v, 100);
188    }
189
190    /// Writes through one clone are observable through another
191    /// clone — symmetric to the test above; covers the case
192    /// where the spawned task holds clone A and the renderer
193    /// reads via clone B.
194    #[test]
195    fn writes_through_one_clone_visible_to_another() {
196        let cache_a: PerBufferCache<&'static str> = empty();
197        let cache_b = cache_a.clone();
198        cache_b.insert_for(BufferId(2), "from-b");
199        assert_eq!(*cache_a.get_for(BufferId(2)).unwrap(), "from-b");
200    }
201}