lru can be changed at runtime

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

@@ -288,6 +288,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 {
@@ -300,6 +301,8 @@ fn wrapper_fns(
             #[allow(dead_code, clippy::needless_lifetimes)]
             #accumulated_fn
 
+            #set_lru_fn
+
             #specify_fn
         }
     };
@@ -415,6 +418,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/tests/lru.rs

@@ -103,3 +103,47 @@ 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);
+
+    for i in 0..128u32 {
+        let input = MyInput::new(&mut db, i);
+        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 in 0..128u32 {
+        let input = MyInput::new(&mut db, i);
+        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 in 0..128u32 {
+        let input = MyInput::new(&mut db, i);
+        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(), 192);
+
+    drop(db);
+    assert_eq!(load_n_potatoes(), 0);
+}