From 26257c4a5c6a09fb29c5423c9ba656b221f3d2df Mon Sep 17 00:00:00 2001
From: Moritz Hoffmann <antiguru@gmail.com>
Date: Tue, 1 Sep 2026 13:14:11 +0200
Subject: [PATCH] spine: expose in-progress merges for memory accounting

A consumer that accounts for a spine's memory footprint by walking
`TraceReader::map_batches` undercounts it. For a layer mid-merge, that method
presents the merge's two input batches and stops, but the layer holds a third
allocation: the merger's partially assembled output. `Merger::new` sizes that
output's containers with `BatchContainer::merge_capacity`, so the allocation
appears in full the moment the merge begins and is invisible for the merge's
whole duration.

Add `Spine::map_mergers`, which applies a closure to the merger of each layer
that is currently mid-merge. It is an inherent method rather than a `TraceReader`
one because `TraceReader` has no `Merger` associated type, and its `Batch` is
bounded by `BatchReader` rather than `Batch`, so the merger type is not nameable
at that level.

Reaching a merger is only useful if its contents can be inspected, so also add
`OrdValMerger::result` and `OrdKeyMerger::result`, which borrow the storage being
assembled, and `RcMerger::inner`, without which `map_mergers` on the default
`Rc`-backed spines yields an opaque wrapper.

All three additions are accessors. No existing signature, trait, or behavior
changes.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
---
 .../src/trace/implementations/ord_neu.rs      | 22 +++++++++++++++++++
 .../src/trace/implementations/spine_fueled.rs | 17 ++++++++++++++
 differential-dataflow/src/trace/mod.rs        |  5 +++++
 3 files changed, 44 insertions(+)

diff --git a/differential-dataflow/src/trace/implementations/ord_neu.rs b/differential-dataflow/src/trace/implementations/ord_neu.rs
index cf0c22a5f..03322643f 100644
--- a/differential-dataflow/src/trace/implementations/ord_neu.rs
+++ b/differential-dataflow/src/trace/implementations/ord_neu.rs
@@ -358,6 +358,17 @@ pub mod val_batch {
         staging: UpdsBuilder<L::TimeContainer, L::DiffContainer>,
     }
 
+    impl<L: Layout> OrdValMerger<L> {
+        /// The partially assembled merge output.
+        ///
+        /// The containers are allocated at their full merged capacity when the merge begins, so
+        /// this storage accounts for the merge's whole memory cost from the outset, even though
+        /// the number of updates in it grows as the merge proceeds.
+        pub fn result(&self) -> &OrdValStorage<L> {
+            &self.result
+        }
+    }
+
     impl<L: Layout> Merger<OrdValBatch<L>> for OrdValMerger<L>
     where
         OrdValBatch<L>: Batch<Time=layout::Time<L>>,
@@ -860,6 +871,17 @@ pub mod key_batch {
         staging: UpdsBuilder<L::TimeContainer, L::DiffContainer>,
     }
 
+    impl<L: Layout> OrdKeyMerger<L> {
+        /// The partially assembled merge output.
+        ///
+        /// The containers are allocated at their full merged capacity when the merge begins, so
+        /// this storage accounts for the merge's whole memory cost from the outset, even though
+        /// the number of updates in it grows as the merge proceeds.
+        pub fn result(&self) -> &OrdKeyStorage<L> {
+            &self.result
+        }
+    }
+
     impl<L: Layout<ValContainer: BatchContainer<Owned: Default>>> Merger<OrdKeyBatch<L>> for OrdKeyMerger<L>
     where
         OrdKeyBatch<L>: Batch<Time=layout::Time<L>>,
diff --git a/differential-dataflow/src/trace/implementations/spine_fueled.rs b/differential-dataflow/src/trace/implementations/spine_fueled.rs
index f4e6ad036..b5e826a78 100644
--- a/differential-dataflow/src/trace/implementations/spine_fueled.rs
+++ b/differential-dataflow/src/trace/implementations/spine_fueled.rs
@@ -223,6 +223,23 @@ impl<B: Batch+Clone+'static> TraceReader for Spine<B> {
     }
 }
 
+impl<B: Batch> Spine<B> {
+    /// Applies `f` to the merger of each layer that is currently mid-merge.
+    ///
+    /// A merger owns the partially assembled output of its merge. That storage is memory the
+    /// spine holds but which [`TraceReader::map_batches`] does not reach: it presents the two
+    /// input batches of a merge and stops there. A caller that accounts for the spine's memory
+    /// footprint must visit both, and must re-read a merger on each observation rather than
+    /// cache what it learns, because a merger's contents change as the merge proceeds.
+    pub fn map_mergers<F: FnMut(&<B as Batch>::Merger)>(&self, mut f: F) {
+        for state in self.merging.iter().rev() {
+            if let MergeState::Double(MergeVariant::InProgress(_, _, merger)) = state {
+                f(merger);
+            }
+        }
+    }
+}
+
 // A trace implementation for any key type that can be borrowed from or converted into `Key`.
 // TODO: Almost all this implementation seems to be generic with respect to the trace and batch types.
 impl<B: Batch+Clone+'static> Trace for Spine<B> {
diff --git a/differential-dataflow/src/trace/mod.rs b/differential-dataflow/src/trace/mod.rs
index c6f864f53..1f6600411 100644
--- a/differential-dataflow/src/trace/mod.rs
+++ b/differential-dataflow/src/trace/mod.rs
@@ -430,6 +430,11 @@ pub mod rc_blanket_impls {
     /// Wrapper type for merging reference counted batches.
     pub struct RcMerger<B:Batch> { merger: B::Merger }
 
+    impl<B:Batch> RcMerger<B> {
+        /// The wrapped batch's merger, which owns the merge's partially assembled output.
+        pub fn inner(&self) -> &B::Merger { &self.merger }
+    }
+
     /// Represents a merge in progress.
     impl<B:Batch> Merger<Rc<B>> for RcMerger<B> {
         fn new(source1: &Rc<B>, source2: &Rc<B>, compaction_frontier: AntichainRef<B::Time>) -> Self { RcMerger { merger: B::begin_merge(source1, source2, compaction_frontier) } }