diff --git a/src/function/memo.rs b/src/function/memo.rs index f982dc33..440fc2a7 100644 --- a/src/function/memo.rs +++ b/src/function/memo.rs @@ -61,30 +61,26 @@ impl IngredientImpl { /// 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; - }; - - match memo.revisions.origin { - QueryOrigin::Assigned(_) - | QueryOrigin::DerivedUntracked(_) - | QueryOrigin::BaseInput => { - // Careful: Cannot evict memos whose values were - // assigned as output of another query - // or those with untracked inputs - // as their values cannot be reconstructed. - } - - QueryOrigin::Derived(_) => { - let memo_evicted = Arc::new(Memo::new( - None::>, - memo.verified_at.load(), - memo.revisions.clone(), - )); - - self.insert_memo_into_table_for(zalsa, id, memo_evicted); - } - } + zalsa + .memo_table_for(id) + .map_memo::>(self.memo_ingredient_index, |memo| { + match memo.revisions.origin { + QueryOrigin::Assigned(_) + | QueryOrigin::DerivedUntracked(_) + | QueryOrigin::BaseInput => { + // Careful: Cannot evict memos whose values were + // assigned as output of another query + // or those with untracked inputs + // as their values cannot be reconstructed. + memo + } + QueryOrigin::Derived(_) => Arc::new(Memo::new( + None::>, + memo.verified_at.load(), + memo.revisions.clone(), + )), + } + }); } } diff --git a/src/table/memo.rs b/src/table/memo.rs index 73bd52bd..f05f7308 100644 --- a/src/table/memo.rs +++ b/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( + &self, + memo_ingredient_index: MemoIngredientIndex, + f: impl FnOnce(Arc) -> Arc, + ) { + // 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::(), + "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::(arc_swap.swap(Self::to_dummy(memo))) }; + } + pub(crate) fn into_memos( mut self, ) -> impl Iterator)> {