Reduce memo lookups needed for eviction

src/function/memo.rs

@@ -61,10 +61,9 @@ 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 {
-            return;
-        };
-
+        zalsa
+            .memo_table_for(id)
+            .map_memo::<Memo<_>>(self.memo_ingredient_index, |memo| {
                 match memo.revisions.origin {
                     QueryOrigin::Assigned(_)
                     | QueryOrigin::DerivedUntracked(_)
@@ -73,18 +72,15 @@ impl<C: Configuration> IngredientImpl<C> {
                         // assigned as output of another query
                         // or those with untracked inputs
                         // as their values cannot be reconstructed.
+                        memo
                     }
-
-            QueryOrigin::Derived(_) => {
-                let memo_evicted = Arc::new(Memo::new(
+                    QueryOrigin::Derived(_) => Arc::new(Memo::new(
                         None::<C::Output<'_>>,
                         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>)> {