refactor: one axis_row_payload for every writer and checker (#809 audit follow-up)

grovedb/src/batch/indexed_tree/mirror.rs

@@ -24,26 +24,10 @@ use grovedb_version::version::GroveVersion;
 
 use super::{read_entry_aggregates, AggregatePair, AggregateTransition};
 use crate::{
-    operations::indexed_tree::{count_value_as_sum, make_axis_secondary_key},
+    operations::indexed_tree::{axis_row_payload, make_axis_secondary_key},
     Element, Error,
 };
 
-/// The per-axis payload stored alongside the sort key. The key encodes the
-/// ordering value; the payload carries what the secondary's own aggregate
-/// must sum to, which is why NO axis can store a bare item: every
-/// secondary is a dual-aggregate `ProvableCountProvableSumTree`, and the
-/// count axis mirrors its `count_value` into the sum half so a band
-/// TOTAL is one committed scalar (issue #806).
-///
-/// Fallible only through [`count_value_as_sum`]'s fail-closed guard.
-fn axis_row_payload(axis: IndexAxis, count: u64, sum: i64) -> Result<Element, Error> {
-    Ok(match axis {
-        IndexAxis::Count => Element::new_sum_item(count_value_as_sum(count)?),
-        IndexAxis::Sum => Element::new_sum_item(sum),
-        IndexAxis::Avg => Element::new_item_with_sum_item(Vec::new(), sum),
-    })
-}
-
 /// Enforce the precise per-axis item-key bound: avg prepends a 16-byte sort
 /// key, count and sum 8, so the same item key can be legal on one axis and
 /// not another. Runs before any secondary write; erroring here aborts the

grovedb/src/estimated_costs/average_case_costs.rs

@@ -32,7 +32,7 @@ use crate::{
 };
 
 /// Upper bound on an indexed secondary row's value. The payload is fixed
-/// per axis — an empty `Item` (count), a `SumItem` (sum) or an empty
+/// per axis — a `SumItem` (count and sum axes) or an empty
 /// `ItemWithSumItem` (avg) — and the largest of those serializes well under
 /// this bound, which also leaves room for the feature type and flags byte.
 pub const INDEXED_SECONDARY_MAX_VALUE_SIZE: u32 = 16;
@@ -482,7 +482,7 @@ impl GroveDb {
     /// - **Key size** is the primary's key size plus the axis sort-key width
     ///   (8 bytes for count/sum, 16 for avg).
     /// - **Value size** is bounded by the fixed per-axis payload shape:
-    ///   an empty `Item` (count), a `SumItem` (sum), or an empty
+    ///   a `SumItem` (count and sum axes), or an empty
     ///   `ItemWithSumItem` (avg) — all under
     ///   [`INDEXED_SECONDARY_MAX_VALUE_SIZE`].
     /// - **Tree type** is fixed per axis.
@@ -542,20 +542,19 @@ impl GroveDb {
             //
             // The inserted row must be sized with the AXIS's real payload
             // shape: `average_case_merk_insert_element` charges non-tree
-            // elements by their own serialized size, and the sum and avg rows
-            // are larger than the count axis's empty `Item`. Sizing all three
-            // as an empty `Item` put the PCPSIT estimate ~25 bytes per key
-            // UNDER actual `added_bytes` — the one dimension a storage-fee
-            // reservation cannot come in under. Sum values are charged at
-            // their fixed worst-case varint width, so `i64::MAX` here is the
-            // upper bound, not an average.
-            let worst_case_row = match axis {
-                grovedb_element::indexed::IndexAxis::Count => Element::new_sum_item(i64::MAX),
-                grovedb_element::indexed::IndexAxis::Sum => Element::new_sum_item(i64::MAX),
-                grovedb_element::indexed::IndexAxis::Avg => {
-                    Element::new_item_with_sum_item(vec![], i64::MAX)
-                }
-            };
+            // elements by their own serialized size, and under-sizing put
+            // the PCPSIT estimate ~25 bytes per key UNDER actual
+            // `added_bytes` — the one dimension a storage-fee reservation
+            // cannot come in under. The shape comes from THE payload
+            // function the mirror writes with (`axis_row_payload`), fed
+            // worst-case aggregates: sum values are charged at their fixed
+            // worst-case varint width, so `i64::MAX` is the upper bound,
+            // not an average (and is in `count_value_as_sum`'s domain, so
+            // the conversion cannot fail here).
+            let worst_case_row = cost_return_on_error_no_add!(
+                cost,
+                crate::operations::indexed_tree::axis_row_payload(*axis, i64::MAX as u64, i64::MAX,)
+            );
             cost_return_on_error!(
                 &mut cost,
                 Self::average_case_merk_delete_element(

grovedb/src/lib.rs

@@ -1691,15 +1691,7 @@ impl GroveDb {
             // The payload the mirror writes for this axis, so a row filed
             // under the right key but carrying the wrong value is caught too
             // — the key alone does not pin the stored aggregate.
-            let payload = match axis {
-                grovedb_element::indexed::IndexAxis::Count => Element::new_sum_item(
-                    crate::operations::indexed_tree::count_value_as_sum(count)?,
-                ),
-                grovedb_element::indexed::IndexAxis::Sum => Element::new_sum_item(sum),
-                grovedb_element::indexed::IndexAxis::Avg => {
-                    Element::new_item_with_sum_item(Vec::new(), sum)
-                }
-            };
+            let payload = crate::operations::indexed_tree::axis_row_payload(axis, count, sum)?;
             expected.insert(
                 p_key.clone(),
                 (make_axis_secondary_key(axis, count, sum, &p_key), payload),

grovedb/src/operations/indexed_tree.rs

@@ -60,6 +60,19 @@ use crate::{util::TxRef, Element, Error, GroveDb, Transaction, TransactionArg};
 /// proof primitive skips whole subtrees via counted node commitments),
 /// and the sum half is what makes a TOTAL over a value band answerable
 /// as one committed scalar ([`AggregateFold::Total`], issue #806).
+///
+/// **Dual-aggregate is a SECURITY requirement here, not a
+/// convenience.** The single-aggregate node hashes share a preimage
+/// layout (`node_hash_with_count` and `node_hash_with_sum` are both
+/// `Blake3(kv ‖ l ‖ r ‖ 8 bytes)` with no domain tag), so on a
+/// single-aggregate secondary a count proof could be node-type
+/// relabeled into a byte-different "sum" proof reconstructing the
+/// IDENTICAL root — making a band Total forgeable from a Population
+/// proof at zero cost. Dual-aggregate nodes hash a 112-byte preimage
+/// (`count_be8 ‖ sum_be8`), which closes the rewrite. Do not
+/// "optimize" any axis back to a single-aggregate tree type without
+/// first adding domain separation to the node-hash functions
+/// (per the #809 security audit, finding C).
 #[inline]
 pub(crate) fn axis_secondary_tree_type(axis: IndexAxis) -> TreeType {
     match axis {
@@ -90,6 +103,35 @@ pub(crate) fn count_value_as_sum(count: u64) -> Result<i64, Error> {
     })
 }
 
+/// THE per-axis secondary row payload — the single definition every
+/// writer and every checker uses. The sort KEY encodes the ordering
+/// value; this payload carries what the secondary's own dual
+/// aggregates must fold to:
+///
+/// - Count → `SumItem(count_value)` — contributes `(1, count)`, so a
+///   band TOTAL is one committed scalar (issue #806)
+/// - Sum   → `SumItem(sum)` — contributes `(1, sum)`
+/// - Avg   → `ItemWithSumItem(empty, sum)` — contributes `(1, sum)`
+///
+/// Callers: the batch mirror row builder, the direct-path mirror, the
+/// reconcile repair loop, `verify_grovedb`'s expected-payload check,
+/// and the average-case cost estimator's worst-case row. They MUST all
+/// go through this function: the mirror writes these bytes into
+/// hash-committed state and the checkers recompute them independently,
+/// so a divergent copy at any site either false-flags healthy state or
+/// makes two entry points commit different root hashes for identical
+/// writes. (This function exists because exactly that drift risk was
+/// flagged by the #809 security audit.)
+///
+/// Fallible only through [`count_value_as_sum`]'s fail-closed guard.
+pub(crate) fn axis_row_payload(axis: IndexAxis, count: u64, sum: i64) -> Result<Element, Error> {
+    Ok(match axis {
+        IndexAxis::Count => Element::new_sum_item(count_value_as_sum(count)?),
+        IndexAxis::Sum => Element::new_sum_item(sum),
+        IndexAxis::Avg => Element::new_item_with_sum_item(Vec::new(), sum),
+    })
+}
+
 /// Build the secondary key bytes for an entry at `item_key` under the
 /// given axis, given the relevant aggregate values:
 /// - count axis → `count_be(8) ‖ item_key`
@@ -1041,14 +1083,8 @@ impl GroveDb {
                 std::collections::BTreeMap::new();
             for (key, (count, sum)) in &entries {
                 let secondary_key = make_axis_secondary_key(axis, *count, *sum, key);
-                let payload = match axis {
-                    IndexAxis::Count => Element::new_sum_item(cost_return_on_error_no_add!(
-                        cost,
-                        count_value_as_sum(*count)
-                    )),
-                    IndexAxis::Sum => Element::new_sum_item(*sum),
-                    IndexAxis::Avg => Element::new_item_with_sum_item(Vec::new(), *sum),
-                };
+                let payload =
+                    cost_return_on_error_no_add!(cost, axis_row_payload(axis, *count, *sum));
                 let payload_bytes = cost_return_on_error_no_add!(
                     cost,
                     payload.serialize(grove_version).map_err(|e| {
@@ -2393,15 +2429,8 @@ impl GroveDb {
 /// from the prior primary state and `new_count`/`new_sum` from the
 /// post-mutation state.
 ///
-/// The secondary entry is a no-payload `Item` whose own sum / count
-/// contribution comes from its position in a sum/count-bearing tree.
-/// For the sum axis the secondary entry is a `SumItem(sum)`, which in
-/// the secondary's `ProvableCountProvableSumTree` contributes
-/// (count = 1, sum); for the avg axis the secondary entry is an
-/// `ItemWithSumItem(empty, sum)` so both count (= 1) and sum (= the
-/// entry's sum_value) propagate to the secondary's
-/// `ProvableCountProvableSumTree`. For the count axis the secondary
-/// entry is a plain `Item` (count = 1, no sum).
+/// The row payload is [`axis_row_payload`] — see its doc for the
+/// per-axis shapes and why every writer must share it.
 #[allow(clippy::too_many_arguments)]
 pub(crate) fn mirror_indexed_axis_to_secondary<'db, S: StorageContext<'db>>(
     secondary: &mut Merk<S>,
@@ -2416,22 +2445,11 @@ pub(crate) fn mirror_indexed_axis_to_secondary<'db, S: StorageContext<'db>>(
     let mut cost = OperationCost::default();
     let secondary_tree_type = axis_secondary_tree_type(axis);
 
-    // The per-axis secondary payload for an entry with the given
-    // aggregate values. Kept in lockstep with the insert branch below:
-    // - Count → constant empty Item (contributes count = 1)
-    // - Sum   → SumItem(sum)
-    // - Avg   → ItemWithSumItem(empty, sum) (contributes (1, sum))
-    // The payload is a function of BOTH aggregates now: the count axis
-    // mirrors count_value into its sum half, so a closure over `sum`
-    // alone could not see a count-only change and the fast path below
-    // would silently skip a real payload update.
-    let axis_payload = |count: u64, sum: i64| -> Result<Element, Error> {
-        Ok(match axis {
-            IndexAxis::Count => Element::new_sum_item(count_value_as_sum(count)?),
-            IndexAxis::Sum => Element::new_sum_item(sum),
-            IndexAxis::Avg => Element::new_item_with_sum_item(Vec::new(), sum),
-        })
-    };
+    // The payload is a function of BOTH aggregates (the count axis
+    // mirrors count_value into its sum half), so the fast-path equality
+    // check below must see the full pair — a sum-only view would
+    // silently skip real payload updates.
+    let axis_payload = |count: u64, sum: i64| axis_row_payload(axis, count, sum);
 
     // Compute old and new sort keys. Either may be None (no entry).
     let old_key = match (old_count, old_sum) {
@@ -2471,7 +2489,18 @@ pub(crate) fn mirror_indexed_axis_to_secondary<'db, S: StorageContext<'db>>(
         }
     }
 
-    if let Some(ok) = &old_key {
+    // A payload change at a FIXED key (avg axis only: a proportional
+    // (count, sum) change keeps the average) must replace in place.
+    // Deleting and reinserting the same key rebalances the AVL twice
+    // and can settle a DIFFERENT shape than the batch mirror's single
+    // replacement write — two write paths committing different
+    // secondary (hence grove) root hashes for identical data, which
+    // `direct_and_batch_agree_on_root_for_a_fixed_key_avg_payload_change`
+    // reproduced on an interior node before this skip existed. The
+    // insert below overwrites the value in place, exactly like the
+    // batch path's put.
+    let key_moved = old_key != new_key;
+    if let (true, Some(ok)) = (key_moved, &old_key) {
         cost_return_on_error!(
             &mut cost,
             Element::delete(
@@ -2614,6 +2643,85 @@ fn corrupted_secondary_key_error(axis: IndexAxis, secondary_key: &[u8]) -> Error
     ))
 }
 
+#[cfg(test)]
+mod axis_row_payload_tests {
+    //! The payload function is THE definition every writer and checker
+    //! shares; this grid pins its output per (axis, count, sum) so any
+    //! change to the shape is a deliberate, reviewed event — the bytes
+    //! land in hash-committed state, so an accidental change here means
+    //! mirrors and checkers disagree about healthy databases.
+
+    use grovedb_element::indexed::IndexAxis;
+    use grovedb_version::version::GroveVersion;
+
+    use super::axis_row_payload;
+    use crate::Element;
+
+    #[test]
+    fn payload_grid_is_pinned_per_axis() {
+        let counts = [0u64, 1, 2, i64::MAX as u64];
+        let sums = [i64::MIN, -1, 0, 1, i64::MAX];
+        for &count in &counts {
+            for &sum in &sums {
+                assert_eq!(
+                    axis_row_payload(IndexAxis::Count, count, sum).unwrap(),
+                    Element::new_sum_item(count as i64),
+                    "count axis stores the COUNT as its sum item; the sum input is ignored"
+                );
+                assert_eq!(
+                    axis_row_payload(IndexAxis::Sum, count, sum).unwrap(),
+                    Element::new_sum_item(sum),
+                    "sum axis stores the sum; the count input is ignored"
+                );
+                assert_eq!(
+                    axis_row_payload(IndexAxis::Avg, count, sum).unwrap(),
+                    Element::new_item_with_sum_item(Vec::new(), sum),
+                    "avg axis stores an empty item carrying the sum"
+                );
+            }
+        }
+        // Above the sum-item domain the count axis fails closed.
+        axis_row_payload(IndexAxis::Count, i64::MAX as u64 + 1, 0)
+            .expect_err("count above i64::MAX must fail closed");
+    }
+
+    #[test]
+    fn serialized_bytes_are_stable() {
+        // The exact bytes the mirror hash-commits, pinned as FIXED
+        // vectors — not re-serialized at assertion time, so a
+        // serialization-format change cannot move both sides of the
+        // comparison and slip through. If this test fails, payload
+        // bytes in authenticated state have changed: that is a
+        // consensus event, not a refactor.
+        let grove_version = GroveVersion::latest();
+        assert_eq!(
+            axis_row_payload(IndexAxis::Count, 7, 0)
+                .unwrap()
+                .serialize(grove_version)
+                .unwrap(),
+            vec![3, 14, 0],
+            "count axis: SumItem(7) as [variant, zigzag-varint 7, no flags]"
+        );
+        assert_eq!(
+            axis_row_payload(IndexAxis::Sum, 1, -3)
+                .unwrap()
+                .serialize(grove_version)
+                .unwrap(),
+            vec![3, 5, 0],
+            "sum axis: SumItem(-3) as [variant, zigzag-varint -3, no flags]"
+        );
+        assert_eq!(
+            axis_row_payload(IndexAxis::Avg, 1, 5)
+                .unwrap()
+                .serialize(grove_version)
+                .unwrap(),
+            vec![9, 0, 10, 0],
+            "avg axis: ItemWithSumItem(empty, 5) as [variant, empty item, \
+             zigzag-varint 5, no flags]"
+        );
+    }
+}
+
 #[cfg(test)]
 mod count_value_as_sum_tests {
     //! The count-axis secondary stores count_value as an i64 sum item;

grovedb/src/operations/proof/indexed_axis/generate.rs

@@ -1318,8 +1318,7 @@ impl GroveDb {
             AxisTraversal::AggregateOverValueRange { lo, hi, fold } => {
                 // classify() rejects wholly-out-of-domain ranges and the
                 // Avg axis before a descent is ever built; the builder
-                // still fails closed on both, plus on (Count, Total)
-                // until issue #806 lands.
+                // still fails closed on both.
                 cost_return_on_error_no_add!(
                     cost,
                     build_aggregate_secondary_proof(