feat(replication)!: complete V1 state sync and bounded-memory restore

docs/book/src/count-indexed-tree.md

@@ -779,16 +779,21 @@ ordering. Top-k descending iteration encounters them last.
 
 ## Limitations and non-goals
 
-- **State sync does not support indexed trees.** A database containing
-  any indexed tree cannot be snapshot-synced: the restorer binds a
-  restored subtree to its parent with the two-input `combine_hash`,
-  which can never reproduce an indexed element's three-input binding,
-  and subtree discovery never enumerates the derived per-axis secondary
-  namespaces. Both the source (`fetch_chunk`) and target (subtree
-  discovery) sides reject with `Error::NotSupported` before any chunk is
-  produced or committed, so the failure is loud and early rather than a
-  half-restored database — but note the rejection is **database-wide**:
-  one indexed tree anywhere disables snapshot sync for the whole grove.
+- **State sync transfers an indexed subtree as one group.** The target
+  requests the primary with a header request built from its
+  hash-verified element (axis tags plus secondary root keys), the
+  source answers with an *indexed header* — the primary root hash and
+  each axis secondary's root hash, which the element itself never
+  stores — bundled with the primary's root chunk, and the per-axis
+  secondaries transfer as ordinary Merk chunks addressed by their
+  derived prefixes (they cannot be rebuilt locally: a secondary's root
+  commits to its write-history-dependent AVL shape). The header is only
+  a hint for per-chunk verification; once the primary and every
+  secondary are restored, the target unconditionally recomputes the
+  three-input binding (`combine_hash_three`, with the canonical
+  `axes_digest` for the multi-axis variant) from the *actual* restored
+  root hashes and requires it to match the element value hash bound
+  into the restored parent.
 - **Generic writes into an indexed primary are rejected.** `db.insert`,
   `db.delete` and `clear_subtree` targeting an indexed primary return
   `Error::NotSupported`, because none of them can mirror the change into

grovedb/Cargo.toml

@@ -55,7 +55,15 @@ grovedb-epoch-based-storage-flags = { version = "5.0.1", path = "../grovedb-epoc
 
 criterion = { workspace = true }
 hex = { workspace = true }
+# Process RSS sampling for the `#[ignore]`d state-sync restore
+# memory-ceiling harness (`tests/replication_scale_tests.rs`). The memory
+# that harness measures lives in RocksDB's C++ `WriteBatchWithIndex`, so it
+# has to be read from the OS rather than from a Rust allocator hook.
+libc = "0.2"
 pretty_assertions = "1.4.0"
+# Adversarial-input properties for the untrusted state-sync decode
+# surfaces (`tests/replication_fuzz_tests.rs`).
+proptest = "1.10.0"
 rand = { workspace = true }
 rand_distr = "0.6"
 assert_matches = { workspace = true }

grovedb/src/lib.rs

@@ -3216,11 +3216,14 @@ impl GroveDb {
     /// - `BulkAppendTree`: `blake3("bulk_state" || mmr_root || dense_root)`
     /// - `MmrTree`: the MMR root hash
     /// - `DenseAppendOnlyFixedSizeTree`: the dense tree root hash
+    /// - `PrivateDocumentStore`: `blake3("pds_state" || config_hash ||
+    ///   bulk_state_root)`
     ///
     /// For empty trees this returns the same conventions the insert path
     /// binds into the parent: `EMPTY_COMMITMENT_TREE_STATE_ROOT` for an
-    /// empty commitment tree, `NULL_HASH` (the empty Merk root) for the
-    /// other three types.
+    /// empty commitment tree, the config-parametrized empty state root for
+    /// an empty private document store, and `NULL_HASH` (the empty Merk
+    /// root) for the other three types.
     ///
     /// Returns an error if `element` is not a non-Merk data tree, or if the
     /// payload cannot be read back as a consistent tree of the declared
@@ -3326,6 +3329,41 @@ impl GroveDb {
                         ))
                     })
             }
