Validate public state reads and writes in native base rollup circuit

circuits/cpp/src/aztec3/circuits/rollup/base/.test.cpp

@@ -2,11 +2,13 @@
 #include "aztec3/circuits/abis/membership_witness.hpp"
 #include "aztec3/circuits/abis/new_contract_data.hpp"
 #include "aztec3/circuits/abis/previous_kernel_data.hpp"
+#include "aztec3/circuits/abis/public_data_read.hpp"
 #include "aztec3/circuits/kernel/private/utils.hpp"
 #include "aztec3/circuits/abis/rollup/nullifier_leaf_preimage.hpp"
 #include "aztec3/constants.hpp"
 #include "barretenberg/crypto/sha256/sha256.hpp"
 #include "barretenberg/ecc/curves/bn254/fr.hpp"
+#include "barretenberg/stdlib/merkle_tree/memory_store.hpp"
 #include "barretenberg/stdlib/merkle_tree/memory_tree.hpp"
 
 #include "aztec3/circuits/rollup/test_utils/utils.hpp"
@@ -99,6 +101,9 @@ using aztec3::circuits::abis::FunctionData;
 using aztec3::circuits::abis::NewContractData;
 using aztec3::circuits::abis::OptionallyRevealedData;
 
+using aztec3::circuits::rollup::test_utils::utils::make_public_read;
+using aztec3::circuits::rollup::test_utils::utils::make_public_write;
+
 using DummyComposer = aztec3::utils::DummyComposer;
 } // namespace
 
@@ -762,4 +767,129 @@ TEST_F(base_rollup_tests, native_cbind_0)
     run_cbind(inputs, ignored_public_inputs, false);
 }
 
