feat: specific membership witness functions in aztec-nr (AztecProtocol#3674)

Fixes AztecProtocol#3663

Author: benesjan

yarn-project/aztec-nr/aztec/src/history/note_inclusion.nr

@@ -1,4 +1,3 @@
-use dep::protocol_types::constants::NOTE_HASH_TREE_HEIGHT;
 use dep::std::merkle::compute_merkle_root;
 
 use crate::{
@@ -8,10 +7,7 @@ use crate::{
         note_header::NoteHeader,
         note_interface::NoteInterface,
     },
-    oracle::get_membership_witness::{
-        get_membership_witness,
-        MembershipWitness,
-    },
+    oracle::get_membership_witness::get_note_hash_membership_witness,
 };
 
 pub fn prove_note_commitment_inclusion(
@@ -23,9 +19,7 @@ pub fn prove_note_commitment_inclusion(
     let block_header = context.get_block_header(block_number);
 
     // 2) Get the membership witness of the note in the note hash tree
-    let note_hash_tree_id = 2; // TODO(#3443)
-    let witness: MembershipWitness<NOTE_HASH_TREE_HEIGHT, NOTE_HASH_TREE_HEIGHT + 1> =
-        get_membership_witness(block_number, note_hash_tree_id, note_commitment);
+    let witness = get_note_hash_membership_witness(block_number, note_commitment);
 
     // 3) Prove that the commitment is in the note hash tree
     assert(

yarn-project/aztec-nr/aztec/src/history/public_value_inclusion.nr

@@ -1,15 +1,12 @@
 use dep::protocol_types::{
-    constants::{
-        PUBLIC_DATA_TREE_HEIGHT,
-        GENERATOR_INDEX__PUBLIC_LEAF_INDEX,
-    },
+    constants::GENERATOR_INDEX__PUBLIC_LEAF_INDEX,
     hash::pedersen_hash,
 };
 use dep::std::merkle::compute_merkle_root;
 
 use crate::{
     context::PrivateContext,
-    oracle::get_sibling_path::get_sibling_path,
+    oracle::get_sibling_path::get_public_data_sibling_path,
 };
 
 pub fn prove_public_value_inclusion(
@@ -31,9 +28,7 @@ pub fn prove_public_value_inclusion(
     );
 
     // 3) Get the sibling path of the value leaf index in the public data tree at block `block_number`.
-    let public_data_tree_id = 3; // TODO(#3443)
-    let path: [Field; PUBLIC_DATA_TREE_HEIGHT] =
-        get_sibling_path(block_number, public_data_tree_id, value_leaf_index);
+    let path = get_public_data_sibling_path(block_number, value_leaf_index);
 
     // 4) Prove that the value provided on input is in the public data tree at the given storage slot.
     assert(

yarn-project/aztec-nr/aztec/src/oracle/get_block_header.nr

@@ -1,18 +1,12 @@
 use dep::std::merkle::compute_merkle_root;
 use dep::protocol_types::{
     abis::block_header::BlockHeader,
-    constants::{
-        BLOCK_HEADER_LENGTH,
-        ARCHIVE_HEIGHT,
-    },
+    constants::BLOCK_HEADER_LENGTH,
 };
 
 use crate::{
     context::PrivateContext,
-    oracle::get_membership_witness::{
-        get_membership_witness,
-        MembershipWitness,
-    },
+    oracle::get_membership_witness::get_archive_membership_witness,
 };
 
 // TODO(#3564) - Nuke this oracle and Inject the number directly to context
@@ -47,8 +41,7 @@ pub fn get_block_header(block_number: u32, context: PrivateContext) -> BlockHead
     let block_hash = block_header.block_hash();
 
     // 5) Get the membership witness of the block in the archive
-    let archive_id = 5; // TODO(#3443)
-    let witness: MembershipWitness<ARCHIVE_HEIGHT, ARCHIVE_HEIGHT + 1> = get_membership_witness(block_header_block_number, archive_id, block_hash);
+    let witness = get_archive_membership_witness(block_header_block_number, block_hash);
 
     // 6) Check that the block is in the archive (i.e. the witness is valid)
     assert(context.block_header.archive_root == compute_merkle_root(block_hash, witness.index, witness.path), "Proving membership of a block in archive failed");

yarn-project/aztec-nr/aztec/src/oracle/get_membership_witness.nr

@@ -1,4 +1,8 @@
-use dep::protocol_types::constants::NOTE_HASH_TREE_HEIGHT;
+use dep::protocol_types::constants::{
+    ARCHIVE_HEIGHT,
+    CONTRACT_TREE_HEIGHT,
+    NOTE_HASH_TREE_HEIGHT,
+};
 use crate::utils::arr_copy_slice;
 
 // Note: We have M here because we need to somehow set it when calling get_membership_witness function and one way to
@@ -20,3 +24,23 @@ unconstrained pub fn get_membership_witness<N, M>(block_number: u32, tree_id: Fi
     let fields: [Field; M] = get_membership_witness_oracle(block_number, tree_id, leaf_value);
     MembershipWitness { index: fields[0], path: arr_copy_slice(fields, [0; N], 1) }
 }
+
+unconstrained pub fn get_contract_membership_witness(block_number: u32, leaf_value: Field) -> MembershipWitness<CONTRACT_TREE_HEIGHT, CONTRACT_TREE_HEIGHT + 1> {
+    let contract_tree_id = 0; // TODO(#3443)
+    get_membership_witness(block_number, contract_tree_id, leaf_value)
+}
+
+// Note: get_nullifier_membership_witness function is implemented in get_nullifier_membership_witness.nr
+
+unconstrained pub fn get_note_hash_membership_witness<N, M>(block_number: u32, leaf_value: Field) -> MembershipWitness<NOTE_HASH_TREE_HEIGHT, NOTE_HASH_TREE_HEIGHT + 1> {
+    let note_hash_tree_id = 2; // TODO(#3443)
+    get_membership_witness(block_number, note_hash_tree_id, leaf_value)
+}
+
+// There is no `get_public_data_membership_witness` function because it doesn't make sense to be getting a membership
+// witness for a value in the public data tree.
+
+unconstrained pub fn get_archive_membership_witness(block_number: u32, leaf_value: Field) -> MembershipWitness<ARCHIVE_HEIGHT, ARCHIVE_HEIGHT + 1> {
+    let archive_tree_id = 5; // TODO(#3443)
+    get_membership_witness(block_number, archive_tree_id, leaf_value)
+}

yarn-project/aztec-nr/aztec/src/oracle/get_sibling_path.nr

@@ -1,5 +1,4 @@
-use dep::protocol_types::constants::NOTE_HASH_TREE_HEIGHT;
-use crate::utils::arr_copy_slice;
+use dep::protocol_types::constants::PUBLIC_DATA_TREE_HEIGHT;
 
 #[oracle(getSiblingPath)]
 fn get_sibling_path_oracle<N>(_block_number: u32, _tree_id: Field, _leaf_index: Field) -> [Field; N] {}
@@ -8,3 +7,11 @@ unconstrained pub fn get_sibling_path<N>(block_number: u32, tree_id: Field, leaf
     let value: [Field; N] = get_sibling_path_oracle(block_number, tree_id, leaf_index);
     value
 }
+
+unconstrained pub fn get_public_data_sibling_path<N>(block_number: u32, leaf_index: Field) -> [Field; PUBLIC_DATA_TREE_HEIGHT] {
+    let public_data_tree_id = 3; // TODO(#3443)
+    get_sibling_path(block_number, public_data_tree_id, leaf_index)
+}
+
+// We don't implement specific function for other trees than public data tree because for the rest it makes sense
+// to use get membership witness function instead.