From 12a7feb51ff298a0d28a0e379ed7a9a58a54e234 Mon Sep 17 00:00:00 2001
From: Santiago Palladino <santiago@aztecprotocol.com>
Date: Thu, 27 Apr 2023 12:05:24 -0300
Subject: [PATCH] Check that public data tree root follows from left to right
 rollup

---
 .../circuits/rollup/components/components.cpp    |  3 +++
 .../src/aztec3/circuits/rollup/merge/.test.cpp   |  6 ++++++
 .../src/aztec3/circuits/rollup/root/.test.cpp    | 16 ++++++++++++----
 circuits/cpp/src/aztec3/utils/circuit_errors.hpp |  3 ++-
 4 files changed, 23 insertions(+), 5 deletions(-)

diff --git a/circuits/cpp/src/aztec3/circuits/rollup/components/components.cpp b/circuits/cpp/src/aztec3/circuits/rollup/components/components.cpp
index e4809910f46..5de96075bf1 100644
--- a/circuits/cpp/src/aztec3/circuits/rollup/components/components.cpp
+++ b/circuits/cpp/src/aztec3/circuits/rollup/components/components.cpp
@@ -132,6 +132,9 @@ void assert_prev_rollups_follow_on_from_each_other(DummyComposer& composer,
     composer.do_assert(left.end_contract_tree_snapshot == right.start_contract_tree_snapshot,
                        "input proofs have different contract tree snapshots",
                        utils::CircuitErrorCode::CONTRACT_TREE_SNAPSHOT_MISMATCH);
+    composer.do_assert(left.end_public_data_tree_snapshot == right.start_public_data_tree_snapshot,
+                       "input proofs have different public data tree snapshots",
+                       utils::CircuitErrorCode::CONTRACT_TREE_SNAPSHOT_MISMATCH);
 }
 
 } // namespace aztec3::circuits::rollup::components
\ No newline at end of file
diff --git a/circuits/cpp/src/aztec3/circuits/rollup/merge/.test.cpp b/circuits/cpp/src/aztec3/circuits/rollup/merge/.test.cpp
index 8a3fa7a4208..21f42262ead 100644
--- a/circuits/cpp/src/aztec3/circuits/rollup/merge/.test.cpp
+++ b/circuits/cpp/src/aztec3/circuits/rollup/merge/.test.cpp
@@ -209,6 +209,12 @@ TEST_F(merge_rollup_tests, native_start_and_end_snapshots)
               inputs.previous_rollup_data[0].base_or_merge_rollup_public_inputs.start_contract_tree_snapshot);
     ASSERT_EQ(outputs.end_contract_tree_snapshot,
               inputs.previous_rollup_data[1].base_or_merge_rollup_public_inputs.end_contract_tree_snapshot);
+
+    ASSERT_EQ(outputs.start_public_data_tree_snapshot,
+              inputs.previous_rollup_data[0].base_or_merge_rollup_public_inputs.start_public_data_tree_snapshot);
+    ASSERT_EQ(outputs.end_public_data_tree_snapshot,
+              inputs.previous_rollup_data[1].base_or_merge_rollup_public_inputs.end_public_data_tree_snapshot);
+
     ASSERT_FALSE(composer.failed());
 }
 
diff --git a/circuits/cpp/src/aztec3/circuits/rollup/root/.test.cpp b/circuits/cpp/src/aztec3/circuits/rollup/root/.test.cpp
index 7f0f3eaaa98..39cdede32e6 100644
--- a/circuits/cpp/src/aztec3/circuits/rollup/root/.test.cpp
+++ b/circuits/cpp/src/aztec3/circuits/rollup/root/.test.cpp
@@ -275,16 +275,24 @@ TEST_F(root_rollup_tests, native_root_missing_nullifier_logic)
     RootRollupPublicInputs outputs =
         aztec3::circuits::rollup::native_root_rollup::root_rollup_circuit(composer, rootRollupInputs);
 
-    // Check data trees
+    // Check private data trees
     ASSERT_EQ(
         outputs.start_private_data_tree_snapshot,
         rootRollupInputs.previous_rollup_data[0].base_or_merge_rollup_public_inputs.start_private_data_tree_snapshot);
     ASSERT_EQ(
         outputs.end_private_data_tree_snapshot,
         rootRollupInputs.previous_rollup_data[1].base_or_merge_rollup_public_inputs.end_private_data_tree_snapshot);
-    AppendOnlyTreeSnapshot<NT> expected_data_tree_snapshot = { .root = data_tree.root(),
-                                                               .next_available_leaf_index = 16 };
-    ASSERT_EQ(outputs.end_private_data_tree_snapshot, expected_data_tree_snapshot);
+    AppendOnlyTreeSnapshot<NT> expected_private_data_tree_snapshot = { .root = data_tree.root(),
+                                                                       .next_available_leaf_index = 16 };
+    ASSERT_EQ(outputs.end_private_data_tree_snapshot, expected_private_data_tree_snapshot);
+
+    // Check public data trees
+    ASSERT_EQ(
+        outputs.start_public_data_tree_snapshot,
+        rootRollupInputs.previous_rollup_data[0].base_or_merge_rollup_public_inputs.start_public_data_tree_snapshot);
+    ASSERT_EQ(
+        outputs.end_public_data_tree_snapshot,
+        rootRollupInputs.previous_rollup_data[1].base_or_merge_rollup_public_inputs.end_public_data_tree_snapshot);
 
     // check contract trees
     ASSERT_EQ(outputs.start_contract_tree_snapshot,
diff --git a/circuits/cpp/src/aztec3/utils/circuit_errors.hpp b/circuits/cpp/src/aztec3/utils/circuit_errors.hpp
index 43e8e9b1a25..1d40cad0a95 100644
--- a/circuits/cpp/src/aztec3/utils/circuit_errors.hpp
+++ b/circuits/cpp/src/aztec3/utils/circuit_errors.hpp
@@ -60,7 +60,8 @@ enum CircuitErrorCode : uint16_t {
     PRIVATE_DATA_TREE_SNAPSHOT_MISMATCH = 7004,
     NULLIFIER_TREE_SNAPSHOT_MISMATCH = 7005,
     CONTRACT_TREE_SNAPSHOT_MISMATCH = 7006,
-    MEMBERSHIP_CHECK_FAILED = 7007,
+    PUBLIC_DATA_TREE_SNAPSHOT_MISMATCH = 7007,
+    MEMBERSHIP_CHECK_FAILED = 7008,
 
     ROOT_CIRCUIT_FAILED = 8000,
 };