refactor(state_store): per-component handle for fn-memo I/O (#1990)

- Introduce `FnMemoAccessor`, owning the read / write / retain operations
  for `FunctionMemoization` entries scoped to a single component path.
- Build it once per `ComponentProcessorContext` and return it by
  reference, so any future per-build state on the accessor (e.g. caching
  or instrumentation) persists across the many fn-memo lookups within a
  single component's processing phase, rather than being re-created per
  call.
- Inline the corresponding `AppStore::{read,write,retain}_fn_memo*`
  methods into the accessor; they had no remaining callers outside it.
  Shrinks `AppStore`'s public surface by ~50 lines.
- Update the three engine call sites (cache-probe read in
  `read_fn_call_memo_with_txn`, post-execute write in
  `write_fn_call_memo`, commit-time retain GC in `Committer`) to route
  through `comp_ctx.fn_memo_accessor()`.

No behavior change.

Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Author: georgeh0

rust/core/src/engine/context.rs

@@ -18,7 +18,7 @@ use crate::state::target_state_path::TargetStatePath;
 use crate::{
     engine::environment::{AppRegistration, Environment},
     state::stable_path::StablePath,
-    state_store::AppStore,
+    state_store::{AppStore, FnMemoAccessor},
 };
 
 struct AppContextInner<Prof: EngineProfile> {
@@ -283,6 +283,13 @@ struct ComponentProcessorContextInner<Prof: EngineProfile> {
 
     /// Opaque per-operation context (e.g. ContextProvider on the Python side).
     host_ctx: Arc<Prof::HostCtx>,
+
+    /// Per-component handle for function-memo I/O. Built once when the
+    /// context is constructed and shared by reference, so any per-build
+    /// state inside the accessor (e.g. a buffer in a future backend)
+    /// persists across the many fn-memo lookups within one component's
+    /// processing phase.
+    fn_memo_accessor: FnMemoAccessor,
 }
 
 #[derive(Clone)]
@@ -298,6 +305,10 @@ impl<Prof: EngineProfile> ComponentProcessorContext<Prof> {
         host_ctx: Arc<Prof::HostCtx>,
         processing_action: ComponentProcessingAction<Prof>,
     ) -> Self {
+        let fn_memo_accessor = FnMemoAccessor::new(
+            component.app_ctx().app_store().clone(),
+            component.stable_path().clone(),
+        );
         Self {
             inner: Arc::new(ComponentProcessorContextInner {
                 component,
@@ -308,6 +319,7 @@ impl<Prof: EngineProfile> ComponentProcessorContext<Prof> {
                 inflight_permit: Mutex::new(None),
                 logic_deps: Mutex::new(HashSet::new()),
                 host_ctx,
+                fn_memo_accessor,
             }),
         }
     }
@@ -328,6 +340,14 @@ impl<Prof: EngineProfile> ComponentProcessorContext<Prof> {
         self.inner.component.stable_path()
     }
 
+    /// Per-component handle for function-memoization I/O. All engine code
+    /// that reads, writes, or retains function memos under this component's
+    /// path should go through this handle rather than calling `AppStore`
+    /// methods directly — see [`FnMemoAccessor`] for why.
+    pub fn fn_memo_accessor(&self) -> &FnMemoAccessor {
+        &self.inner.fn_memo_accessor
+    }
+
     pub(crate) fn parent_context(&self) -> Option<&ComponentProcessorContext<Prof>> {
         self.inner.parent_context.as_ref()
     }

rust/core/src/engine/execution.rs

@@ -199,7 +199,6 @@ pub(crate) async fn update_component_memo_states<Prof: EngineProfile>(
 
 async fn write_fn_call_memo<Prof: EngineProfile>(
     wtxn: &mut WriteTxn<'_>,
-    app_store: &AppStore,
     comp_ctx: &ComponentProcessorContext<Prof>,
     memo_fp: Fingerprint,
     memo: FnCallMemo<Prof>,
@@ -217,21 +216,18 @@ async fn write_fn_call_memo<Prof: EngineProfile>(
         memo_states: memo_states_serialized,
         context_memo_states: context_memo_states_serialized,
     };
-    app_store
-        .write_fn_memo(wtxn, comp_ctx.stable_path(), memo_fp, &fn_call_memo)
+    comp_ctx
+        .fn_memo_accessor()
+        .write(wtxn, memo_fp, &fn_call_memo)
         .await
 }
 
 async fn read_fn_call_memo_with_txn<Prof: EngineProfile, T: AnyTxn>(
     rtxn: &mut T,
-    app_store: &AppStore,
     comp_ctx: &ComponentProcessorContext<Prof>,
     memo_fp: Fingerprint,
 ) -> Result<Option<FnCallMemo<Prof>>> {
-    let Some(bytes) = app_store
-        .read_fn_memo(rtxn, comp_ctx.stable_path(), memo_fp)
-        .await?
-    else {
+    let Some(bytes) = comp_ctx.fn_memo_accessor().read(rtxn, memo_fp).await? else {
         return Ok(None);
     };
     let fn_call_memo: db_schema::FunctionMemoizationEntry<'_> = from_msgpack_slice(&bytes)?;
@@ -270,7 +266,7 @@ pub(crate) async fn read_fn_call_memo<Prof: EngineProfile>(
         return Ok(None);
     }
     let mut rtxn = comp_ctx.app_ctx().env().read_txn().await?;
-    read_fn_call_memo_with_txn(&mut rtxn, comp_ctx.app_ctx().app_store(), comp_ctx, memo_fp).await
+    read_fn_call_memo_with_txn(&mut rtxn, comp_ctx, memo_fp).await
 }
 
 pub fn declare_target_state<Prof: EngineProfile>(
@@ -480,12 +476,13 @@ impl<Prof: EngineProfile> Committer<Prof> {
 
             // Write memos.
             for (fp, memo) in memos_without_mounts_to_store {
-                write_fn_call_memo(wtxn, &self.app_store, &self.component_ctx, fp, memo).await?;
+                write_fn_call_memo(wtxn, &self.component_ctx, fp, memo).await?;
             }
 
             // Delete all function memo entries that are not in the all_memo_fps.
-            self.app_store
-                .retain_fn_memos(wtxn, &self.component_path, all_memo_fps)
+            self.component_ctx
+                .fn_memo_accessor()
+                .retain(wtxn, all_memo_fps)
                 .await?;
 
             if !self.demote_component_only {
@@ -1496,13 +1493,11 @@ async fn finalize_fn_call_memoization<Prof: EngineProfile>(
     // Use a single read transaction for all DB reads.
     if !deps_to_process.is_empty() {
         let mut rtxn = comp_ctx.app_ctx().env().read_txn().await?;
-        let app_store = comp_ctx.app_ctx().app_store();
         while let Some(fp) = deps_to_process.pop_front() {
             if !result.all_memos_fps.insert(fp) {
                 continue;
             }
-            let Some(memo) = read_fn_call_memo_with_txn(&mut rtxn, app_store, comp_ctx, fp).await?
-            else {
+            let Some(memo) = read_fn_call_memo_with_txn(&mut rtxn, comp_ctx, fp).await? else {
                 continue;
             };
             result

rust/core/src/state_store/app_store.rs

@@ -181,44 +181,68 @@ impl AppStore {
     }
 }
 
-// --- Function memoization ------------------------------------------------
+// --- Function memoization (per-component, via FnMemoAccessor) -----------
+
+/// Per-component handle that mediates function-memoization reads, writes,
+/// and retention during a single component build.
+///
+/// Engine code routes all per-component function-memo I/O through this type
+/// rather than calling [`AppStore`] methods directly, so storage backends
+/// that benefit from prefetching (e.g. a future Postgres backend that wants
+/// to load all entries for a component in one prefix scan and serve all
+/// reads from memory) can introduce a per-build buffer here without
+/// touching the engine call sites.
+///
+/// The accessor is constructed once per component build and held by
+/// [`ComponentProcessorContext`](crate::engine::context::ComponentProcessorContext)
+/// — engine code reaches it via `comp_ctx.fn_memo_accessor()`, which
+/// returns a borrow rather than a fresh value, so any per-build state
+/// (a future buffer) persists across the many fn-memo lookups inside a
+/// single component's processing phase.
+///
+/// The current LMDB implementation is a pure passthrough — every call
+/// delegates directly to the corresponding [`AppStore`] method, since LMDB's
+/// random reads are already memory speed.
+pub struct FnMemoAccessor {
+    app_store: AppStore,
+    component_path: StablePath,
+}
 
-impl AppStore {
-    /// Read raw function-memo bytes. See [`Self::read_tracking_info`] for
-    /// the owned-bytes return rationale.
-    pub async fn read_fn_memo<T: AnyTxn>(
-        &self,
-        txn: &mut T,
-        path: &StablePath,
-        fp: Fingerprint,
-    ) -> Result<Option<Vec<u8>>> {
-        let key = key_fn_memo(path, fp)?;
-        Ok(txn.db_get_bytes(self.db(), &key)?.map(<[u8]>::to_vec))
+impl FnMemoAccessor {
+    pub fn new(app_store: AppStore, component_path: StablePath) -> Self {
+        Self {
+            app_store,
+            component_path,
+        }
     }
 
-    pub async fn write_fn_memo(
+    /// Read raw function-memo bytes for the given fingerprint. Returns
+    /// owned bytes; see [`AppStore::read_tracking_info`] for the rationale.
+    pub async fn read<T: AnyTxn>(&self, txn: &mut T, fp: Fingerprint) -> Result<Option<Vec<u8>>> {
+        let key = key_fn_memo(&self.component_path, fp)?;
+        Ok(txn
+            .db_get_bytes(self.app_store.db(), &key)?
+            .map(<[u8]>::to_vec))
+    }
+
+    pub async fn write(
         &self,
-        txn: &mut WriteTxn<'_>,
-        path: &StablePath,
+        wtxn: &mut WriteTxn<'_>,
         fp: Fingerprint,
         entry: &FunctionMemoizationEntry<'_>,
     ) -> Result<()> {
-        let key = key_fn_memo(path, fp)?;
+        let key = key_fn_memo(&self.component_path, fp)?;
         let value = rmp_serde::to_vec_named(entry)?;
-        self.db().put(&mut **txn, &key, &value)?;
+        self.app_store.db().put(&mut **wtxn, &key, &value)?;
         Ok(())
     }
 
-    /// GC: delete all function memos for `path` whose fingerprint is NOT in `keep`.
-    pub async fn retain_fn_memos(
-        &self,
-        txn: &mut WriteTxn<'_>,
-        path: &StablePath,
-        keep: &HashSet<Fingerprint>,
-    ) -> Result<()> {
-        let prefix = key_fn_memo_prefix(path)?;
-        let db = self.db();
-        let mut iter = db.prefix_iter_mut(&mut **txn, &prefix)?;
+    /// GC: delete all function memos for this component whose fingerprint
+    /// is NOT in `keep`.
+    pub async fn retain(&self, wtxn: &mut WriteTxn<'_>, keep: &HashSet<Fingerprint>) -> Result<()> {
+        let prefix = key_fn_memo_prefix(&self.component_path)?;
+        let db = self.app_store.db();
+        let mut iter = db.prefix_iter_mut(&mut **wtxn, &prefix)?;
         while let Some((raw_key, _)) = iter.next().transpose()? {
             let fp: Fingerprint = storekey::decode(raw_key[prefix.len()..].as_ref())?;
             if keep.contains(&fp) {

rust/core/src/state_store/mod.rs

@@ -12,6 +12,6 @@ mod app_store;
 mod storage;
 mod txn;
 
-pub use app_store::AppStore;
+pub use app_store::{AppStore, FnMemoAccessor};
 pub use storage::{Storage, StorageSettings};
 pub use txn::{AnyTxn, ReadTxn, WriteTxn};