+            Element::PrivateDocumentStore(total_count, entry_size, chunk_power, _) => {
+                // The state root binds the committed config even when the
+                // store is empty, so the empty case is the
+                // config-parametrized empty root rather than NULL_HASH.
+                if *total_count == 0 {
+                    return Ok(
+                        grovedb_private_document_store::empty_private_document_store_state_root(
+                            *entry_size,
+                            *chunk_power,
+                        ),
+                    );
+                }
+                let storage_ctx = self
+                    .db
+                    .get_transactional_storage_context(subtree_path, None, transaction)
+                    .unwrap();
+                let store = grovedb_private_document_store::PrivateDocumentStore::from_state(
+                    *total_count,
+                    *entry_size,
+                    *chunk_power,
+                    storage_ctx,
+                )
+                .unwrap()
+                .map_err(|e| {
+                    Error::CorruptedData(format!(
+                        "cannot open private document store of {total_count} entries from \
+                         payload: {e}"
+                    ))
+                })?;
+                store.compute_current_state_root_from_values().map_err(|e| {
+                    Error::CorruptedData(format!(
+                        "cannot compute private document store state root from payload: {e}"
+                    ))
+                })
+            }
             _ => Err(Error::InternalError(format!(
                 "compute_non_merk_state_root called on a non append-only element: {}",
                 element.type_str()

grovedb/src/operations/indexed_tree.rs

@@ -115,6 +115,40 @@ pub(crate) fn axis_secondary_tree_type(axis: IndexAxis) -> TreeType {
     }
 }
 
+/// Decode the configured axes of an indexed-tree element as
+/// `(axis, secondary_root_key)` pairs in canonical element order: the
+/// single implicit axis for PCIT / PSIT, the 1..=3 entry TLV for PCPSIT.
+///
+/// Errors if `element` is not an indexed-tree variant or a PCPSIT axis tag
+/// is invalid. The one intentional non-caller is
+/// `cleanup_dedicated_indexed_child_storage`, which must stay tolerant of
+/// an invalid tag (cleanup of a corrupt element should still clear the
+/// valid axes rather than fail).
+pub(crate) fn indexed_element_axes(
+    element: &Element,
+) -> Result<Vec<(IndexAxis, Option<Vec<u8>>)>, Error> {
+    match element.underlying() {
+        Element::ProvableCountIndexedTree(_, s, ..) => Ok(vec![(IndexAxis::Count, s.clone())]),
+        Element::ProvableSumIndexedTree(_, s, ..) => Ok(vec![(IndexAxis::Sum, s.clone())]),
+        Element::ProvableCountProvableSumIndexedTree(_, _, _, axes, _) => axes
+            .iter()
+            .map(|(tag, root_key)| {
+                IndexAxis::try_from_tag(*tag)
+                    .map(|axis| (axis, root_key.clone()))
+                    .map_err(|e| {
+                        Error::CorruptedData(format!(
+                            "invalid axis tag in indexed-tree element: {e}"
+                        ))
+                    })
+            })
+            .collect(),
+        other => Err(Error::CorruptedData(format!(
+            "expected an indexed-tree element, got {}",
+            other.type_str()
+        ))),
+    }
+}
+
 /// One primary entry's mirror-relevant state.
 ///
 /// The aggregates decide the row's sort key and carried sum; the value
@@ -421,31 +455,12 @@ impl GroveDb {
                     .map_err(Error::MerkError)
             )
         };
-        let axes: Vec<(IndexAxis, Option<Vec<u8>>)> = match element.underlying() {
-            Element::ProvableCountIndexedTree(_, s, ..) => vec![(IndexAxis::Count, s.clone())],
-            Element::ProvableSumIndexedTree(_, s, ..) => vec![(IndexAxis::Sum, s.clone())],
-            Element::ProvableCountProvableSumIndexedTree(_, _, _, axes, _) => {
-                let mut out = Vec::with_capacity(axes.len());
-                for (tag, root_key) in axes {
-                    let axis = cost_return_on_error_no_add!(
-                        cost,
-                        IndexAxis::try_from_tag(*tag).map_err(|e| Error::CorruptedData(format!(
-                            "open_indexed_secondaries_for_batch: invalid axis tag: {e}"
-                        )))
-                    );
-                    out.push((axis, root_key.clone()));
-                }
-                out
-            }
-            other => {
-                return Err(Error::CorruptedData(format!(
-                    "open_indexed_secondaries_for_batch: parent element is not an indexed tree, \
-                     got {}",
-                    other.type_str()
-                )))
-                .wrap_with_cost(cost);
-            }
-        };
+        let axes: Vec<(IndexAxis, Option<Vec<u8>>)> = cost_return_on_error_no_add!(
+            cost,
+            indexed_element_axes(&element).map_err(|e| Error::CorruptedData(format!(
+                "open_indexed_secondaries_for_batch: {e}"
+            )))
+        );
 
         let mut merks = Vec::with_capacity(axes.len());
         for (axis, root_key) in axes {
@@ -943,29 +958,12 @@ impl GroveDb {
             &mut cost,
             Element::get(&parent_merk, indexed_key, true, grove_version).map_err(Error::MerkError)
         );
-        let axes: Vec<(IndexAxis, Option<Vec<u8>>)> = match indexed_element.underlying() {
-            Element::ProvableCountIndexedTree(_, s, ..) => vec![(IndexAxis::Count, s.clone())],
-            Element::ProvableSumIndexedTree(_, s, ..) => vec![(IndexAxis::Sum, s.clone())],
-            Element::ProvableCountProvableSumIndexedTree(_, _, _, axes_tlv, _) => {
-                let mut out = Vec::with_capacity(axes_tlv.len());
-                for (tag, root_key) in axes_tlv {
-                    let axis = cost_return_on_error_no_add!(
-                        cost,
-                        IndexAxis::try_from_tag(*tag).map_err(|e| Error::CorruptedData(format!(
-                            "reconcile_indexed_tree_secondaries: invalid axis tag: {e}"
-                        )))
-                    );
-                    out.push((axis, root_key.clone()));
-                }
-                out
-            }
-            _ => {
-                return Err(Error::CorruptedData(
-                    "parent element at the indexed key is not an indexed tree".to_string(),
-                ))
-                .wrap_with_cost(cost);
-            }
-        };
+        let axes: Vec<(IndexAxis, Option<Vec<u8>>)> = cost_return_on_error_no_add!(
+            cost,
+            indexed_element_axes(&indexed_element).map_err(|e| Error::CorruptedData(format!(
+                "reconcile_indexed_tree_secondaries: {e}"
+            )))
+        );
         // The tightest ceiling across the configured axes (avg prepends a
         // 16-byte sort key against count/sum's 8).
         let max_item_key_len = axes
@@ -4589,3 +4587,61 @@ mod direct_axis_mirror_tests {
         );
     }
 }