+TEST_F(base_rollup_tests, native_single_public_state_read)
+{
+    DummyComposer composer = DummyComposer();
+    native_base_rollup::MerkleTree private_data_tree(PRIVATE_DATA_TREE_HEIGHT);
+    native_base_rollup::MerkleTree contract_tree(CONTRACT_TREE_HEIGHT);
+    stdlib::merkle_tree::MemoryStore public_data_tree_store;
+    native_base_rollup::SparseTree public_data_tree(public_data_tree_store, PUBLIC_DATA_TREE_HEIGHT);
+
+    auto data_read = abis::PublicDataRead<NT>{
+        .leaf_index = fr(1),
+        .value = fr(42),
+    };
+
+    std::array<PreviousKernelData<NT>, 2> kernel_data = { get_empty_kernel(), get_empty_kernel() };
+    kernel_data[0].public_inputs.end.state_reads[0] = data_read;
+    auto inputs = test_utils::utils::base_rollup_inputs_from_kernels(
+        kernel_data, private_data_tree, contract_tree, public_data_tree);
+
+    BaseOrMergeRollupPublicInputs outputs =
+        aztec3::circuits::rollup::native_base_rollup::base_rollup_circuit(composer, inputs);
+
+    ASSERT_EQ(outputs.start_public_data_tree_snapshot, inputs.start_public_data_tree_snapshot);
+    ASSERT_EQ(outputs.end_public_data_tree_snapshot.root, public_data_tree.root());
+    ASSERT_EQ(outputs.end_public_data_tree_snapshot, outputs.start_public_data_tree_snapshot);
+    ASSERT_FALSE(composer.failed());
+    run_cbind(inputs, outputs);
+}
+
+TEST_F(base_rollup_tests, native_single_public_state_write)
+{
+    DummyComposer composer = DummyComposer();
+    native_base_rollup::MerkleTree private_data_tree(PRIVATE_DATA_TREE_HEIGHT);
+    native_base_rollup::MerkleTree contract_tree(CONTRACT_TREE_HEIGHT);
+    stdlib::merkle_tree::MemoryStore public_data_tree_store;
+    native_base_rollup::SparseTree public_data_tree(public_data_tree_store, PUBLIC_DATA_TREE_HEIGHT);
+
+    auto data_write = abis::PublicDataWrite<NT>{
+        .leaf_index = fr(1),
+        .old_value = fr(2),
+        .new_value = fr(42),
+    };
+
+    std::array<PreviousKernelData<NT>, 2> kernel_data = { get_empty_kernel(), get_empty_kernel() };
+    kernel_data[0].public_inputs.end.state_transitions[0] = data_write;
+
+    auto inputs = test_utils::utils::base_rollup_inputs_from_kernels(
+        kernel_data, private_data_tree, contract_tree, public_data_tree);
+
+    BaseOrMergeRollupPublicInputs outputs =
+        aztec3::circuits::rollup::native_base_rollup::base_rollup_circuit(composer, inputs);
+
+    ASSERT_EQ(outputs.start_public_data_tree_snapshot, inputs.start_public_data_tree_snapshot);
+    ASSERT_EQ(outputs.end_public_data_tree_snapshot.root, public_data_tree.root());
+    ASSERT_NE(outputs.end_public_data_tree_snapshot, outputs.start_public_data_tree_snapshot);
+    ASSERT_FALSE(composer.failed());
+    run_cbind(inputs, outputs);
+}
+
+TEST_F(base_rollup_tests, native_multiple_public_state_read_writes)
+{
+    DummyComposer composer = DummyComposer();
+    native_base_rollup::MerkleTree private_data_tree(PRIVATE_DATA_TREE_HEIGHT);
+    native_base_rollup::MerkleTree contract_tree(CONTRACT_TREE_HEIGHT);
+    stdlib::merkle_tree::MemoryStore public_data_tree_store;
+    native_base_rollup::SparseTree public_data_tree(public_data_tree_store, PUBLIC_DATA_TREE_HEIGHT);
+
+    std::array<PreviousKernelData<NT>, 2> kernel_data = { get_empty_kernel(), get_empty_kernel() };
+
+    // We set up reads and writes such that the right tx will read or write to indices already modified by the left tx
+    kernel_data[0].public_inputs.end.state_reads[0] = make_public_read(fr(1), fr(101));
+    kernel_data[0].public_inputs.end.state_reads[1] = make_public_read(fr(2), fr(102));
+    kernel_data[0].public_inputs.end.state_transitions[0] = make_public_write(fr(3), fr(103), fr(203));
+    kernel_data[0].public_inputs.end.state_transitions[1] = make_public_write(fr(4), fr(104), fr(204));
+    kernel_data[0].public_inputs.end.state_transitions[2] = make_public_write(fr(5), fr(105), fr(205));
+
+    kernel_data[1].public_inputs.end.state_reads[0] = make_public_read(fr(3), fr(203));
+    kernel_data[1].public_inputs.end.state_reads[1] = make_public_read(fr(11), fr(211));
+    kernel_data[1].public_inputs.end.state_transitions[0] = make_public_write(fr(12), fr(212), fr(312));
+    kernel_data[1].public_inputs.end.state_transitions[1] = make_public_write(fr(4), fr(204), fr(304));
+
+    auto inputs = test_utils::utils::base_rollup_inputs_from_kernels(
+        kernel_data, private_data_tree, contract_tree, public_data_tree);
+
+    BaseOrMergeRollupPublicInputs outputs =
+        aztec3::circuits::rollup::native_base_rollup::base_rollup_circuit(composer, inputs);
+
+    ASSERT_EQ(outputs.start_public_data_tree_snapshot, inputs.start_public_data_tree_snapshot);
+    ASSERT_EQ(outputs.end_public_data_tree_snapshot.root, public_data_tree.root());
+    ASSERT_NE(outputs.end_public_data_tree_snapshot, outputs.start_public_data_tree_snapshot);
+    ASSERT_FALSE(composer.failed());
+    run_cbind(inputs, outputs);
+}
+
+TEST_F(base_rollup_tests, native_invalid_public_state_read)
+{
+    DummyComposer composer = DummyComposer();
+    native_base_rollup::MerkleTree private_data_tree(PRIVATE_DATA_TREE_HEIGHT);
+    native_base_rollup::MerkleTree contract_tree(CONTRACT_TREE_HEIGHT);
+    stdlib::merkle_tree::MemoryStore public_data_tree_store;
+    native_base_rollup::SparseTree public_data_tree(public_data_tree_store, PUBLIC_DATA_TREE_HEIGHT);
+
+    auto data_read = abis::PublicDataRead<NT>{
+        .leaf_index = fr(1),
+        .value = fr(42),
+    };
+
+    std::array<PreviousKernelData<NT>, 2> kernel_data = { get_empty_kernel(), get_empty_kernel() };
+    kernel_data[0].public_inputs.end.state_reads[0] = data_read;
+    auto inputs = test_utils::utils::base_rollup_inputs_from_kernels(
+        kernel_data, private_data_tree, contract_tree, public_data_tree);
+
+    // We change the initial tree root so the read value does not match
+    public_data_tree.update_element(1, fr(43));
+    inputs.start_public_data_tree_snapshot.root = public_data_tree.root();
+
+    BaseOrMergeRollupPublicInputs outputs =
+        aztec3::circuits::rollup::native_base_rollup::base_rollup_circuit(composer, inputs);
+
+    ASSERT_EQ(outputs.start_public_data_tree_snapshot, inputs.start_public_data_tree_snapshot);
+    ASSERT_EQ(outputs.end_public_data_tree_snapshot.root, public_data_tree.root());
+    ASSERT_EQ(outputs.end_public_data_tree_snapshot, outputs.start_public_data_tree_snapshot);
+    ASSERT_TRUE(composer.failed());
+    run_cbind(inputs, outputs);
+}
+
 } // namespace aztec3::circuits::rollup::base::native_base_rollup_circuit

