spine: expose in-progress merges for memory accounting

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>>,

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> {

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) } }