Implemented support for note trees in Miden node

block-producer/src/batch_builder/batch.rs

@@ -2,12 +2,10 @@ use std::collections::BTreeMap;
 
 use miden_objects::{
     accounts::AccountId,
-    crypto::{
-        hash::blake::{Blake3Digest, Blake3_256},
-        merkle::SimpleSmt,
-    },
+    batches::BatchNoteTree,
+    crypto::hash::blake::{Blake3Digest, Blake3_256},
     notes::{NoteEnvelope, Nullifier},
-    Digest, BATCH_OUTPUT_NOTES_TREE_DEPTH, MAX_NOTES_PER_BATCH,
+    Digest, MAX_NOTES_PER_BATCH,
 };
 use tracing::instrument;
 
@@ -27,7 +25,7 @@ pub struct TransactionBatch {
     id: BatchId,
     updated_accounts: BTreeMap<AccountId, AccountStates>,
     produced_nullifiers: Vec<Nullifier>,
-    created_notes_smt: SimpleSmt<BATCH_OUTPUT_NOTES_TREE_DEPTH>,
+    created_notes_smt: BatchNoteTree,
     /// The notes stored `created_notes_smt`
     created_notes: Vec<NoteEnvelope>,
 }
@@ -74,10 +72,10 @@ impl TransactionBatch {
             // TODO: document under what circumstances SMT creating can fail
             (
                 created_notes.clone(),
-                SimpleSmt::<BATCH_OUTPUT_NOTES_TREE_DEPTH>::with_contiguous_leaves(
-                    created_notes.into_iter().flat_map(|note_envelope| {
-                        [note_envelope.note_id().into(), note_envelope.metadata().into()]
-                    }),
+                BatchNoteTree::with_contiguous_leaves(
+                    created_notes
+                        .iter()
+                        .map(|note_envelope| (note_envelope.note_id(), note_envelope.metadata())),
                 )
                 .map_err(|e| BuildBatchError::NotesSmtError(e, txs))?,
             )

block-producer/src/block.rs

@@ -18,7 +18,7 @@ use crate::store::BlockInputsError;
 pub struct Block {
     pub header: BlockHeader,
     pub updated_accounts: Vec<(AccountId, Digest)>,
-    pub created_notes: BTreeMap<u64, NoteEnvelope>,
+    pub created_notes: Vec<(usize, usize, NoteEnvelope)>,
     pub produced_nullifiers: Vec<Nullifier>,
     // TODO:
     // - full states for updated public accounts

block-producer/src/block_builder/mod.rs

@@ -2,7 +2,7 @@ use std::sync::Arc;
 
 use async_trait::async_trait;
 use miden_node_utils::formatting::{format_array, format_blake3_digest};
-use miden_objects::{accounts::AccountId, notes::Nullifier, Digest, MAX_NOTES_PER_BATCH};
+use miden_objects::{accounts::AccountId, notes::Nullifier, Digest};
 use tracing::{debug, info, instrument};
 
 use crate::{
@@ -83,10 +83,10 @@ where
             .iter()
             .enumerate()
             .flat_map(|(batch_idx, batch)| {
-                batch.created_notes().enumerate().map(move |(note_idx_in_batch, note)| {
-                    let note_idx_in_block = batch_idx * MAX_NOTES_PER_BATCH + note_idx_in_batch;
-                    (note_idx_in_block as u64, *note)
-                })
+                batch
+                    .created_notes()
+                    .enumerate()
+                    .map(move |(note_idx_in_batch, note)| (batch_idx, note_idx_in_batch, *note))
             })
             .collect();
         let produced_nullifiers: Vec<Nullifier> =

block-producer/src/test_utils/block.rs

@@ -1,11 +1,9 @@
-use std::collections::BTreeMap;
-
 use miden_objects::{
     accounts::AccountId,
+    block::BlockNoteTree,
     crypto::merkle::{Mmr, SimpleSmt},
     notes::{NoteEnvelope, Nullifier},
-    BlockHeader, Digest, ACCOUNT_TREE_DEPTH, BLOCK_OUTPUT_NOTES_TREE_DEPTH, MAX_NOTES_PER_BATCH,
-    ONE, ZERO,
+    BlockHeader, Digest, ACCOUNT_TREE_DEPTH, ONE, ZERO,
 };
 
 use super::MockStoreSuccess;
@@ -92,7 +90,7 @@ pub struct MockBlockBuilder {
     last_block_header: BlockHeader,
 
     updated_accounts: Option<Vec<(AccountId, Digest)>>,
-    created_notes: Option<BTreeMap<u64, NoteEnvelope>>,
+    created_notes: Option<Vec<(usize, usize, NoteEnvelope)>>,
     produced_nullifiers: Option<Vec<Nullifier>>,
 }
 
@@ -124,7 +122,7 @@ impl MockBlockBuilder {
 
     pub fn created_notes(
         mut self,
-        created_notes: BTreeMap<u64, NoteEnvelope>,
+        created_notes: Vec<(usize, usize, NoteEnvelope)>,
     ) -> Self {
         self.created_notes = Some(created_notes);
 
@@ -149,7 +147,7 @@ impl MockBlockBuilder {
             self.store_chain_mmr.peaks(self.store_chain_mmr.forest()).unwrap().hash_peaks(),
             self.store_accounts.root(),
             Digest::default(),
-            note_created_smt_from_envelopes(created_notes.iter()).root(),
+            note_created_smt_from_envelopes(created_notes.iter().cloned()).root(),
             Digest::default(),
             Digest::default(),
             ZERO,
@@ -165,28 +163,23 @@ impl MockBlockBuilder {
     }
 }
 
-pub(crate) fn note_created_smt_from_envelopes<'a>(
-    note_iterator: impl Iterator<Item = (&'a u64, &'a NoteEnvelope)>
-) -> SimpleSmt<BLOCK_OUTPUT_NOTES_TREE_DEPTH> {
-    SimpleSmt::<BLOCK_OUTPUT_NOTES_TREE_DEPTH>::with_leaves(note_iterator.flat_map(
-        |(note_idx_in_block, note)| {
-            let index = note_idx_in_block * 2;
-            [(index, note.note_id().into()), (index + 1, note.metadata().into())]
-        },
-    ))
+pub(crate) fn note_created_smt_from_envelopes(
+    note_iterator: impl Iterator<Item = (usize, usize, NoteEnvelope)>
+) -> BlockNoteTree {
+    BlockNoteTree::with_entries(note_iterator.map(|(batch_idx, note_idx_in_batch, note)| {
+        (batch_idx, note_idx_in_batch, (note.note_id().into(), *note.metadata()))
+    }))
     .unwrap()
 }
 
 pub(crate) fn note_created_smt_from_batches<'a>(
     batches: impl Iterator<Item = &'a TransactionBatch>
-) -> SimpleSmt<BLOCK_OUTPUT_NOTES_TREE_DEPTH> {
-    let note_leaf_iterator = batches.enumerate().flat_map(|(batch_index, batch)| {
-        let subtree_index = batch_index * MAX_NOTES_PER_BATCH * 2;
-        batch.created_notes().enumerate().flat_map(move |(note_index, note)| {
-            let index = (subtree_index + note_index * 2) as u64;
-            [(index, note.note_id().into()), (index + 1, note.metadata().into())]
+) -> BlockNoteTree {
+    let note_leaf_iterator = batches.enumerate().flat_map(|(batch_idx, batch)| {
+        batch.created_notes().enumerate().map(move |(note_idx_in_batch, note)| {
+            (batch_idx, note_idx_in_batch, (note.note_id().into(), *note.metadata()))
         })
     });
 
-    SimpleSmt::<BLOCK_OUTPUT_NOTES_TREE_DEPTH>::with_leaves(note_leaf_iterator).unwrap()
+    BlockNoteTree::with_entries(note_leaf_iterator).unwrap()
 }

block-producer/src/test_utils/store.rs

@@ -2,9 +2,10 @@ use std::collections::BTreeSet;
 
 use async_trait::async_trait;
 use miden_objects::{
+    block::BlockNoteTree,
     crypto::merkle::{Mmr, SimpleSmt, Smt, ValuePath},
     notes::{NoteEnvelope, Nullifier},
-    BlockHeader, ACCOUNT_TREE_DEPTH, BLOCK_OUTPUT_NOTES_TREE_DEPTH, EMPTY_WORD, ONE, ZERO,
+    BlockHeader, ACCOUNT_TREE_DEPTH, EMPTY_WORD, ONE, ZERO,
 };
 
 use super::*;
@@ -22,7 +23,7 @@ use crate::{
 #[derive(Debug)]
 pub struct MockStoreSuccessBuilder {
     accounts: Option<SimpleSmt<ACCOUNT_TREE_DEPTH>>,
-    notes: Option<SimpleSmt<BLOCK_OUTPUT_NOTES_TREE_DEPTH>>,
+    notes: Option<BlockNoteTree>,
     produced_nullifiers: Option<BTreeSet<Digest>>,
     chain_mmr: Option<Mmr>,
     block_num: Option<u32>,
@@ -68,9 +69,9 @@ impl MockStoreSuccessBuilder {
         }
     }
 
-    pub fn initial_notes<'a>(
+    pub fn initial_notes(
         mut self,
-        notes: impl Iterator<Item = (&'a u64, &'a NoteEnvelope)>,
+        notes: impl Iterator<Item = (usize, usize, NoteEnvelope)>,
     ) -> Self {
         self.notes = Some(note_created_smt_from_envelopes(notes));
 
@@ -107,7 +108,7 @@ impl MockStoreSuccessBuilder {
     pub fn build(self) -> MockStoreSuccess {
         let block_num = self.block_num.unwrap_or(1);
         let accounts_smt = self.accounts.unwrap_or(SimpleSmt::new().unwrap());
-        let notes_smt = self.notes.unwrap_or(SimpleSmt::new().unwrap());
+        let notes_smt = self.notes.unwrap_or_default();
         let chain_mmr = self.chain_mmr.unwrap_or_default();
         let nullifiers_smt = self
             .produced_nullifiers

proto/proto/note.proto

@@ -23,8 +23,9 @@ message NoteSyncRecord {
 }
 
 message NoteCreated {
-    uint32 note_index = 1;
-    digest.Digest note_id = 2;
-    account.AccountId sender  = 3;
-    fixed64 tag = 4;
+    uint32 batch_index = 1;
+    uint32 note_index = 2;
+    digest.Digest note_id = 3;
+    account.AccountId sender  = 4;
+    fixed64 tag = 5;
 }

proto/src/domain/notes.rs

@@ -20,13 +20,14 @@ impl From<note::Note> for note::NoteSyncRecord {
 // NoteCreated
 // ================================================================================================
 
-impl From<(u64, NoteEnvelope)> for note::NoteCreated {
-    fn from((note_idx, note): (u64, NoteEnvelope)) -> Self {
+impl From<(usize, usize, NoteEnvelope)> for note::NoteCreated {
+    fn from((batch_idx, note_idx, note): (usize, usize, NoteEnvelope)) -> Self {
         Self {
+            batch_index: batch_idx as u32,
+            note_index: note_idx as u32,
             note_id: Some(note.note_id().into()),
             sender: Some(note.metadata().sender().into()),
             tag: note.metadata().tag().into(),
-            note_index: note_idx as u32,
         }
     }
 }

proto/src/generated/note.rs

@@ -35,11 +35,13 @@ pub struct NoteSyncRecord {
 #[derive(Clone, PartialEq, ::prost::Message)]
 pub struct NoteCreated {
     #[prost(uint32, tag = "1")]
+    pub batch_index: u32,
+    #[prost(uint32, tag = "2")]
     pub note_index: u32,
-    #[prost(message, optional, tag = "2")]
-    pub note_id: ::core::option::Option<super::digest::Digest>,
     #[prost(message, optional, tag = "3")]
+    pub note_id: ::core::option::Option<super::digest::Digest>,
+    #[prost(message, optional, tag = "4")]
     pub sender: ::core::option::Option<super::account::AccountId>,
-    #[prost(fixed64, tag = "4")]
+    #[prost(fixed64, tag = "5")]
     pub tag: u64,
 }

store/src/db/migrations.rs

@@ -18,7 +18,8 @@ pub static MIGRATIONS: Lazy<Migrations> = Lazy::new(|| {
             notes
         (
             block_num INTEGER NOT NULL,
-            note_index INTEGER NOT NULL,
+            batch_index INTEGER NOT NULL,   -- Index of batch in block, starting from 0
+            note_index INTEGER NOT NULL,    -- Index of note in batch, starting from 0
             note_hash BLOB NOT NULL,
             sender INTEGER NOT NULL,
             tag INTEGER NOT NULL,
@@ -27,6 +28,7 @@ pub static MIGRATIONS: Lazy<Migrations> = Lazy::new(|| {
             PRIMARY KEY (block_num, note_index),
             CONSTRAINT fk_block_num FOREIGN KEY (block_num) REFERENCES block_headers (block_num),
             CONSTRAINT notes_block_number_is_u32 CHECK (block_num >= 0 AND block_num < 4294967296),
+            CONSTRAINT notes_batch_index_is_u32 CHECK (batch_index BETWEEN 0 AND 0xFFFFFFFF)
             CONSTRAINT notes_note_index_is_u32 CHECK (note_index >= 0 AND note_index < 4294967296)
         ) STRICT, WITHOUT ROWID;
 

store/src/db/mod.rs

@@ -45,6 +45,7 @@ pub struct NullifierInfo {
 
 #[derive(Debug, Clone, PartialEq)]
 pub struct NoteCreated {
+    pub batch_index: u32,
     pub note_index: u32,
     pub note_id: RpoDigest,
     pub sender: AccountId,

store/src/db/sql.rs

@@ -248,24 +248,40 @@ pub fn select_nullifiers_by_block_range(
 ///
 /// A vector with notes, or an error.
 pub fn select_notes(conn: &mut Connection) -> Result<Vec<Note>> {
-    let mut stmt = conn.prepare("SELECT * FROM notes ORDER BY block_num ASC;")?;
+    let mut stmt = conn.prepare(
+        "
+        SELECT
+            block_num,
+            batch_index,
+            note_index,
+            note_hash,
+            sender,
+            tag,
+            merkle_path
+        FROM
+            notes
+        ORDER BY
+            block_num ASC;
+        ",
+    )?;
     let mut rows = stmt.query([])?;
 
     let mut notes = vec![];
     while let Some(row) = rows.next()? {
-        let note_id_data = row.get_ref(2)?.as_blob()?;
+        let note_id_data = row.get_ref(3)?.as_blob()?;
         let note_id = deserialize(note_id_data)?;
 
-        let merkle_path_data = row.get_ref(5)?.as_blob()?;
+        let merkle_path_data = row.get_ref(6)?.as_blob()?;
         let merkle_path = deserialize(merkle_path_data)?;
 
         notes.push(Note {
             block_num: row.get(0)?,
             note_created: NoteCreated {
-                note_index: row.get(1)?,
+                batch_index: row.get(1)?,
+                note_index: row.get(2)?,
                 note_id,
-                sender: column_value_as_u64(row, 3)?,
-                tag: column_value_as_u64(row, 4)?,
+                sender: column_value_as_u64(row, 4)?,
+                tag: column_value_as_u64(row, 5)?,
             },
             merkle_path,
         })
@@ -293,6 +309,7 @@ pub fn insert_notes(
         notes
         (
             block_num,
+            batch_index,
             note_index,
             note_hash,
             sender,
@@ -301,14 +318,15 @@ pub fn insert_notes(
         )
         VALUES
         (
-            ?1, ?2, ?3, ?4, ?5, ?6 
+            ?1, ?2, ?3, ?4, ?5, ?6, ?7
         );",
     )?;
 
     let mut count = 0;
     for note in notes.iter() {
         count += stmt.execute(params![
             note.block_num,
+            note.note_created.batch_index,
             note.note_created.note_index,
             note.note_created.note_id.to_bytes(),
             u64_to_value(note.note_created.sender),
@@ -345,6 +363,7 @@ pub fn select_notes_since_block_by_tag_and_sender(
         "
         SELECT
             block_num,
+            batch_index,
             note_index,
             note_hash,
             sender,
@@ -376,17 +395,19 @@ pub fn select_notes_since_block_by_tag_and_sender(
     let mut res = Vec::new();
     while let Some(row) = rows.next()? {
         let block_num = row.get(0)?;
-        let note_index = row.get(1)?;
-        let note_id_data = row.get_ref(2)?.as_blob()?;
+        let batch_index = row.get(1)?;
+        let note_index = row.get(2)?;
+        let note_id_data = row.get_ref(3)?.as_blob()?;
         let note_id = deserialize(note_id_data)?;
-        let sender = column_value_as_u64(row, 3)?;
-        let tag = column_value_as_u64(row, 4)?;
-        let merkle_path_data = row.get_ref(5)?.as_blob()?;
+        let sender = column_value_as_u64(row, 4)?;
+        let tag = column_value_as_u64(row, 5)?;
+        let merkle_path_data = row.get_ref(6)?.as_blob()?;
         let merkle_path = deserialize(merkle_path_data)?;
 
         let note = Note {
             block_num,
             note_created: NoteCreated {
+                batch_index,
                 note_index,
                 note_id,
                 sender,