circuits/cpp/src/aztec3/circuits/rollup/base/init.hpp

@@ -6,7 +6,9 @@
 #include "aztec3/circuits/abis/rollup/base/base_rollup_inputs.hpp"
 #include "aztec3/circuits/abis/rollup/base/base_or_merge_rollup_public_inputs.hpp"
 #include "aztec3/utils/circuit_errors.hpp"
+#include "barretenberg/stdlib/merkle_tree/memory_store.hpp"
 #include "barretenberg/stdlib/merkle_tree/memory_tree.hpp"
+#include "barretenberg/stdlib/merkle_tree/merkle_tree.hpp"
 #include "barretenberg/stdlib/merkle_tree/nullifier_tree/nullifier_memory_tree.hpp"
 #include <aztec3/circuits/recursion/aggregator.hpp>
 #include <aztec3/circuits/abis/private_circuit_public_inputs.hpp>
@@ -39,5 +41,6 @@ using AppendOnlySnapshot = abis::AppendOnlyTreeSnapshot<NT>;
 using MerkleTree = stdlib::merkle_tree::MemoryTree;
 using NullifierTree = stdlib::merkle_tree::NullifierMemoryTree;
 using NullifierLeaf = stdlib::merkle_tree::nullifier_leaf;
+using SparseTree = stdlib::merkle_tree::MerkleTree<stdlib::merkle_tree::MemoryStore>;
 
 } // namespace aztec3::circuits::rollup::native_base_rollup

circuits/cpp/src/aztec3/circuits/rollup/base/native_base_rollup_circuit.cpp

@@ -1,3 +1,6 @@
+#include "aztec3/circuits/abis/membership_witness.hpp"
+#include "aztec3/circuits/abis/public_data_read.hpp"
+#include "aztec3/circuits/abis/public_data_write.hpp"
 #include "aztec3/circuits/hash.hpp"
 #include "aztec3/constants.hpp"
 #include "aztec3/utils/circuit_errors.hpp"
@@ -428,15 +431,99 @@ AppendOnlySnapshot check_nullifier_tree_non_membership_and_insert_to_tree(DummyC
     };
 }
 
