Merge pull request #337 from BitGo/CSHLD-783-configurable-checkpoints

CSHLD-783: add support for custom checkpoints

Author: Ranjna-G

packages/wasm-privacy-coin/README.md

@@ -130,7 +130,9 @@ memory and reads a `Response` proto from the `LAST_RESULT` buffer.
 
 **Persistence.** `save()` / `fromState()` use serde JSON internally (the
 `PersistedShardTreeState` format). This is the on-disk/DB format, not the Java↔WASM
-wire format. The Java layer sees it as opaque `TreeState` bytes.
+wire format. The Java layer sees it as opaque `TreeState` bytes. `PersistedShardTreeState`
+carries `max_checkpoints`, so a restored tree keeps its original checkpoint retention
+capacity across `save()`/`fromState()` round-trips.
 
 **One instance = one tree.** Each `ShieldedMerkleTree` owns a dedicated Chicory
 `Instance` with its own WASM linear memory. Two instances never share state.
@@ -150,10 +152,11 @@ Implements `AutoCloseable`. Always use in try-with-resources.
 
 #### Factory methods
 
-| Method                                                   | Description                                                                    |
-| -------------------------------------------------------- | ------------------------------------------------------------------------------ |
-| `static fromFrontier(byte[] frontier, long blockHeight)` | Initialize from a CommitmentTree v0 frontier (raw bytes from `z_gettreestate`) |
-| `static fromState(TreeState state)`                      | Restore from a `TreeState` previously returned by `save()`                     |
+| Method                                                                           | Description                                                                                                                            |
+| -------------------------------------------------------------------------------- | -------------------------------------------------------------------------------------------------------------------------------------- |
+| `static fromFrontier(byte[] frontier, long blockHeight)`                         | Initialize from a CommitmentTree v0 frontier (raw bytes from `z_gettreestate`); convenience overload, `maxCheckpoints` defaults to 100 |
+| `static fromFrontier(byte[] frontier, long blockHeight, Integer maxCheckpoints)` | Initialize from a frontier with a custom checkpoint retention capacity; `null` = default (100)                                         |
+| `static fromState(TreeState state)`                                              | Restore from a `TreeState` previously returned by `save()`                                                                             |
 
 #### Instance methods
 
@@ -216,11 +219,11 @@ TreeState restored = TreeState.of(blob); // load from DB
 
 Immutable snapshot returned by `getInfo()`.
 
-| Field             | Type   | Description                                                              |
-| ----------------- | ------ | ------------------------------------------------------------------------ |
-| `tipHeight`       | `Long` | Most recently checkpointed block height; `null` if no block appended yet |
-| `leafCount`       | `long` | Total Orchard commitments appended across all blocks                     |
-| `checkpointCount` | `int`  | Number of checkpoints currently retained (max 100)                       |
+| Field             | Type   | Description                                                                        |
+| ----------------- | ------ | ---------------------------------------------------------------------------------- |
+| `tipHeight`       | `Long` | Most recently checkpointed block height; `null` if no block appended yet           |
+| `leafCount`       | `long` | Total Orchard commitments appended across all blocks                               |
+| `checkpointCount` | `int`  | Number of checkpoints currently retained, capped by `maxCheckpoints` (default 100) |
 
 ---
 
