Merge #298

298: when evicting LRU data, keep dep information r=nikomatsakis a=nikomatsakis

and add a test that we do so!

cc `@Veykril` -- this should fix Rust-Analyzer performance

Co-authored-by: Niko Matsakis <niko@alum.mit.edu>

src/derived/fetch.rs

@@ -104,14 +104,6 @@ where
     }
 
     fn evict(&self, key_index: DerivedKeyIndex) {
-        if let Some(memo) = self.memo_map.get(key_index) {
+        self.memo_map.evict(key_index);
-            // Careful: we can't evict memos with untracked inputs
-            // as their values cannot be reconstructed.
-            if let QueryInputs::Untracked = memo.revisions.inputs {
-                return;
-            }
-
-            self.memo_map.remove(key_index);
-        }
     }
 }

src/derived/memo.rs

@@ -2,10 +2,12 @@ use std::sync::Arc;
 
 use arc_swap::{ArcSwap, Guard};
 use crossbeam_utils::atomic::AtomicCell;
+use dashmap::mapref::entry::Entry;
 
 use crate::{
-    hash::FxDashMap, runtime::local_state::QueryRevisions, DatabaseKeyIndex, Event, EventKind,
+    hash::FxDashMap,
-    Revision, Runtime,
+    runtime::local_state::{QueryInputs, QueryRevisions},
+    DatabaseKeyIndex, Event, EventKind, Revision, Runtime,
 };
 
 use super::DerivedKeyIndex;
@@ -28,9 +30,34 @@ impl<V> MemoMap<V> {
         self.map.insert(key, ArcSwap::from(Arc::new(memo)));
     }
 