-AppendOnlySnapshot insert_state_transitions(BaseRollupInputs const& baseRollupInputs)
+fr insert_state_transitions(
+    DummyComposer& composer,
+    fr tree_root,
+    std::array<abis::PublicDataWrite<NT>, STATE_TRANSITIONS_LENGTH> const& state_transitions,
+    size_t witnesses_offset,
+    std::array<abis::MembershipWitness<NT, PUBLIC_DATA_TREE_HEIGHT>, 2 * STATE_TRANSITIONS_LENGTH> const& witnesses)
 {
-    // TODO: Implement me
-    auto root = baseRollupInputs.start_public_data_tree_snapshot.root;
+    auto root = tree_root;
 
-    return {
-        .root = root,
-        .next_available_leaf_index = 0,
-    };
+    for (size_t i = 0; i < STATE_TRANSITIONS_LENGTH; ++i) {
+        const auto& state_write = state_transitions[i];
+        const auto& witness = witnesses[i + witnesses_offset];
+
+        if (state_write.is_empty())
+            continue;
+
+        composer.do_assert(
+            witness.leaf_index == state_write.leaf_index,
+            format("mismatch state write ", state_write.leaf_index, " and witness leaf index ", witness.leaf_index),
+            CircuitErrorCode::BASE__INVALID_PUBLIC_READS);
+
+        check_membership<NT>(composer,
+                             state_write.old_value,
+                             state_write.leaf_index,
+                             witness.sibling_path,
+                             root,
+                             format("validate_state_reads index ", i));
+
+        root = root_from_sibling_path<NT>(state_write.new_value, state_write.leaf_index, witness.sibling_path);
+    }
+
+    return root;
+}
+
+void validate_state_reads(
+    DummyComposer& composer,
+    fr tree_root,
+    std::array<abis::PublicDataRead<NT>, STATE_READS_LENGTH> const& state_reads,
+    size_t witnesses_offset,
+    std::array<abis::MembershipWitness<NT, PUBLIC_DATA_TREE_HEIGHT>, 2 * STATE_READS_LENGTH> const& witnesses)
+{
+    for (size_t i = 0; i < STATE_READS_LENGTH; ++i) {
+        const auto& state_read = state_reads[i];
+        const auto& witness = witnesses[i + witnesses_offset];
+
+        if (state_read.is_empty())
+            continue;
+
+        composer.do_assert(
+            witness.leaf_index == state_read.leaf_index,
+            format("mismatch state read ", state_read.leaf_index, " and witness leaf index ", witness.leaf_index),
+            CircuitErrorCode::BASE__INVALID_PUBLIC_READS);
+
+        check_membership<NT>(composer,
+                             state_read.value,
+                             state_read.leaf_index,
+                             witness.sibling_path,
+                             tree_root,
+                             format("validate_state_reads index ", i + witnesses_offset));
+    }
+};
+
+fr validate_and_process_public_state(DummyComposer& composer, BaseRollupInputs const& baseRollupInputs)
+{
+    // Process state reads and transitions for left input
+    validate_state_reads(composer,
+                         baseRollupInputs.start_public_data_tree_snapshot.root,
+                         baseRollupInputs.kernel_data[0].public_inputs.end.state_reads,
+                         0,
+                         baseRollupInputs.new_state_reads_sibling_paths);
+
+    auto mid_public_data_tree_root =
+        insert_state_transitions(composer,
+                                 baseRollupInputs.start_public_data_tree_snapshot.root,
+                                 baseRollupInputs.kernel_data[0].public_inputs.end.state_transitions,
+                                 0,
+                                 baseRollupInputs.new_state_transitions_sibling_paths);
+
+    // Process state reads and transitions for right input using the resulting tree root from the left one
+    validate_state_reads(composer,
+                         mid_public_data_tree_root,
+                         baseRollupInputs.kernel_data[1].public_inputs.end.state_reads,
+                         STATE_READS_LENGTH,
+                         baseRollupInputs.new_state_reads_sibling_paths);
+
+    auto end_public_data_tree_root =
+        insert_state_transitions(composer,
+                                 mid_public_data_tree_root,
+                                 baseRollupInputs.kernel_data[1].public_inputs.end.state_transitions,
+                                 STATE_TRANSITIONS_LENGTH,
+                                 baseRollupInputs.new_state_transitions_sibling_paths);
+
+    return end_public_data_tree_root;
 }
 
 BaseOrMergeRollupPublicInputs base_rollup_circuit(DummyComposer& composer, BaseRollupInputs const& baseRollupInputs)
@@ -482,8 +569,13 @@ BaseOrMergeRollupPublicInputs base_rollup_circuit(DummyComposer& composer, BaseR
     AppendOnlySnapshot end_nullifier_tree_snapshot =
         check_nullifier_tree_non_membership_and_insert_to_tree(composer, baseRollupInputs);
 
-    // Insert state transitions
-    auto end_public_data_tree_snapshot = insert_state_transitions(baseRollupInputs);
+    // Validate public state reads and transitions, and update public data tree
+    fr end_public_data_tree_root = validate_and_process_public_state(composer, baseRollupInputs);
+
+    AppendOnlySnapshot end_public_data_tree_snapshot = {
+        .root = end_public_data_tree_root,
+        .next_available_leaf_index = 0,
+    };
 
     // Calculate the overall calldata hash
     std::array<NT::fr, 2> calldata_hash = calculate_calldata_hash(baseRollupInputs, contract_leaves);