Merge #363

363: Add method to change LRU capacity at runtime r=nikomatsakis a=XFFXFF

closes #355 

Co-authored-by: XFFXFF <1247714429@qq.com>

components/salsa-2022-macros/src/tracked_fn.rs

@@ -290,6 +290,7 @@ fn wrapper_fns(
     let accumulated_fn = accumulated_fn(args, item_fn, config_ty)?;
     let setter_fn = setter_fn(args, item_fn, config_ty)?;
     let specify_fn = specify_fn(args, item_fn, config_ty)?.map(|f| quote! { #f });
+    let set_lru_fn = set_lru_capacity_fn(args, config_ty)?.map(|f| quote! { #f });
 
     let setter_impl: syn::ItemImpl = parse_quote! {
         impl #config_ty {
@@ -302,6 +303,8 @@ fn wrapper_fns(
             #[allow(dead_code, clippy::needless_lifetimes)]
             #accumulated_fn
 
+            #set_lru_fn
+
             #specify_fn
         }
     };
@@ -417,6 +420,39 @@ fn setter_fn(
     })
 }
 
+/// Create a `set_lru_capacity` associated function that can be used to change LRU
+/// capacity at runtime.
+/// Note that this function is only generated if the tracked function has the lru option set.
+///
+/// # Examples
+///
+/// ```rust
+/// #[salsa::tracked(lru=32)]
+/// fn my_tracked_fn(db: &dyn crate::Db, ...) { }
+///
+/// my_tracked_fn::set_lru_capacity(16)
+/// ```
+fn set_lru_capacity_fn(
+    args: &Args,
+    config_ty: &syn::Type,
+) -> syn::Result<Option<syn::ImplItemMethod>> {
+    if args.lru.is_none() {
+        return Ok(None);
+    }
+
+    let jar_ty = args.jar_ty();
+    let lru_fn = parse_quote! {
+        #[allow(dead_code, clippy::needless_lifetimes)]
+        fn set_lru_capacity(__db: &salsa::function::DynDb<Self>, __value: usize) {
+            let (__jar, __runtime) = <_ as salsa::storage::HasJar<#jar_ty>>::jar(__db);
+            let __ingredients =
+                <_ as salsa::storage::HasIngredientsFor<#config_ty>>::ingredient(__jar);
+            __ingredients.function.set_capacity(__value);
+        }
+    };
+    Ok(Some(lru_fn))
+}
+
 fn specify_fn(
     args: &Args,
     item_fn: &syn::ItemFn,

salsa-2022-tests/src/lib.rs

@@ -21,4 +21,12 @@ pub trait HasLogger {
         let logs = std::mem::replace(&mut *self.logger().logs.lock().unwrap(), vec![]);
         expected.assert_eq(&format!("{:#?}", logs));
     }
+
+    /// Asserts the length of the logs,
+    /// clearing the logged events. This takes `&mut self` because
+    /// it is meant to be run from outside any tracked functions.
+    fn assert_logs_len(&mut self, expected: usize) {
+        let logs = std::mem::replace(&mut *self.logger().logs.lock().unwrap(), vec![]);
+        assert_eq!(logs.len(), expected);
+    }
 }

salsa-2022-tests/tests/expect_reuse_field_x_of_a_tracked_struct_changes_but_fn_depends_on_field_y.rs

@@ -11,8 +11,8 @@ use expect_test::expect;
 struct Jar(
     MyInput,
     MyTracked,
-    final_result_deponds_on_x,
-    final_result_deponds_on_y,
+    final_result_depends_on_x,
+    final_result_depends_on_y,
     intermediate_result,
 );
 
@@ -24,14 +24,14 @@ struct MyInput {
 }
 
 #[salsa::tracked(jar = Jar)]
-fn final_result_deponds_on_x(db: &dyn Db, input: MyInput) -> u32 {
-    db.push_log(format!("final_result_deponds_on_x({:?})", input));
+fn final_result_depends_on_x(db: &dyn Db, input: MyInput) -> u32 {
+    db.push_log(format!("final_result_depends_on_x({:?})", input));
     intermediate_result(db, input).x(db) * 2
 }
 
 #[salsa::tracked(jar = Jar)]
-fn final_result_deponds_on_y(db: &dyn Db, input: MyInput) -> u32 {
-    db.push_log(format!("final_result_deponds_on_y({:?})", input));
+fn final_result_depends_on_y(db: &dyn Db, input: MyInput) -> u32 {
+    db.push_log(format!("final_result_depends_on_y({:?})", input));
     intermediate_result(db, input).y(db) * 2
 }
 
@@ -71,39 +71,39 @@ impl HasLogger for Database {
 fn execute() {
     // x = (input.field + 1) / 2
     // y = input.field / 2
-    // final_result_deponds_on_x = x * 2 = (input.field + 1) / 2 * 2
-    // final_result_deponds_on_y = y * 2 = input.field / 2 * 2
+    // final_result_depends_on_x = x * 2 = (input.field + 1) / 2 * 2
+    // final_result_depends_on_y = y * 2 = input.field / 2 * 2
     let mut db = Database::default();
 
     // intermediate results:
     // x = (22 + 1) / 2 = 11
     // y = 22 / 2 = 11
     let input = MyInput::new(&mut db, 22);
-    assert_eq!(final_result_deponds_on_x(&db, input), 22);
+    assert_eq!(final_result_depends_on_x(&db, input), 22);
     db.assert_logs(expect![[r#"
         [
-            "final_result_deponds_on_x(MyInput(Id { value: 1 }))",
+            "final_result_depends_on_x(MyInput(Id { value: 1 }))",
         ]"#]]);
 
-    assert_eq!(final_result_deponds_on_y(&db, input), 22);
+    assert_eq!(final_result_depends_on_y(&db, input), 22);
     db.assert_logs(expect![[r#"
         [
-            "final_result_deponds_on_y(MyInput(Id { value: 1 }))",
+            "final_result_depends_on_y(MyInput(Id { value: 1 }))",
         ]"#]]);
 
     input.set_field(&mut db, 23);
     // x = (23 + 1) / 2 = 12
     // Intermediate result x changes, so final result depends on x
     // needs to be recomputed;
-    assert_eq!(final_result_deponds_on_x(&db, input), 24);
+    assert_eq!(final_result_depends_on_x(&db, input), 24);
     db.assert_logs(expect![[r#"
         [
-            "final_result_deponds_on_x(MyInput(Id { value: 1 }))",
+            "final_result_depends_on_x(MyInput(Id { value: 1 }))",
         ]"#]]);
 
     // y = 23 / 2 = 11
     // Intermediate result y is the same, so final result depends on y
     // does not need to be recomputed;
-    assert_eq!(final_result_deponds_on_y(&db, input), 22);
+    assert_eq!(final_result_depends_on_y(&db, input), 22);
     db.assert_logs(expect!["[]"]);
 }

salsa-2022-tests/tests/expect_reuse_field_x_of_an_input_changes_but_fn_depends_on_field_y.rs

@@ -8,7 +8,7 @@ use salsa_2022_tests::{HasLogger, Logger};
 use expect_test::expect;
 
 #[salsa::jar(db = Db)]
-struct Jar(MyInput, result_deponds_on_x, result_deponds_on_y);
+struct Jar(MyInput, result_depends_on_x, result_depends_on_y);
 
 trait Db: salsa::DbWithJar<Jar> + HasLogger {}
 
@@ -19,14 +19,14 @@ struct MyInput {
 }
 
 #[salsa::tracked(jar = Jar)]
-fn result_deponds_on_x(db: &dyn Db, input: MyInput) -> u32 {
-    db.push_log(format!("result_deponds_on_x({:?})", input));
+fn result_depends_on_x(db: &dyn Db, input: MyInput) -> u32 {
+    db.push_log(format!("result_depends_on_x({:?})", input));
     input.x(db) + 1
 }
 
 #[salsa::tracked(jar = Jar)]
-fn result_deponds_on_y(db: &dyn Db, input: MyInput) -> u32 {
-    db.push_log(format!("result_deponds_on_y({:?})", input));
+fn result_depends_on_y(db: &dyn Db, input: MyInput) -> u32 {
+    db.push_log(format!("result_depends_on_y({:?})", input));
     input.y(db) - 1
 }
 
@@ -53,33 +53,33 @@ impl HasLogger for Database {
 
 #[test]
 fn execute() {
-    // result_deponds_on_x = x + 1
-    // result_deponds_on_y = y - 1
+    // result_depends_on_x = x + 1
+    // result_depends_on_y = y - 1
     let mut db = Database::default();
 
     let input = MyInput::new(&mut db, 22, 33);
-    assert_eq!(result_deponds_on_x(&db, input), 23);
+    assert_eq!(result_depends_on_x(&db, input), 23);
     db.assert_logs(expect![[r#"
         [
-            "result_deponds_on_x(MyInput(Id { value: 1 }))",
+            "result_depends_on_x(MyInput(Id { value: 1 }))",
         ]"#]]);
 
-    assert_eq!(result_deponds_on_y(&db, input), 32);
+    assert_eq!(result_depends_on_y(&db, input), 32);
     db.assert_logs(expect![[r#"
         [
-            "result_deponds_on_y(MyInput(Id { value: 1 }))",
+            "result_depends_on_y(MyInput(Id { value: 1 }))",
         ]"#]]);
 
     input.set_x(&mut db, 23);
     // input x changes, so result depends on x needs to be recomputed;
-    assert_eq!(result_deponds_on_x(&db, input), 24);
+    assert_eq!(result_depends_on_x(&db, input), 24);
     db.assert_logs(expect![[r#"
         [
-            "result_deponds_on_x(MyInput(Id { value: 1 }))",
+            "result_depends_on_x(MyInput(Id { value: 1 }))",
         ]"#]]);
 
     // input y is the same, so result depends on y
     // does not need to be recomputed;
-    assert_eq!(result_deponds_on_y(&db, input), 32);
+    assert_eq!(result_depends_on_y(&db, input), 32);
     db.assert_logs(expect!["[]"]);
 }

salsa-2022-tests/tests/lru.rs

@@ -6,12 +6,13 @@ use std::sync::{
     Arc,
 };
 
+use salsa_2022_tests::{HasLogger, Logger};
 use test_log::test;
 
 #[salsa::jar(db = Db)]
-struct Jar(MyInput, get_hot_potato, get_volatile);
+struct Jar(MyInput, get_hot_potato, get_hot_potato2, get_volatile);
 
-trait Db: salsa::DbWithJar<Jar> {}
+trait Db: salsa::DbWithJar<Jar> + HasLogger {}
 
 #[derive(Debug, PartialEq, Eq)]
 struct HotPotato(u32);
@@ -40,9 +41,16 @@ struct MyInput {
 
 #[salsa::tracked(jar = Jar, lru = 32)]
 fn get_hot_potato(db: &dyn Db, input: MyInput) -> Arc<HotPotato> {
+    db.push_log(format!("get_hot_potato({:?})", input.field(db)));
     Arc::new(HotPotato::new(input.field(db)))
 }
 
+#[salsa::tracked(jar = Jar)]
+fn get_hot_potato2(db: &dyn Db, input: MyInput) -> u32 {
+    db.push_log(format!("get_hot_potato2({:?})", input.field(db)));
+    get_hot_potato(db, input).0
+}
+
 #[salsa::tracked(jar = Jar, lru = 32)]
 fn get_volatile(db: &dyn Db, _input: MyInput) -> usize {
     static COUNTER: AtomicUsize = AtomicUsize::new(0);
@@ -54,6 +62,7 @@ fn get_volatile(db: &dyn Db, _input: MyInput) -> usize {
 #[derive(Default)]
 struct Database {
     storage: salsa::Storage<Self>,
+    logger: Logger,
 }
 
 impl salsa::Database for Database {
@@ -64,6 +73,12 @@ impl salsa::Database for Database {
 
 impl Db for Database {}
 
+impl HasLogger for Database {
+    fn logger(&self) -> &Logger {
+        &self.logger
+    }
+}
+
 fn load_n_potatoes() -> usize {
     N_POTATOES.with(|n| n.load(Ordering::SeqCst))
 }
@@ -103,3 +118,73 @@ fn lru_doesnt_break_volatile_queries() {
         }
     }
 }
+
+#[test]
+fn lru_can_be_changed_at_runtime() {
+    let mut db = Database::default();
+    assert_eq!(load_n_potatoes(), 0);
+
+    let inputs: Vec<(u32, MyInput)> = (0..128).map(|i| (i, MyInput::new(&mut db, i))).collect();
+
+    for &(i, input) in inputs.iter() {
+        let p = get_hot_potato(&db, input);
+        assert_eq!(p.0, i)
+    }
+
+    // Create a new input to change the revision, and trigger the GC
+    MyInput::new(&mut db, 0);
+    assert_eq!(load_n_potatoes(), 32);
+
+    get_hot_potato::set_lru_capacity(&db, 64);
+    assert_eq!(load_n_potatoes(), 32);
+    for &(i, input) in inputs.iter() {
+        let p = get_hot_potato(&db, input);
+        assert_eq!(p.0, i)
+    }
+
+    // Create a new input to change the revision, and trigger the GC
+    MyInput::new(&mut db, 0);
+    assert_eq!(load_n_potatoes(), 64);
+
+    // Special case: setting capacity to zero disables LRU
+    get_hot_potato::set_lru_capacity(&db, 0);
+    assert_eq!(load_n_potatoes(), 64);
+    for &(i, input) in inputs.iter() {
+        let p = get_hot_potato(&db, input);
+        assert_eq!(p.0, i)
+    }
+
+    // Create a new input to change the revision, and trigger the GC
+    MyInput::new(&mut db, 0);
+    assert_eq!(load_n_potatoes(), 128);
+
+    drop(db);
+    assert_eq!(load_n_potatoes(), 0);
+}
+
+#[test]
+fn lru_keeps_dependency_info() {
+    let mut db = Database::default();
+    let capacity = 32;
+
+    // Invoke `get_hot_potato2` 33 times. This will (in turn) invoke
+    // `get_hot_potato`, which will trigger LRU after 32 executions.
+    let inputs: Vec<MyInput> = (0..(capacity + 1))
+        .map(|i| MyInput::new(&mut db, i as u32))
+        .collect();
+
+    for (i, input) in inputs.iter().enumerate() {
+        let x = get_hot_potato2(&db, *input);
+        assert_eq!(x as usize, i);
+    }
+
+    // We want to test that calls to `get_hot_potato2` are still considered
+    // clean. Check that no new executions occur as we go here.
+    db.assert_logs_len((capacity + 1) * 2);
+
+    // calling `get_hot_potato2(0)` has to check that `get_hot_potato(0)` is still valid;
+    // even though we've evicted it (LRU), we find that it is still good
+    let p = get_hot_potato2(&db, *inputs.first().unwrap());
+    assert_eq!(p, 0);
+    db.assert_logs_len(0);
+}