@@ -255,12 +258,21 @@ try (ShieldedMerkleTree tree = ShieldedMerkleTree.fromFrontier(frontier, blockHe
 }
 ```
 
+To use a non-default checkpoint retention capacity (e.g. a shallower reorg-depth
+tolerance), pass an explicit `maxCheckpoints`:
+
+```java
+try (ShieldedMerkleTree tree = ShieldedMerkleTree.fromFrontier(frontier, blockHeight, 20)) {
+    // tree retains at most 20 checkpoints before pruning the oldest
+}
+```
+
 ### 2. Initialize from an empty state
 
 ```java
 TreeState emptyState = new TreeState(
     "{\"shards\":[],\"cap\":{\"type\":\"Nil\"},\"checkpoints\":[],"
-    + "\"tip_height\":null,\"leaf_count\":0}");
+    + "\"tip_height\":null,\"leaf_count\":0,\"max_checkpoints\":100}");
 
 try (ShieldedMerkleTree tree = ShieldedMerkleTree.fromState(emptyState)) {
     tree.ping();
@@ -330,7 +342,7 @@ The wire format between Java and Rust is declared in `proto/privacy_coin.proto`.
 
 | Message                    | Export                   | Fields                                                                                 |
 | -------------------------- | ------------------------ | -------------------------------------------------------------------------------------- |
-| `FromFrontierRequest`      | `from_frontier`          | `frontier: bytes`, `block_height: uint32`                                              |
+| `FromFrontierRequest`      | `from_frontier`          | `frontier: bytes`, `block_height: uint32`, `max_checkpoints: optional uint32`          |
 | `AppendCommitmentsRequest` | `append_commitments`     | `block_height: uint32`, `commitments: repeated bytes`, `expected_root: optional bytes` |
 | `TruncateRequest`          | `truncate_to_checkpoint` | `block_height: uint32`                                                                 |
 

packages/wasm-privacy-coin/proto/privacy_coin.proto

@@ -9,6 +9,7 @@ option java_multiple_files = true;
 message FromFrontierRequest {
   bytes frontier      = 1;  // raw CommitmentTree v0 bytes
   uint32 block_height = 2;
+  optional uint32 max_checkpoints = 3;  // shardtree checkpoint capacity; absent = default (100)
 }
 
 message AppendCommitmentsRequest {

packages/wasm-privacy-coin/src/lib.rs

@@ -122,7 +122,11 @@ pub unsafe extern "C" fn from_frontier(ptr: *const u8, len: u32) -> i32 {
     let bytes = unsafe { std::slice::from_raw_parts(ptr, len as usize) };
     match FromFrontierRequest::decode(bytes) {
         Err(e) => write_error("DECODE_ERROR", &e.to_string()),
-        Ok(req) => match zcash::tree::OwnedTree::from_frontier(&req.frontier, req.block_height) {
+        Ok(req) => match zcash::tree::OwnedTree::from_frontier(
+            &req.frontier,
+            req.block_height,
+            req.max_checkpoints.map(|v| v as usize),
+        ) {
             Ok(tree) => {
                 TREE.with(|t| *t.borrow_mut() = Some(tree));
                 write_ok();

packages/wasm-privacy-coin/src/main/java/com/bitgo/wasm/privacycoin/zcash/ShieldedMerkleTree.java

@@ -59,16 +59,36 @@ private static ShieldedMerkleTree create(Consumer<WasmBridge> init) {
    * @throws WasmException if the frontier is invalid
    */
   public static ShieldedMerkleTree fromFrontier(byte[] frontier, long blockHeight) {
+    return fromFrontier(frontier, blockHeight, null);
+  }
+
+  /**
+   * Initializes a new tree from a CommitmentTree v0 frontier
+   * (the {@code orchardTree} value from {@code z_gettreestate}), with a custom
+   * checkpoint retention capacity.
+   *
+   * @param frontier        raw CommitmentTree v0 bytes
+   * @param blockHeight     block height at which the frontier was captured (u32 range)
+   * @param maxCheckpoints  number of checkpoints to retain before pruning the oldest;
+   *                        {@code null} to use the default (100)
+   * @return initialized tree instance
+   * @throws WasmException if the frontier is invalid
+   */
+  public static ShieldedMerkleTree fromFrontier(byte[] frontier, long blockHeight, Integer maxCheckpoints) {
     Objects.requireNonNull(frontier, "frontier must not be null");
     requireU32(blockHeight, "blockHeight");
+    if (maxCheckpoints != null) {
+      requireU32(maxCheckpoints, "maxCheckpoints");
+    }
     return create(bridge -> {
-      byte[] reqBytes = FromFrontierRequest.newBuilder()
+      FromFrontierRequest.Builder builder = FromFrontierRequest.newBuilder()
           .setFrontier(ByteString.copyFrom(frontier))
           // safe: requireU32 guarantees blockHeight is in [0, 0xFFFF_FFFF]
-          .setBlockHeight((int) blockHeight)
-          .build()
-          .toByteArray();
-      unwrapVoid(bridge.call("from_frontier", reqBytes));
+          .setBlockHeight((int) blockHeight);
+      if (maxCheckpoints != null) {
+        builder.setMaxCheckpoints(maxCheckpoints);
+      }
+      unwrapVoid(bridge.call("from_frontier", builder.build().toByteArray()));
     });
   }
 

packages/wasm-privacy-coin/src/test/java/com/bitgo/wasm/privacycoin/zcash/ShieldedMerkleTreeTest.java

@@ -24,7 +24,7 @@ class ShieldedMerkleTreeTest {
    */
   private static final TreeState EMPTY_STATE = new TreeState(
       "{\"shards\":[],\"cap\":{\"type\":\"Nil\"},\"checkpoints\":[],"
-      + "\"tip_height\":null,\"leaf_count\":0}");
+      + "\"tip_height\":null,\"leaf_count\":0,\"max_checkpoints\":100}");
 
   /**
    * CommitmentTree v0 frontier encoding for a single-leaf tree.
@@ -102,6 +102,45 @@ void fromFrontier_setsCheckpointCountToOne() {
     }
   }
 
+  @Test
+  void fromFrontier_withMaxCheckpoints_capsCheckpointRetention() {
+    try (ShieldedMerkleTree tree = ShieldedMerkleTree.fromFrontier(FRONTIER, 1L, 2)) {
+      tree.appendCommitments(2L, Collections.emptyList(), List.of(), null);
+      tree.appendCommitments(3L, Collections.emptyList(), List.of(), null);
+      tree.appendCommitments(4L, Collections.emptyList(), List.of(), null);
+
+      MerkleTreeInfo info = tree.getInfo();
+      assertEquals(2, info.checkpointCount);
+
+      WasmException ex = assertThrows(WasmException.class, () -> tree.truncateToCheckpoint(1L));
+      assertEquals("CHECKPOINT_NOT_FOUND", ex.getErrorCode());
+    }
+  }
+
+  @Test
+  void fromFrontier_negativeMaxCheckpoints_throwsIllegalArgumentException() {
+    assertThrows(IllegalArgumentException.class, () ->
+        ShieldedMerkleTree.fromFrontier(FRONTIER, 1L, -1));
+  }
+
+  @Test
+  void saveAndLoad_preservesMaxCheckpoints() {
+    TreeState savedState;
+    try (ShieldedMerkleTree tree = ShieldedMerkleTree.fromFrontier(FRONTIER, 1L, 2)) {
+      tree.appendCommitments(2L, Collections.emptyList(), List.of(), null);
+      tree.appendCommitments(3L, Collections.emptyList(), List.of(), null);
+      savedState = tree.save();
+    }
+    try (ShieldedMerkleTree restored = ShieldedMerkleTree.fromState(savedState)) {
+      restored.appendCommitments(4L, Collections.emptyList(), List.of(), null);
+      MerkleTreeInfo info = restored.getInfo();
+      assertEquals(2, info.checkpointCount);
+
+      WasmException ex = assertThrows(WasmException.class, () -> restored.truncateToCheckpoint(1L));
+      assertEquals("CHECKPOINT_NOT_FOUND", ex.getErrorCode());
+    }
+  }
+
   // -------------------------------------------------------------------------
   // fromState
   // -------------------------------------------------------------------------

packages/wasm-privacy-coin/src/zcash/tree.rs

@@ -62,6 +62,7 @@ pub struct PersistedShardTreeState {
     pub checkpoints: Vec<SerializedCheckpoint>,
     pub tip_height: Option<u32>,
     pub leaf_count: u64,
+    pub max_checkpoints: usize,
 }
 
 // ---------------------------------------------------------------------------
@@ -121,6 +122,7 @@ pub fn extract_state(
     tree: &ShieldedShardTree,
     tip_height: Option<u32>,
     leaf_count: u64,
+    max_checkpoints: usize,
 ) -> Result<PersistedShardTreeState, String> {
     let store = tree.store();
 
@@ -175,6 +177,7 @@ pub fn extract_state(
         checkpoints,
         tip_height,
         leaf_count,
+        max_checkpoints,
     })
 }
 
@@ -212,7 +215,7 @@ pub fn restore_state(state: &PersistedShardTreeState) -> Result<ShieldedShardTre
             .map_err(|e| format!("add_checkpoint error: {:?}", e))?;
     }
 
-    Ok(ShardTree::new(store, MAX_CHECKPOINTS))
+    Ok(ShardTree::new(store, state.max_checkpoints))
 }
 
 // ---------------------------------------------------------------------------
@@ -321,6 +324,7 @@ pub struct OwnedTree {
     tree: ShieldedShardTree,
     tip_height: Option<u32>,
     leaf_count: u64,
+    max_checkpoints: usize,
 }
 
 impl OwnedTree {
@@ -334,11 +338,17 @@ impl OwnedTree {
             tree,
             tip_height: persisted.tip_height,
             leaf_count: persisted.leaf_count,
+            max_checkpoints: persisted.max_checkpoints,
         })
     }
 
     /// Initialize from a CommitmentTree v0 frontier (raw bytes, not hex-encoded).
-    pub fn from_frontier(frontier: &[u8], block_height: u32) -> Result<Self, String> {
+    pub fn from_frontier(
+        frontier: &[u8],
+        block_height: u32,
+        max_checkpoints: Option<usize>,
+    ) -> Result<Self, String> {
+        let max_checkpoints = max_checkpoints.unwrap_or(MAX_CHECKPOINTS);
         use incrementalmerkletree::frontier::NonEmptyFrontier;
 
         let mut offset = 0;
@@ -395,7 +405,7 @@ impl OwnedTree {
 
         let leaf_count = u64::from(nef.position()) + 1;
 
-        let mut tree = ShardTree::new(MemoryShardStore::empty(), MAX_CHECKPOINTS);
+        let mut tree = ShardTree::new(MemoryShardStore::empty(), max_checkpoints);
         tree.insert_frontier_nodes(
             nef,
             Retention::Checkpoint {
@@ -409,12 +419,18 @@ impl OwnedTree {
             tree,
             tip_height: Some(block_height),
             leaf_count,
+            max_checkpoints,
         })
     }
 
     /// Serialize the tree state to bytes (UTF-8 JSON of `PersistedShardTreeState`).
     pub fn save(&self) -> Result<Vec<u8>, String> {
-        let state = extract_state(&self.tree, self.tip_height, self.leaf_count)?;
+        let state = extract_state(
+            &self.tree,
+            self.tip_height,
+            self.leaf_count,
+            self.max_checkpoints,
+        )?;
         serde_json::to_vec(&state).map_err(|e| format!("JSON serialize error: {}", e))
     }
 
@@ -540,7 +556,7 @@ mod tests {
     const F_CHECKPOINT_MARKED: u8 = 3;
 
     fn empty_tree() -> OwnedTree {
-        let json = r#"{"shards":[],"cap":{"type":"Nil"},"checkpoints":[],"tip_height":null,"leaf_count":0}"#;
+        let json = r#"{"shards":[],"cap":{"type":"Nil"},"checkpoints":[],"tip_height":null,"leaf_count":0,"max_checkpoints":100}"#;
         OwnedTree::from_state(json.as_bytes()).expect("empty state")
     }
 
@@ -582,7 +598,13 @@ mod tests {
             None,
         )
         .unwrap();
-        let state = extract_state(&tree.tree, tree.tip_height, tree.leaf_count).unwrap();
+        let state = extract_state(
+            &tree.tree,
+            tree.tip_height,
+            tree.leaf_count,
+            tree.max_checkpoints,
+        )
+        .unwrap();
 
         assert_eq!(
             find_f_in_state(&state, cmx1_hex),
@@ -603,7 +625,13 @@ mod tests {
         let mut tree = empty_tree();
         tree.append_commitments(1, vec![cmx(1)], vec![false], None)
             .unwrap();
-        let state = extract_state(&tree.tree, tree.tip_height, tree.leaf_count).unwrap();
+        let state = extract_state(
+            &tree.tree,
+            tree.tip_height,
+            tree.leaf_count,
+            tree.max_checkpoints,
+        )
+        .unwrap();
 
         assert_eq!(
             find_f_in_state(&state, cmx_hex),
@@ -619,4 +647,53 @@ mod tests {
             tree.append_commitments(1, vec![cmx(1), cmx(2)], vec![true, false, true], None);
         assert!(result.is_err());
     }
+
+    // -------------------------------------------------------------------------
+    // max_checkpoints
+    // -------------------------------------------------------------------------
+
+    fn frontier_bytes() -> Vec<u8> {
+        hex::decode("0101000000000000000000000000000000000000000000000000000000000000000000")
+            .unwrap()
+    }
+
+    #[test]
+    fn from_frontier_defaults_max_checkpoints_when_not_specified() {
+        let tree = OwnedTree::from_frontier(&frontier_bytes(), 1, None).unwrap();
+        assert_eq!(tree.max_checkpoints, MAX_CHECKPOINTS);
+    }
+
+    #[test]
+    fn from_frontier_uses_custom_max_checkpoints_when_specified() {
+        let tree = OwnedTree::from_frontier(&frontier_bytes(), 1, Some(5)).unwrap();
+        assert_eq!(tree.max_checkpoints, 5);
+    }
+
+    #[test]
+    fn max_checkpoints_enforced_evicts_oldest_checkpoint() {
+        let mut tree = OwnedTree::from_frontier(&frontier_bytes(), 1, Some(2)).unwrap();
+        tree.append_commitments(2, vec![], vec![], None).unwrap();
+        tree.append_commitments(3, vec![], vec![], None).unwrap();
+        tree.append_commitments(4, vec![], vec![], None).unwrap();
+
+        let (_, _, checkpoint_count) = tree.get_info().unwrap();
+        assert_eq!(
+            checkpoint_count, 2,
+            "checkpoint_count should be capped at 2"
+        );
+
+        let result = tree.truncate_to_checkpoint(1);
+        assert!(
+            result.is_err(),
+            "oldest checkpoint should have been evicted"
+        );
+    }
+
+    #[test]
+    fn save_round_trip_preserves_max_checkpoints() {
+        let tree = OwnedTree::from_frontier(&frontier_bytes(), 1, Some(7)).unwrap();
+        let bytes = tree.save().unwrap();
+        let restored = OwnedTree::from_state(&bytes).unwrap();
+        assert_eq!(restored.max_checkpoints, 7);
+    }
 }