-    /// Removes any existing memo for the given key.
+    /// Evicts the existing memo for the given key, replacing it
-    pub(super) fn remove(&self, key: DerivedKeyIndex) {
+    /// with an equivalent memo that has no value. If the memo
-        self.map.remove(&key);
+    /// has untracked inputs, this has no effect.
+    pub(super) fn evict(&self, key: DerivedKeyIndex) {
+        // Nit: this function embodies a touch more "business logic"
+        // than I would like (specifically the check about "query-input::untracked"),
+        // but I can't see a clean way to encapsulate it otherwise. I suppose
+        // it could take a closure, but it seems silly.
+        match self.map.entry(key) {
+            Entry::Vacant(_) => return,
+            Entry::Occupied(entry) => {
+                let memo = entry.get().load();
+
+                // Careful: we can't evict memos with untracked inputs
+                // as their values cannot be reconstructed.
+                if let QueryInputs::Untracked = memo.revisions.inputs {
+                    return;
+                }
+
+                let memo_evicted = Arc::new(Memo::new(
+                    None::<V>,
+                    memo.verified_at.load(),
+                    memo.revisions.clone(),
+                ));
+
+                entry.get().store(memo_evicted);
+            }
+        }
     }
 
     /// Loads the current memo for `key_index`. This does not hold any sort of

tests/lru.rs

@@ -1,37 +1,59 @@
 //! Test setting LRU actually limits the number of things in the database;
-use std::sync::{
+use std::{
-    atomic::{AtomicUsize, Ordering},
+    cell::RefCell,
-    Arc,
+    sync::{
+        atomic::{AtomicUsize, Ordering},
+        Arc,
+    },
 };
 
+use salsa::{Database as _, Durability};
+
+trait LruPeek {
+    fn log(&self, event: String);
+}
+
 #[derive(Debug, PartialEq, Eq)]
 struct HotPotato(u32);
 
-static N_POTATOES: AtomicUsize = AtomicUsize::new(0);
+thread_local! {
+    static N_POTATOES: AtomicUsize = AtomicUsize::new(0)
+}
 
 impl HotPotato {
     fn new(id: u32) -> HotPotato {
-        N_POTATOES.fetch_add(1, Ordering::SeqCst);
+        N_POTATOES.with(|n| n.fetch_add(1, Ordering::SeqCst));
         HotPotato(id)
     }
 }
 
 impl Drop for HotPotato {
     fn drop(&mut self) {
-        N_POTATOES.fetch_sub(1, Ordering::SeqCst);
+        N_POTATOES.with(|n| n.fetch_sub(1, Ordering::SeqCst));
     }
 }
 
 #[salsa::query_group(QueryGroupStorage)]
-trait QueryGroup: salsa::Database {
+trait QueryGroup: salsa::Database + LruPeek {
+    fn get2(&self, x: u32) -> u32;
     fn get(&self, x: u32) -> Arc<HotPotato>;
     fn get_volatile(&self, x: u32) -> usize;
 }
 
-fn get(_db: &dyn QueryGroup, x: u32) -> Arc<HotPotato> {
+/// Create a hotpotato (this will increment the counter above)
+fn get(db: &dyn QueryGroup, x: u32) -> Arc<HotPotato> {
+    db.log(format!("get({x})"));
     Arc::new(HotPotato::new(x))
 }
 
+/// Forward to the `get` query
+fn get2(db: &dyn QueryGroup, x: u32) -> u32 {
+    db.log(format!("get2({x})"));
+    db.get(x).0
+}
+
+// Like `get`, but with a volatile input, which means it can't
+// be LRU'd.
 fn get_volatile(db: &dyn QueryGroup, _x: u32) -> usize {
     static COUNTER: AtomicUsize = AtomicUsize::new(0);
     db.salsa_runtime().report_untracked_read();
@@ -42,50 +64,62 @@ fn get_volatile(db: &dyn QueryGroup, _x: u32) -> usize {
 #[derive(Default)]
 struct Database {
     storage: salsa::Storage<Self>,
+    logs: RefCell<Vec<String>>,
 }
 
 impl salsa::Database for Database {}
 
+impl LruPeek for Database {
+    fn log(&self, event: String) {
+        eprintln!("{event}");
+        self.logs.borrow_mut().push(event);
+    }
+}
+
+fn load_n_potatoes() -> usize {
+    N_POTATOES.with(|n| n.load(Ordering::SeqCst))
+}
+
 #[test]
 fn lru_works() {
     let mut db = Database::default();
     GetQuery.in_db_mut(&mut db).set_lru_capacity(32);
-    assert_eq!(N_POTATOES.load(Ordering::SeqCst), 0);
+    assert_eq!(load_n_potatoes(), 0);
 
     for i in 0..128u32 {
         let p = db.get(i);
         assert_eq!(p.0, i)
     }
-    assert_eq!(N_POTATOES.load(Ordering::SeqCst), 32);
+    assert_eq!(load_n_potatoes(), 32);
 
     for i in 0..128u32 {
         let p = db.get(i);
         assert_eq!(p.0, i)
     }
-    assert_eq!(N_POTATOES.load(Ordering::SeqCst), 32);
+    assert_eq!(load_n_potatoes(), 32);
 
     GetQuery.in_db_mut(&mut db).set_lru_capacity(32);
-    assert_eq!(N_POTATOES.load(Ordering::SeqCst), 32);
+    assert_eq!(load_n_potatoes(), 32);
 
     GetQuery.in_db_mut(&mut db).set_lru_capacity(64);
-    assert_eq!(N_POTATOES.load(Ordering::SeqCst), 32);
+    assert_eq!(load_n_potatoes(), 32);
     for i in 0..128u32 {
         let p = db.get(i);
         assert_eq!(p.0, i)
     }
-    assert_eq!(N_POTATOES.load(Ordering::SeqCst), 64);
+    assert_eq!(load_n_potatoes(), 64);
 
     // Special case: setting capacity to zero disables LRU
     GetQuery.in_db_mut(&mut db).set_lru_capacity(0);
-    assert_eq!(N_POTATOES.load(Ordering::SeqCst), 64);
+    assert_eq!(load_n_potatoes(), 64);
     for i in 0..128u32 {
         let p = db.get(i);
         assert_eq!(p.0, i)
     }
-    assert_eq!(N_POTATOES.load(Ordering::SeqCst), 128);
+    assert_eq!(load_n_potatoes(), 128);
 
     drop(db);
-    assert_eq!(N_POTATOES.load(Ordering::SeqCst), 0);
+    assert_eq!(load_n_potatoes(), 0);
 }
 
 #[test]
@@ -100,3 +134,30 @@ fn lru_doesnt_break_volatile_queries() {
         assert_eq!(x, i)
     }
 }
+
+#[test]
+fn lru_keeps_dependency_info() {
+    let mut db = Database::default();
+    let capacity = 4;
+    GetQuery.in_db_mut(&mut db).set_lru_capacity(capacity);
+
+    // Invoke `get2` 128 times. This will (in turn) invoke
+    // `get`, which will trigger LRU after 32 executions.
+    for i in 0..(capacity + 1) {
+        let p = db.get2(i as u32);
+        assert_eq!(p, i as u32);
+    }
+
+    db.salsa_runtime_mut().synthetic_write(Durability::HIGH);
+
+    // We want to test that calls to `get2` are still considered
+    // clean. Check that no new executions occur as we go here.
+    let events = db.logs.borrow().len();
+    assert_eq!(events, (capacity + 1) * 2);
+
+    // calling `get2(0)` has to check that `get(0)` is still valid;
+    // even though we've evicted it (LRU), we find that it is still good
+    let p = db.get2(0);
+    assert_eq!(p, 0);
+    assert_eq!(db.logs.borrow().len(), events);
+}