+
+#[cfg(test)]
+mod indexed_element_axes_tests {
+    use super::*;
+
+    /// Every indexed variant decodes to its configured axes in canonical
+    /// element order: the single implicit axis for PCIT / PSIT, the TLV
+    /// entries for PCPSIT.
+    #[test]
+    fn indexed_element_axes_decodes_every_indexed_variant() {
+        assert_eq!(
+            indexed_element_axes(&Element::empty_provable_count_indexed_tree()).unwrap(),
+            vec![(IndexAxis::Count, None)]
+        );
+        assert_eq!(
+            indexed_element_axes(&Element::empty_provable_sum_indexed_tree()).unwrap(),
+            vec![(IndexAxis::Sum, None)]
+        );
+        let pcpsit = Element::empty_provable_count_provable_sum_indexed_tree(vec![
+            (IndexAxis::Count.tag(), Some(b"count_root".to_vec())),
+            (IndexAxis::Sum.tag(), None),
+        ])
+        .expect("valid axes");
+        assert_eq!(
+            indexed_element_axes(&pcpsit).unwrap(),
+            vec![
+                (IndexAxis::Count, Some(b"count_root".to_vec())),
+                (IndexAxis::Sum, None),
+            ]
+        );
+    }
+
+    /// The decode fails closed on anything that is not an indexed tree and
+    /// on a stored PCPSIT axis tag no axis maps to (the constructors
+    /// validate tags, but a corrupt stored element can carry any byte).
+    #[test]
+    fn indexed_element_axes_rejects_non_indexed_elements_and_invalid_tags() {
+        for element in [
+            Element::empty_tree(),
+            Element::new_item(vec![1]),
+            Element::empty_sum_tree(),
+        ] {
+            let err = indexed_element_axes(&element).expect_err("not an indexed tree");
+            assert!(
+                matches!(&err, Error::CorruptedData(m) if m.contains("expected an indexed-tree element")),
+                "{err}"
+            );
+        }
+
+        let corrupt =
+            Element::ProvableCountProvableSumIndexedTree(None, 0, 0, vec![(0xFF, None)], None);
+        let err = indexed_element_axes(&corrupt).expect_err("invalid axis tag");
+        assert!(
+            matches!(&err, Error::CorruptedData(m) if m.contains("invalid axis tag")),
+            "{err}"
+        );
+    }
+}