Reduce memo lookups needed for eviction

src/function/memo.rs

@@ -61,30 +61,26 @@ impl<C: Configuration> IngredientImpl<C> {
     /// with an equivalent memo that has no value. If the memo is untracked, BaseInput,
     /// or has values assigned as output of another query, this has no effect.
     pub(super) fn evict_value_from_memo_for<'db>(&'db self, zalsa: &'db Zalsa, id: Id) {
-        let Some(memo) = self.get_memo_from_table_for(zalsa, id) else {
+        zalsa
-            return;
+            .memo_table_for(id)
-        };
+            .map_memo::<Memo<_>>(self.memo_ingredient_index, |memo| {
-
+                match memo.revisions.origin {
-        match memo.revisions.origin {
+                    QueryOrigin::Assigned(_)
-            QueryOrigin::Assigned(_)
+                    | QueryOrigin::DerivedUntracked(_)
-            | QueryOrigin::DerivedUntracked(_)
+                    | QueryOrigin::BaseInput => {
-            | QueryOrigin::BaseInput => {
+                        // Careful: Cannot evict memos whose values were
-                // Careful: Cannot evict memos whose values were
+                        // assigned as output of another query
-                // assigned as output of another query
+                        // or those with untracked inputs
-                // or those with untracked inputs
+                        // as their values cannot be reconstructed.
-                // as their values cannot be reconstructed.
+                        memo
-            }
+                    }
-
+                    QueryOrigin::Derived(_) => Arc::new(Memo::new(
-            QueryOrigin::Derived(_) => {
+                        None::<C::Output<'_>>,
-                let memo_evicted = Arc::new(Memo::new(
+                        memo.verified_at.load(),
-                    None::<C::Output<'_>>,
+                        memo.revisions.clone(),
-                    memo.verified_at.load(),
+                    )),
-                    memo.revisions.clone(),
+                }
-                ));
+            });
-
-                self.insert_memo_into_table_for(zalsa, id, memo_evicted);
-            }
-        }
     }
 }
 

src/table/memo.rs

@@ -163,6 +163,37 @@ impl MemoTable {
         unsafe { Some(Self::from_dummy(arc_swap.load_full())) }
     }
 
+    /// Calls `f` on the memo at `memo_ingredient_index` and replaces the memo with the result of `f`.
+    /// If the memo is not present, `f` is not called.
+    pub(crate) fn map_memo<M: Memo>(
+        &self,
+        memo_ingredient_index: MemoIngredientIndex,
+        f: impl FnOnce(Arc<M>) -> Arc<M>,
+    ) {
+        // If the memo slot is already occupied, it must already have the
+        // right type info etc, and we only need the read-lock.
+        let memos = self.memos.read();
+        let Some(MemoEntry {
+            data:
+                Some(MemoEntryData {
+                    type_id,
+                    to_dyn_fn: _,
+                    arc_swap,
+                }),
+        }) = memos.get(memo_ingredient_index.as_usize())
+        else {
+            return;
+        };
+        assert_eq!(
+            *type_id,
+            TypeId::of::<M>(),
+            "inconsistent type-id for `{memo_ingredient_index:?}`"
+        );
+        // SAFETY: type_id check asserted above
+        let memo = f(unsafe { Self::from_dummy(arc_swap.load_full()) });
+        unsafe { Self::from_dummy::<M>(arc_swap.swap(Self::to_dummy(memo))) };
+    }
+
     pub(crate) fn into_memos(
         mut self,
     ) -> impl Iterator<Item = (MemoIngredientIndex, Arc<dyn Memo>)> {