Merge origin/dev

.github/workflows/ci.yml

@@ -2,6 +2,7 @@ on:
   push:
     branches:
       - main
+      - dev
     paths-ignore:
       - "**.md"
       - "!.github/workflows/*.yml"
@@ -94,7 +95,7 @@ jobs:
         uses: Swatinem/rust-cache@v2
         with:
           shared-key: ci-rust-cache
-          save-if: ${{ github.ref == 'refs/heads/main' }}
+          save-if: ${{ github.ref == 'refs/heads/main' || github.ref == 'refs/heads/dev' }}
 
       - name: Lint workspace
         env:
@@ -125,7 +126,7 @@ jobs:
         uses: Swatinem/rust-cache@v2
         with:
           shared-key: ci-rust-cache
-          save-if: ${{ github.ref == 'refs/heads/main' }}
+          save-if: ${{ github.ref == 'refs/heads/main' || github.ref == 'refs/heads/dev' }}
 
       - name: Install nextest
         run: cargo install --locked cargo-nextest
@@ -156,7 +157,7 @@ jobs:
         uses: Swatinem/rust-cache@v2
         with:
           shared-key: ci-rust-cache
-          save-if: ${{ github.ref == 'refs/heads/main' }}
+          save-if: ${{ github.ref == 'refs/heads/main' || github.ref == 'refs/heads/dev' }}
 
       - name: Install nextest
         run: cargo install --locked cargo-nextest
@@ -245,7 +246,7 @@ jobs:
         uses: Swatinem/rust-cache@v2
         with:
           shared-key: ci-rust-cache
-          save-if: ${{ github.ref == 'refs/heads/main' }}
+          save-if: ${{ github.ref == 'refs/heads/main' || github.ref == 'refs/heads/dev' }}
 
       - name: Test valid proof
         env:
@@ -268,7 +269,7 @@ jobs:
         uses: Swatinem/rust-cache@v2
         with:
           shared-key: ci-rust-cache
-          save-if: ${{ github.ref == 'refs/heads/main' }}
+          save-if: ${{ github.ref == 'refs/heads/main' || github.ref == 'refs/heads/dev' }}
 
       - name: Install just
         run: cargo install --locked just

CONTRIBUTING.md

@@ -57,12 +57,16 @@ Before merging a PR, consider squashing non-meaningful commits. E.g.:
 
 Could be squashed to an empty commit if they belong to the same PR.
 
+## Default branch
+
+By default all PRs must be directed into the `dev` branch. This helps us to keep releases stable. 
+
 ## Branch workflow
 
-When bringing your feature branch up to date, prefer rebasing on top of `main`.
+When bringing your feature branch up to date, prefer rebasing on top of `dev`.
 
-- Preferred: `git rebase main`
-- Avoid: `git merge main` in feature branches
+- Preferred: `git rebase dev`
+- Avoid: `git merge dev` in feature branches
 
 This keeps commit history cleaner and makes reviews easier.
 

Cargo.lock

@@ -206,9 +206,9 @@ dependencies = [
 
 [[package]]
 name = "anyhow"
-version = "1.0.102"
+version = "1.0.103"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "7f202df86484c868dbad7eaa557ef785d5c66295e41b460ef922eca0723b842c"
+checksum = "2a4385e2e34eb35d6b3efe798b9eb88096925d87726c0798709bf56d9ed84af3"
 
 [[package]]
 name = "arbitrary"
@@ -9087,8 +9087,10 @@ name = "sequencer_service_protocol"
 version = "0.1.0"
 dependencies = [
  "common",
+ "hex",
  "lee",
  "lee_core",
+ "serde_with",
 ]
 
 [[package]]

integration_tests/src/lib.rs

@@ -6,11 +6,168 @@
 use std::time::Duration;
 
 use anyhow::{Context as _, Result};
+use key_protocol::key_management::key_tree::chain_index::ChainIndex;
+use lee::AccountId;
 use log::info;
+use sequencer_service_rpc::RpcClient as _;
 pub use test_fixtures::*;
+use wallet::{
+    cli::{
+        CliAccountMention, Command, SubcommandReturnValue,
+        account::{AccountSubcommand, NewSubcommand},
+        programs::{
+            native_token_transfer::AuthTransferSubcommand, token::TokenProgramAgnosticSubcommand,
+        },
+    },
+    storage::key_chain::FoundPrivateAccount,
+};
 
 /// Maximum time to wait for the indexer to catch up to the sequencer.
-pub const L2_TO_L1_TIMEOUT: Duration = Duration::from_mins(7);
+pub const L2_TO_L1_TIMEOUT: Duration = Duration::from_mins(6);
+
+/// Create a private or public account at the given chain index and return its ID.
+/// Pass `cci: None` to use the wallet's next available chain index.
+pub async fn new_account(
+    ctx: &mut TestContext,
+    private: bool,
+    cci: Option<ChainIndex>,
+) -> Result<AccountId> {
+    let subcommand = if private {
+        NewSubcommand::Private { cci, label: None }
+    } else {
+        NewSubcommand::Public { cci, label: None }
+    };
+    let result = wallet::cli::execute_subcommand(
+        ctx.wallet_mut(),
+        Command::Account(AccountSubcommand::New(subcommand)),
+    )
+    .await?;
+    let SubcommandReturnValue::RegisterAccount { account_id } = result else {
+        anyhow::bail!("Expected RegisterAccount return value");
+    };
+    Ok(account_id)
+}
+
+/// Send `amount` from `from` to `to` via an authenticated transfer (identifier 0).
+pub async fn send(
+    ctx: &mut TestContext,
+    from: CliAccountMention,
+    to: CliAccountMention,
+    amount: u128,
+) -> Result<()> {
+    let command = Command::AuthTransfer(AuthTransferSubcommand::Send {
+        from,
+        to: Some(to),
+        to_npk: None,
+        to_vpk: None,
+        to_keys: None,
+        to_identifier: Some(0),
+        amount,
+    });
+    wallet::cli::execute_subcommand(ctx.wallet_mut(), command).await?;
+    Ok(())
+}
+
+/// Create a token (New) and wait for the block to be included.
+pub async fn create_token(
+    ctx: &mut TestContext,
+    definition_account_id: CliAccountMention,
+    supply_account_id: CliAccountMention,
+    name: impl Into<String>,
+    total_supply: u128,
+) -> Result<()> {
+    let subcommand = TokenProgramAgnosticSubcommand::New {
+        definition_account_id,
+        supply_account_id,
+        name: name.into(),
+        total_supply,
+    };
+    wallet::cli::execute_subcommand(ctx.wallet_mut(), Command::Token(subcommand)).await?;
+    info!("Waiting for next block creation");
+    tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await;
+    Ok(())
+}
+
+/// Send tokens and wait for the block to be included.
+pub async fn token_send(
+    ctx: &mut TestContext,
+    from: CliAccountMention,
+    to: CliAccountMention,
+    amount: u128,
+) -> Result<()> {
+    let subcommand = TokenProgramAgnosticSubcommand::Send {
+        from,
+        to: Some(to),
+        to_npk: None,
+        to_vpk: None,
+        to_keys: None,
+        to_identifier: Some(0),
+        amount,
+    };
+    wallet::cli::execute_subcommand(ctx.wallet_mut(), Command::Token(subcommand)).await?;
+    info!("Waiting for next block creation");
+    tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await;
+    Ok(())
+}
+
+/// Retrieve the native token balance for `account_id`.
+pub async fn account_balance(ctx: &TestContext, account_id: AccountId) -> Result<u128> {
+    Ok(ctx
+        .sequencer_client()
+        .get_account_balance(account_id)
+        .await?)
+}
+
+/// Fetch the full account state for `account_id` from the sequencer.
+pub async fn get_account(ctx: &TestContext, account_id: AccountId) -> Result<lee::Account> {
+    Ok(ctx.sequencer_client().get_account(account_id).await?)
+}
+
+/// Fetch the current commitment for `account_id` and assert it is present in the sequencer state.
+pub async fn assert_private_commitment_in_state(
+    ctx: &TestContext,
+    account_id: AccountId,
+    label: &str,
+) -> Result<()> {
+    let commitment = ctx
+        .wallet()
+        .get_private_account_commitment(account_id)
+        .with_context(|| format!("Failed to get commitment for {label}"))?;
+    assert!(verify_commitment_is_in_state(commitment, ctx.sequencer_client()).await);
+    Ok(())
+}
+
+/// Sync the wallet's private accounts.
+pub async fn sync_private(ctx: &mut TestContext) -> Result<()> {
+    wallet::cli::execute_subcommand(
+        ctx.wallet_mut(),
+        Command::Account(AccountSubcommand::SyncPrivate {}),
+    )
+    .await?;
+    Ok(())
+}
+
+/// Look up a restored private account for `account_id`, panicking with `label` if absent.
+pub fn restored_private_account<'ctx>(
+    ctx: &'ctx TestContext,
+    account_id: AccountId,
+    label: &str,
+) -> FoundPrivateAccount<'ctx> {
+    ctx.wallet()
+        .storage()
+        .key_chain()
+        .private_account(account_id)
+        .unwrap_or_else(|| panic!("{label} should be restored"))
+}
+
+/// Assert that a restored public account's signing key exists, panicking with `label` if absent.
+pub fn assert_public_account_restored(ctx: &TestContext, account_id: AccountId, label: &str) {
+    ctx.wallet()
+        .storage()
+        .key_chain()
+        .pub_account_signing_key(account_id)
+        .unwrap_or_else(|| panic!("{label} should be restored"));
+}
 
 /// Poll the indexer until its last finalized block id reaches the sequencer's
 /// current last block id or until [`L2_TO_L1_TIMEOUT`] elapses.

integration_tests/tests/account.rs

@@ -4,12 +4,11 @@
 )]
 
 use anyhow::{Context as _, Result};
-use integration_tests::{TestContext, private_mention};
+use integration_tests::{TestContext, get_account, new_account, private_mention};
 use key_protocol::key_management::KeyChain;
 use lee::Data;
 use lee_core::account::Nonce;
 use log::info;
-use sequencer_service_rpc::RpcClient as _;
 use tokio::test;
 use wallet::{
     account::{AccountIdWithPrivacy, HumanReadableAccount, Label},
@@ -24,10 +23,7 @@ use wallet::{
 async fn get_existing_account() -> Result<()> {
     let ctx = TestContext::new().await?;
 
-    let account = ctx
-        .sequencer_client()
-        .get_account(ctx.existing_public_accounts()[0])
-        .await?;
+    let account = get_account(&ctx, ctx.existing_public_accounts()[0]).await?;
 
     assert_eq!(
         account.program_owner,
@@ -95,18 +91,7 @@ async fn add_label_to_existing_account() -> Result<()> {
 async fn new_public_account_without_label() -> Result<()> {
     let mut ctx = TestContext::new().await?;
 
-    let command = Command::Account(AccountSubcommand::New(NewSubcommand::Public {
-        cci: None,
-        label: None,
-    }));
-
-    let result = execute_subcommand(ctx.wallet_mut(), command).await?;
-
-    // Extract the account_id from the result
-
-    let wallet::cli::SubcommandReturnValue::RegisterAccount { account_id } = result else {
-        panic!("Expected RegisterAccount return value")
-    };
+    let account_id = new_account(&mut ctx, false, None).await?;
 
     // Verify no label was stored for the account id
     assert!(

integration_tests/tests/amm.rs

@@ -7,16 +7,18 @@
 use std::time::Duration;
 
 use anyhow::Result;
-use integration_tests::{TIME_TO_WAIT_FOR_BLOCK_SECONDS, TestContext, public_mention};
+use integration_tests::{
+    TIME_TO_WAIT_FOR_BLOCK_SECONDS, TestContext, create_token, get_account, new_account,
+    public_mention, token_send,
+};
 use log::info;
-use sequencer_service_rpc::RpcClient as _;
 use tokio::test;
 use wallet::{
     account::Label,
     cli::{
         Command, SubcommandReturnValue,
         account::{AccountSubcommand, NewSubcommand},
-        programs::{amm::AmmProgramAgnosticSubcommand, token::TokenProgramAgnosticSubcommand},
+        programs::amm::AmmProgramAgnosticSubcommand,
     },
 };
 
@@ -25,148 +27,60 @@ async fn amm_public() -> Result<()> {
     let mut ctx = TestContext::new().await?;
 
     // Create new account for the token definition
-    let SubcommandReturnValue::RegisterAccount {
-        account_id: definition_account_id_1,
-    } = wallet::cli::execute_subcommand(
-        ctx.wallet_mut(),
-        Command::Account(AccountSubcommand::New(NewSubcommand::Public {
-            cci: None,
-            label: None,
-        })),
-    )
-    .await?
-    else {
-        anyhow::bail!("Expected RegisterAccount return value");
-    };
+    let definition_account_id_1 = new_account(&mut ctx, false, None).await?;
 
     // Create new account for the token supply holder
-    let SubcommandReturnValue::RegisterAccount {
-        account_id: supply_account_id_1,
-    } = wallet::cli::execute_subcommand(
-        ctx.wallet_mut(),
-        Command::Account(AccountSubcommand::New(NewSubcommand::Public {
-            cci: None,
-            label: None,
-        })),
-    )
-    .await?
-    else {
-        anyhow::bail!("Expected RegisterAccount return value");
-    };
+    let supply_account_id_1 = new_account(&mut ctx, false, None).await?;
 
     // Create new account for receiving a token transaction
-    let SubcommandReturnValue::RegisterAccount {
-        account_id: recipient_account_id_1,
-    } = wallet::cli::execute_subcommand(
-        ctx.wallet_mut(),
-        Command::Account(AccountSubcommand::New(NewSubcommand::Public {
-            cci: None,
-            label: None,
-        })),
-    )
-    .await?
-    else {
-        anyhow::bail!("Expected RegisterAccount return value");
-    };
+    let recipient_account_id_1 = new_account(&mut ctx, false, None).await?;
 
     // Create new account for the token definition
-    let SubcommandReturnValue::RegisterAccount {
-        account_id: definition_account_id_2,
-    } = wallet::cli::execute_subcommand(
-        ctx.wallet_mut(),
-        Command::Account(AccountSubcommand::New(NewSubcommand::Public {
-            cci: None,
-            label: None,
-        })),
-    )
-    .await?
-    else {
-        anyhow::bail!("Expected RegisterAccount return value");
-    };
+    let definition_account_id_2 = new_account(&mut ctx, false, None).await?;
 
     // Create new account for the token supply holder
-    let SubcommandReturnValue::RegisterAccount {
-        account_id: supply_account_id_2,
-    } = wallet::cli::execute_subcommand(
-        ctx.wallet_mut(),
-        Command::Account(AccountSubcommand::New(NewSubcommand::Public {
-            cci: None,
-            label: None,
-        })),
-    )
-    .await?
-    else {
-        anyhow::bail!("Expected RegisterAccount return value");
-    };
+    let supply_account_id_2 = new_account(&mut ctx, false, None).await?;
 
     // Create new account for receiving a token transaction
-    let SubcommandReturnValue::RegisterAccount {
-        account_id: recipient_account_id_2,
-    } = wallet::cli::execute_subcommand(
-        ctx.wallet_mut(),
-        Command::Account(AccountSubcommand::New(NewSubcommand::Public {
-            cci: None,
-            label: None,
-        })),
-    )
-    .await?
-    else {
-        anyhow::bail!("Expected RegisterAccount return value");
-    };
+    let recipient_account_id_2 = new_account(&mut ctx, false, None).await?;
 
     // Create new token
-    let subcommand = TokenProgramAgnosticSubcommand::New {
-        definition_account_id: public_mention(definition_account_id_1),
-        supply_account_id: public_mention(supply_account_id_1),
-        name: "A NAM1".to_owned(),
-
-        total_supply: 37,
-    };
-    wallet::cli::execute_subcommand(ctx.wallet_mut(), Command::Token(subcommand)).await?;
-    info!("Waiting for next block creation");
-    tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await;
+    create_token(
+        &mut ctx,
+        public_mention(definition_account_id_1),
+        public_mention(supply_account_id_1),
+        "A NAM1".to_owned(),
+        37,
+    )
+    .await?;
 
     // Transfer 7 tokens from `supply_acc` to the account at account_id `recipient_account_id_1`
-    let subcommand = TokenProgramAgnosticSubcommand::Send {
-        from: public_mention(supply_account_id_1),
-        to: Some(public_mention(recipient_account_id_1)),
-        to_npk: None,
-        to_vpk: None,
-        to_keys: None,
-        to_identifier: Some(0),
-        amount: 7,
-    };
-
-    wallet::cli::execute_subcommand(ctx.wallet_mut(), Command::Token(subcommand)).await?;
-    info!("Waiting for next block creation");
-    tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await;
+    token_send(
+        &mut ctx,
+        public_mention(supply_account_id_1),
+        public_mention(recipient_account_id_1),
+        7,
+    )
+    .await?;
 
     // Create new token
-    let subcommand = TokenProgramAgnosticSubcommand::New {
-        definition_account_id: public_mention(definition_account_id_2),
-        supply_account_id: public_mention(supply_account_id_2),
-        name: "A NAM2".to_owned(),
-
-        total_supply: 37,
-    };
-    wallet::cli::execute_subcommand(ctx.wallet_mut(), Command::Token(subcommand)).await?;
-    info!("Waiting for next block creation");
-    tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await;
+    create_token(
+        &mut ctx,
+        public_mention(definition_account_id_2),
+        public_mention(supply_account_id_2),
+        "A NAM2".to_owned(),
+        37,
+    )
+    .await?;
 
     // Transfer 7 tokens from `supply_acc` to the account at account_id `recipient_account_id_2`
-    let subcommand = TokenProgramAgnosticSubcommand::Send {
-        from: public_mention(supply_account_id_2),
-        to: Some(public_mention(recipient_account_id_2)),
-        to_npk: None,
-        to_vpk: None,
-        to_keys: None,
-        to_identifier: Some(0),
-        amount: 7,
-    };
-
-    wallet::cli::execute_subcommand(ctx.wallet_mut(), Command::Token(subcommand)).await?;
-    info!("Waiting for next block creation");
-    tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await;
+    token_send(
+        &mut ctx,
+        public_mention(supply_account_id_2),
+        public_mention(recipient_account_id_2),
+        7,
+    )
+    .await?;
 
     info!("=================== SETUP FINISHED ===============");
 
@@ -174,19 +88,7 @@ async fn amm_public() -> Result<()> {
 
     // Setup accounts
     // Create new account for the user holding lp
-    let SubcommandReturnValue::RegisterAccount {
-        account_id: user_holding_lp,
-    } = wallet::cli::execute_subcommand(
-        ctx.wallet_mut(),
-        Command::Account(AccountSubcommand::New(NewSubcommand::Public {
-            cci: None,
-            label: None,
-        })),
-    )
-    .await?
-    else {
-        anyhow::bail!("Expected RegisterAccount return value");
-    };
+    let user_holding_lp = new_account(&mut ctx, false, None).await?;
 
     // Send creation tx
     let subcommand = AmmProgramAgnosticSubcommand::New {
@@ -201,17 +103,11 @@ async fn amm_public() -> Result<()> {
     info!("Waiting for next block creation");
     tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await;
 
-    let user_holding_a_acc = ctx
-        .sequencer_client()
-        .get_account(recipient_account_id_1)
-        .await?;
+    let user_holding_a_acc = get_account(&ctx, recipient_account_id_1).await?;
 
-    let user_holding_b_acc = ctx
-        .sequencer_client()
-        .get_account(recipient_account_id_2)
-        .await?;
+    let user_holding_b_acc = get_account(&ctx, recipient_account_id_2).await?;
 
-    let user_holding_lp_acc = ctx.sequencer_client().get_account(user_holding_lp).await?;
+    let user_holding_lp_acc = get_account(&ctx, user_holding_lp).await?;
 
     assert_eq!(
         u128::from_le_bytes(user_holding_a_acc.data[33..].try_into().unwrap()),
@@ -244,17 +140,11 @@ async fn amm_public() -> Result<()> {
     info!("Waiting for next block creation");
     tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await;
 
-    let user_holding_a_acc = ctx
-        .sequencer_client()
-        .get_account(recipient_account_id_1)
-        .await?;
+    let user_holding_a_acc = get_account(&ctx, recipient_account_id_1).await?;
 
-    let user_holding_b_acc = ctx
-        .sequencer_client()
-        .get_account(recipient_account_id_2)
-        .await?;
+    let user_holding_b_acc = get_account(&ctx, recipient_account_id_2).await?;
 
-    let user_holding_lp_acc = ctx.sequencer_client().get_account(user_holding_lp).await?;
+    let user_holding_lp_acc = get_account(&ctx, user_holding_lp).await?;
 
     assert_eq!(
         u128::from_le_bytes(user_holding_a_acc.data[33..].try_into().unwrap()),
@@ -287,17 +177,11 @@ async fn amm_public() -> Result<()> {
     info!("Waiting for next block creation");
     tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await;
 
-    let user_holding_a_acc = ctx
-        .sequencer_client()
-        .get_account(recipient_account_id_1)
-        .await?;
+    let user_holding_a_acc = get_account(&ctx, recipient_account_id_1).await?;
 
-    let user_holding_b_acc = ctx
-        .sequencer_client()
-        .get_account(recipient_account_id_2)
-        .await?;
+    let user_holding_b_acc = get_account(&ctx, recipient_account_id_2).await?;
 
-    let user_holding_lp_acc = ctx.sequencer_client().get_account(user_holding_lp).await?;
+    let user_holding_lp_acc = get_account(&ctx, user_holding_lp).await?;
 
     assert_eq!(
         u128::from_le_bytes(user_holding_a_acc.data[33..].try_into().unwrap()),
@@ -331,17 +215,11 @@ async fn amm_public() -> Result<()> {
     info!("Waiting for next block creation");
     tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await;
 
-    let user_holding_a_acc = ctx
-        .sequencer_client()
-        .get_account(recipient_account_id_1)
-        .await?;
+    let user_holding_a_acc = get_account(&ctx, recipient_account_id_1).await?;
 
-    let user_holding_b_acc = ctx
-        .sequencer_client()
-        .get_account(recipient_account_id_2)
-        .await?;
+    let user_holding_b_acc = get_account(&ctx, recipient_account_id_2).await?;
 
-    let user_holding_lp_acc = ctx.sequencer_client().get_account(user_holding_lp).await?;
+    let user_holding_lp_acc = get_account(&ctx, user_holding_lp).await?;
 
     assert_eq!(
         u128::from_le_bytes(user_holding_a_acc.data[33..].try_into().unwrap()),
@@ -375,17 +253,11 @@ async fn amm_public() -> Result<()> {
     info!("Waiting for next block creation");
     tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await;
 
-    let user_holding_a_acc = ctx
-        .sequencer_client()
-        .get_account(recipient_account_id_1)
-        .await?;
+    let user_holding_a_acc = get_account(&ctx, recipient_account_id_1).await?;
 
-    let user_holding_b_acc = ctx
-        .sequencer_client()
-        .get_account(recipient_account_id_2)
-        .await?;
+    let user_holding_b_acc = get_account(&ctx, recipient_account_id_2).await?;
 
-    let user_holding_lp_acc = ctx.sequencer_client().get_account(user_holding_lp).await?;
+    let user_holding_lp_acc = get_account(&ctx, user_holding_lp).await?;
 
     assert_eq!(
         u128::from_le_bytes(user_holding_a_acc.data[33..].try_into().unwrap()),
@@ -412,33 +284,9 @@ async fn amm_new_pool_using_labels() -> Result<()> {
     let mut ctx = TestContext::new().await?;
 
     // Create token 1 accounts
-    let SubcommandReturnValue::RegisterAccount {
-        account_id: definition_account_id_1,
-    } = wallet::cli::execute_subcommand(
-        ctx.wallet_mut(),
-        Command::Account(AccountSubcommand::New(NewSubcommand::Public {
-            cci: None,
-            label: None,
-        })),
-    )
-    .await?
-    else {
-        anyhow::bail!("Expected RegisterAccount return value");
-    };
+    let definition_account_id_1 = new_account(&mut ctx, false, None).await?;
 
-    let SubcommandReturnValue::RegisterAccount {
-        account_id: supply_account_id_1,
-    } = wallet::cli::execute_subcommand(
-        ctx.wallet_mut(),
-        Command::Account(AccountSubcommand::New(NewSubcommand::Public {
-            cci: None,
-            label: None,
-        })),
-    )
-    .await?
-    else {
-        anyhow::bail!("Expected RegisterAccount return value");
-    };
+    let supply_account_id_1 = new_account(&mut ctx, false, None).await?;
 
     // Create holding_a with a label
     let holding_a_label = Label::new("amm-holding-a-label");
@@ -457,33 +305,9 @@ async fn amm_new_pool_using_labels() -> Result<()> {
     };
 
     // Create token 2 accounts
-    let SubcommandReturnValue::RegisterAccount {
-        account_id: definition_account_id_2,
-    } = wallet::cli::execute_subcommand(
-        ctx.wallet_mut(),
-        Command::Account(AccountSubcommand::New(NewSubcommand::Public {
-            cci: None,
-            label: None,
-        })),
-    )
-    .await?
-    else {
-        anyhow::bail!("Expected RegisterAccount return value");
-    };
+    let definition_account_id_2 = new_account(&mut ctx, false, None).await?;
 
-    let SubcommandReturnValue::RegisterAccount {
-        account_id: supply_account_id_2,
-    } = wallet::cli::execute_subcommand(
-        ctx.wallet_mut(),
-        Command::Account(AccountSubcommand::New(NewSubcommand::Public {
-            cci: None,
-            label: None,
-        })),
-    )
-    .await?
-    else {
-        anyhow::bail!("Expected RegisterAccount return value");
-    };
+    let supply_account_id_2 = new_account(&mut ctx, false, None).await?;
 
     // Create holding_b with a label
     let holding_b_label = Label::new("amm-holding-b-label");
@@ -518,48 +342,40 @@ async fn amm_new_pool_using_labels() -> Result<()> {
     };
 
     // Create token 1 and distribute to holding_a
-    let subcommand = TokenProgramAgnosticSubcommand::New {
-        definition_account_id: public_mention(definition_account_id_1),
-        supply_account_id: public_mention(supply_account_id_1),
-        name: "TOKEN1".to_owned(),
-        total_supply: 10,
-    };
-    wallet::cli::execute_subcommand(ctx.wallet_mut(), Command::Token(subcommand)).await?;
-    tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await;
+    create_token(
+        &mut ctx,
+        public_mention(definition_account_id_1),
+        public_mention(supply_account_id_1),
+        "TOKEN1".to_owned(),
+        10,
+    )
+    .await?;
 
-    let subcommand = TokenProgramAgnosticSubcommand::Send {
-        from: public_mention(supply_account_id_1),
-        to: Some(public_mention(holding_a_id)),
-        to_npk: None,
-        to_vpk: None,
-        to_keys: None,
-        to_identifier: Some(0),
-        amount: 5,
-    };
-    wallet::cli::execute_subcommand(ctx.wallet_mut(), Command::Token(subcommand)).await?;
-    tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await;
+    token_send(
+        &mut ctx,
+        public_mention(supply_account_id_1),
+        public_mention(holding_a_id),
+        5,
+    )
+    .await?;
 
     // Create token 2 and distribute to holding_b
-    let subcommand = TokenProgramAgnosticSubcommand::New {
-        definition_account_id: public_mention(definition_account_id_2),
-        supply_account_id: public_mention(supply_account_id_2),
-        name: "TOKEN2".to_owned(),
-        total_supply: 10,
-    };
-    wallet::cli::execute_subcommand(ctx.wallet_mut(), Command::Token(subcommand)).await?;
-    tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await;
+    create_token(
+        &mut ctx,
+        public_mention(definition_account_id_2),
+        public_mention(supply_account_id_2),
+        "TOKEN2".to_owned(),
+        10,
+    )
+    .await?;
 
-    let subcommand = TokenProgramAgnosticSubcommand::Send {
-        from: public_mention(supply_account_id_2),
-        to: Some(public_mention(holding_b_id)),
-        to_npk: None,
-        to_vpk: None,
-        to_keys: None,
-        to_identifier: Some(0),
-        amount: 5,
-    };
-    wallet::cli::execute_subcommand(ctx.wallet_mut(), Command::Token(subcommand)).await?;
-    tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await;
+    token_send(
+        &mut ctx,
+        public_mention(supply_account_id_2),
+        public_mention(holding_b_id),
+        5,
+    )
+    .await?;
 
     // Create AMM pool using account labels instead of IDs
     let subcommand = AmmProgramAgnosticSubcommand::New {
@@ -572,7 +388,7 @@ async fn amm_new_pool_using_labels() -> Result<()> {
     wallet::cli::execute_subcommand(ctx.wallet_mut(), Command::AMM(subcommand)).await?;
     tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await;
 
-    let holding_lp_acc = ctx.sequencer_client().get_account(holding_lp_id).await?;
+    let holding_lp_acc = get_account(&ctx, holding_lp_id).await?;
 
     // LP balance should be 3 (geometric mean of 3, 3)
     assert_eq!(

integration_tests/tests/ata.rs

@@ -9,75 +9,34 @@ use std::time::Duration;
 use anyhow::{Context as _, Result};
 use associated_token_account_core::{compute_ata_seed, get_associated_token_account_id};
 use integration_tests::{
-    TIME_TO_WAIT_FOR_BLOCK_SECONDS, TestContext, private_mention, public_mention,
-    verify_commitment_is_in_state,
+    TIME_TO_WAIT_FOR_BLOCK_SECONDS, TestContext, create_token, get_account, new_account,
+    private_mention, public_mention, token_send, verify_commitment_is_in_state,
 };
 use log::info;
 use sequencer_service_rpc::RpcClient as _;
 use token_core::{TokenDefinition, TokenHolding};
 use tokio::test;
-use wallet::cli::{
-    Command, SubcommandReturnValue,
-    account::{AccountSubcommand, NewSubcommand},
-    programs::{ata::AtaSubcommand, token::TokenProgramAgnosticSubcommand},
-};
-
-/// Create a public account and return its ID.
-async fn new_public_account(ctx: &mut TestContext) -> Result<lee::AccountId> {
-    let result = wallet::cli::execute_subcommand(
-        ctx.wallet_mut(),
-        Command::Account(AccountSubcommand::New(NewSubcommand::Public {
-            cci: None,
-            label: None,
-        })),
-    )
-    .await?;
-    let SubcommandReturnValue::RegisterAccount { account_id } = result else {
-        anyhow::bail!("Expected RegisterAccount return value");
-    };
-    Ok(account_id)
-}
-
-/// Create a private account and return its ID.
-async fn new_private_account(ctx: &mut TestContext) -> Result<lee::AccountId> {
-    let result = wallet::cli::execute_subcommand(
-        ctx.wallet_mut(),
-        Command::Account(AccountSubcommand::New(NewSubcommand::Private {
-            cci: None,
-            label: None,
-        })),
-    )
-    .await?;
-    let SubcommandReturnValue::RegisterAccount { account_id } = result else {
-        anyhow::bail!("Expected RegisterAccount return value");
-    };
-    Ok(account_id)
-}
+use wallet::cli::{Command, programs::ata::AtaSubcommand};
 
 #[test]
 async fn create_ata_initializes_holding_account() -> Result<()> {
     let mut ctx = TestContext::new().await?;
 
-    let definition_account_id = new_public_account(&mut ctx).await?;
-    let supply_account_id = new_public_account(&mut ctx).await?;
-    let owner_account_id = new_public_account(&mut ctx).await?;
+    let definition_account_id = new_account(&mut ctx, false, None).await?;
+    let supply_account_id = new_account(&mut ctx, false, None).await?;
+    let owner_account_id = new_account(&mut ctx, false, None).await?;
 
     // Create a fungible token
     let total_supply = 100_u128;
-    wallet::cli::execute_subcommand(
-        ctx.wallet_mut(),
-        Command::Token(TokenProgramAgnosticSubcommand::New {
-            definition_account_id: public_mention(definition_account_id),
-            supply_account_id: public_mention(supply_account_id),
-            name: "TEST".to_owned(),
-            total_supply,
-        }),
+    create_token(
+        &mut ctx,
+        public_mention(definition_account_id),
+        public_mention(supply_account_id),
+        "TEST".to_owned(),
+        total_supply,
     )
     .await?;
 
-    info!("Waiting for next block creation");
-    tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await;
-
     // Create the ATA for owner + definition
     wallet::cli::execute_subcommand(
         ctx.wallet_mut(),
@@ -121,25 +80,20 @@ async fn create_ata_initializes_holding_account() -> Result<()> {
 async fn create_ata_is_idempotent() -> Result<()> {
     let mut ctx = TestContext::new().await?;
 
-    let definition_account_id = new_public_account(&mut ctx).await?;
-    let supply_account_id = new_public_account(&mut ctx).await?;
-    let owner_account_id = new_public_account(&mut ctx).await?;
+    let definition_account_id = new_account(&mut ctx, false, None).await?;
+    let supply_account_id = new_account(&mut ctx, false, None).await?;
+    let owner_account_id = new_account(&mut ctx, false, None).await?;
 
     // Create a fungible token
-    wallet::cli::execute_subcommand(
-        ctx.wallet_mut(),
-        Command::Token(TokenProgramAgnosticSubcommand::New {
-            definition_account_id: public_mention(definition_account_id),
-            supply_account_id: public_mention(supply_account_id),
-            name: "TEST".to_owned(),
-            total_supply: 100,
-        }),
+    create_token(
+        &mut ctx,
+        public_mention(definition_account_id),
+        public_mention(supply_account_id),
+        "TEST".to_owned(),
+        100,
     )
     .await?;
 
-    info!("Waiting for next block creation");
-    tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await;
-
     // Create the ATA once
     wallet::cli::execute_subcommand(
         ctx.wallet_mut(),
@@ -196,28 +150,23 @@ async fn create_ata_is_idempotent() -> Result<()> {
 async fn transfer_and_burn_via_ata() -> Result<()> {
     let mut ctx = TestContext::new().await?;
 
-    let definition_account_id = new_public_account(&mut ctx).await?;
-    let supply_account_id = new_public_account(&mut ctx).await?;
-    let sender_account_id = new_public_account(&mut ctx).await?;
-    let recipient_account_id = new_public_account(&mut ctx).await?;
+    let definition_account_id = new_account(&mut ctx, false, None).await?;
+    let supply_account_id = new_account(&mut ctx, false, None).await?;
+    let sender_account_id = new_account(&mut ctx, false, None).await?;
+    let recipient_account_id = new_account(&mut ctx, false, None).await?;
 
     let total_supply = 1000_u128;
 
     // Create a fungible token, supply goes to supply_account_id
-    wallet::cli::execute_subcommand(
-        ctx.wallet_mut(),
-        Command::Token(TokenProgramAgnosticSubcommand::New {
-            definition_account_id: public_mention(definition_account_id),
-            supply_account_id: public_mention(supply_account_id),
-            name: "TEST".to_owned(),
-            total_supply,
-        }),
+    create_token(
+        &mut ctx,
+        public_mention(definition_account_id),
+        public_mention(supply_account_id),
+        "TEST".to_owned(),
+        total_supply,
     )
     .await?;
 
-    info!("Waiting for next block creation");
-    tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await;
-
     // Derive ATA addresses
     let ata_program_id = programs::ata().id();
     let sender_ata_id = get_associated_token_account_id(
@@ -252,23 +201,14 @@ async fn transfer_and_burn_via_ata() -> Result<()> {
 
     // Fund sender's ATA from the supply account (direct token transfer)
     let fund_amount = 200_u128;
-    wallet::cli::execute_subcommand(
-        ctx.wallet_mut(),
-        Command::Token(TokenProgramAgnosticSubcommand::Send {
-            from: public_mention(supply_account_id),
-            to: Some(public_mention(sender_ata_id)),
-            to_npk: None,
-            to_vpk: None,
-            to_keys: None,
-            to_identifier: Some(0),
-            amount: fund_amount,
-        }),
+    token_send(
+        &mut ctx,
+        public_mention(supply_account_id),
+        public_mention(sender_ata_id),
+        fund_amount,
     )
     .await?;
 
-    info!("Waiting for next block creation");
-    tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await;
-
     // Transfer from sender's ATA to recipient's ATA via the ATA program
     let transfer_amount = 50_u128;
     wallet::cli::execute_subcommand(
@@ -286,7 +226,7 @@ async fn transfer_and_burn_via_ata() -> Result<()> {
     tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await;
 
     // Verify sender ATA balance decreased
-    let sender_ata_acc = ctx.sequencer_client().get_account(sender_ata_id).await?;
+    let sender_ata_acc = get_account(&ctx, sender_ata_id).await?;
     let sender_holding = TokenHolding::try_from(&sender_ata_acc.data)?;
     assert_eq!(
         sender_holding,
@@ -297,7 +237,7 @@ async fn transfer_and_burn_via_ata() -> Result<()> {
     );
 
     // Verify recipient ATA balance increased
-    let recipient_ata_acc = ctx.sequencer_client().get_account(recipient_ata_id).await?;
+    let recipient_ata_acc = get_account(&ctx, recipient_ata_id).await?;
     let recipient_holding = TokenHolding::try_from(&recipient_ata_acc.data)?;
     assert_eq!(
         recipient_holding,
@@ -323,7 +263,7 @@ async fn transfer_and_burn_via_ata() -> Result<()> {
     tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await;
 
     // Verify sender ATA balance after burn
-    let sender_ata_acc = ctx.sequencer_client().get_account(sender_ata_id).await?;
+    let sender_ata_acc = get_account(&ctx, sender_ata_id).await?;
     let sender_holding = TokenHolding::try_from(&sender_ata_acc.data)?;
     assert_eq!(
         sender_holding,
@@ -334,10 +274,7 @@ async fn transfer_and_burn_via_ata() -> Result<()> {
     );
 
     // Verify the token definition total_supply decreased by burn_amount
-    let definition_acc = ctx
-        .sequencer_client()
-        .get_account(definition_account_id)
-        .await?;
+    let definition_acc = get_account(&ctx, definition_account_id).await?;
     let token_definition = TokenDefinition::try_from(&definition_acc.data)?;
     assert_eq!(
         token_definition,
@@ -355,25 +292,20 @@ async fn transfer_and_burn_via_ata() -> Result<()> {
 async fn create_ata_with_private_owner() -> Result<()> {
     let mut ctx = TestContext::new().await?;
 
-    let definition_account_id = new_public_account(&mut ctx).await?;
-    let supply_account_id = new_public_account(&mut ctx).await?;
-    let owner_account_id = new_private_account(&mut ctx).await?;
+    let definition_account_id = new_account(&mut ctx, false, None).await?;
+    let supply_account_id = new_account(&mut ctx, false, None).await?;
+    let owner_account_id = new_account(&mut ctx, true, None).await?;
 
     // Create a fungible token
-    wallet::cli::execute_subcommand(
-        ctx.wallet_mut(),
-        Command::Token(TokenProgramAgnosticSubcommand::New {
-            definition_account_id: public_mention(definition_account_id),
-            supply_account_id: public_mention(supply_account_id),
-            name: "TEST".to_owned(),
-            total_supply: 100,
-        }),
+    create_token(
+        &mut ctx,
+        public_mention(definition_account_id),
+        public_mention(supply_account_id),
+        "TEST".to_owned(),
+        100,
     )
     .await?;
 
-    info!("Waiting for next block creation");
-    tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await;
-
     // Create the ATA for the private owner + definition
     wallet::cli::execute_subcommand(
         ctx.wallet_mut(),
@@ -424,28 +356,23 @@ async fn create_ata_with_private_owner() -> Result<()> {
 async fn transfer_via_ata_private_owner() -> Result<()> {
     let mut ctx = TestContext::new().await?;
 
-    let definition_account_id = new_public_account(&mut ctx).await?;
-    let supply_account_id = new_public_account(&mut ctx).await?;
-    let sender_account_id = new_private_account(&mut ctx).await?;
-    let recipient_account_id = new_public_account(&mut ctx).await?;
+    let definition_account_id = new_account(&mut ctx, false, None).await?;
+    let supply_account_id = new_account(&mut ctx, false, None).await?;
+    let sender_account_id = new_account(&mut ctx, true, None).await?;
+    let recipient_account_id = new_account(&mut ctx, false, None).await?;
 
     let total_supply = 1000_u128;
 
     // Create a fungible token
-    wallet::cli::execute_subcommand(
-        ctx.wallet_mut(),
-        Command::Token(TokenProgramAgnosticSubcommand::New {
-            definition_account_id: public_mention(definition_account_id),
-            supply_account_id: public_mention(supply_account_id),
-            name: "TEST".to_owned(),
-            total_supply,
-        }),
+    create_token(
+        &mut ctx,
+        public_mention(definition_account_id),
+        public_mention(supply_account_id),
+        "TEST".to_owned(),
+        total_supply,
     )
     .await?;
 
-    info!("Waiting for next block creation");
-    tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await;
-
     // Derive ATA addresses
     let ata_program_id = programs::ata().id();
     let sender_ata_id = get_associated_token_account_id(
@@ -480,23 +407,14 @@ async fn transfer_via_ata_private_owner() -> Result<()> {
 
     // Fund sender's ATA from the supply account (direct token transfer)
     let fund_amount = 200_u128;
-    wallet::cli::execute_subcommand(
-        ctx.wallet_mut(),
-        Command::Token(TokenProgramAgnosticSubcommand::Send {
-            from: public_mention(supply_account_id),
-            to: Some(public_mention(sender_ata_id)),
-            to_npk: None,
-            to_vpk: None,
-            to_keys: None,
-            to_identifier: Some(0),
-            amount: fund_amount,
-        }),
+    token_send(
+        &mut ctx,
+        public_mention(supply_account_id),
+        public_mention(sender_ata_id),
+        fund_amount,
     )
     .await?;
 
-    info!("Waiting for next block creation");
-    tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await;
-
     // Transfer from sender's ATA (private owner) to recipient's ATA
     let transfer_amount = 50_u128;
     wallet::cli::execute_subcommand(
@@ -514,7 +432,7 @@ async fn transfer_via_ata_private_owner() -> Result<()> {
     tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await;
 
     // Verify sender ATA balance decreased
-    let sender_ata_acc = ctx.sequencer_client().get_account(sender_ata_id).await?;
+    let sender_ata_acc = get_account(&ctx, sender_ata_id).await?;
     let sender_holding = TokenHolding::try_from(&sender_ata_acc.data)?;
     assert_eq!(
         sender_holding,
@@ -525,7 +443,7 @@ async fn transfer_via_ata_private_owner() -> Result<()> {
     );
 
     // Verify recipient ATA balance increased
-    let recipient_ata_acc = ctx.sequencer_client().get_account(recipient_ata_id).await?;
+    let recipient_ata_acc = get_account(&ctx, recipient_ata_id).await?;
     let recipient_holding = TokenHolding::try_from(&recipient_ata_acc.data)?;
     assert_eq!(
         recipient_holding,
@@ -549,27 +467,22 @@ async fn transfer_via_ata_private_owner() -> Result<()> {
 async fn burn_via_ata_private_owner() -> Result<()> {
     let mut ctx = TestContext::new().await?;
 
-    let definition_account_id = new_public_account(&mut ctx).await?;
-    let supply_account_id = new_public_account(&mut ctx).await?;
-    let holder_account_id = new_private_account(&mut ctx).await?;
+    let definition_account_id = new_account(&mut ctx, false, None).await?;
+    let supply_account_id = new_account(&mut ctx, false, None).await?;
+    let holder_account_id = new_account(&mut ctx, true, None).await?;
 
     let total_supply = 500_u128;
 
     // Create a fungible token
-    wallet::cli::execute_subcommand(
-        ctx.wallet_mut(),
-        Command::Token(TokenProgramAgnosticSubcommand::New {
-            definition_account_id: public_mention(definition_account_id),
-            supply_account_id: public_mention(supply_account_id),
-            name: "TEST".to_owned(),
-            total_supply,
-        }),
+    create_token(
+        &mut ctx,
+        public_mention(definition_account_id),
+        public_mention(supply_account_id),
+        "TEST".to_owned(),
+        total_supply,
     )
     .await?;
 
-    info!("Waiting for next block creation");
-    tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await;
-
     // Derive holder's ATA address
     let ata_program_id = programs::ata().id();
     let holder_ata_id = get_associated_token_account_id(
@@ -592,23 +505,14 @@ async fn burn_via_ata_private_owner() -> Result<()> {
 
     // Fund holder's ATA from the supply account
     let fund_amount = 300_u128;
-    wallet::cli::execute_subcommand(
-        ctx.wallet_mut(),
-        Command::Token(TokenProgramAgnosticSubcommand::Send {
-            from: public_mention(supply_account_id),
-            to: Some(public_mention(holder_ata_id)),
-            to_npk: None,
-            to_vpk: None,
-            to_keys: None,
-            to_identifier: Some(0),
-            amount: fund_amount,
-        }),
+    token_send(
+        &mut ctx,
+        public_mention(supply_account_id),
+        public_mention(holder_ata_id),
+        fund_amount,
     )
     .await?;
 
-    info!("Waiting for next block creation");
-    tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await;
-
     // Burn from holder's ATA (private owner)
     let burn_amount = 100_u128;
     wallet::cli::execute_subcommand(
@@ -625,7 +529,7 @@ async fn burn_via_ata_private_owner() -> Result<()> {
     tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await;
 
     // Verify holder ATA balance after burn
-    let holder_ata_acc = ctx.sequencer_client().get_account(holder_ata_id).await?;
+    let holder_ata_acc = get_account(&ctx, holder_ata_id).await?;
     let holder_holding = TokenHolding::try_from(&holder_ata_acc.data)?;
     assert_eq!(
         holder_holding,
@@ -636,10 +540,7 @@ async fn burn_via_ata_private_owner() -> Result<()> {
     );
 
     // Verify the token definition total_supply decreased by burn_amount
-    let definition_acc = ctx
-        .sequencer_client()
-        .get_account(definition_account_id)
-        .await?;
+    let definition_acc = get_account(&ctx, definition_account_id).await?;
     let token_definition = TokenDefinition::try_from(&definition_acc.data)?;
     assert_eq!(
         token_definition,

integration_tests/tests/auth_transfer/private.rs

@@ -3,16 +3,19 @@ use std::time::Duration;
 use anyhow::{Context as _, Result};
 use common::transaction::LeeTransaction;
 use integration_tests::{
-    TIME_TO_WAIT_FOR_BLOCK_SECONDS, TestContext, fetch_privacy_preserving_tx, private_mention,
-    public_mention, verify_commitment_is_in_state,
+    TIME_TO_WAIT_FOR_BLOCK_SECONDS, TestContext, account_balance,
+    assert_private_commitment_in_state, fetch_privacy_preserving_tx, get_account, new_account,
+    private_mention, public_mention, send, sync_private, verify_commitment_is_in_state,
 };
 use lee::{
     AccountId, SharedSecretKey, execute_and_prove,
     privacy_preserving_transaction::circuit::ProgramWithDependencies, program::Program,
 };
 use lee_core::{
-    EncryptedAccountData, InputAccountIdentity, NullifierPublicKey,
-    account::AccountWithMetadata,
+    DUMMY_COMMITMENT, DUMMY_COMMITMENT_HASH, EncryptedAccountData, InputAccountIdentity,
+    ML_KEM_768_CIPHERTEXT_LEN, Nullifier, NullifierPublicKey,
+    account::{Account, AccountWithMetadata},
+    compute_digest_for_path,
     encryption::{EphemeralPublicKey, ViewingPublicKey},
 };
 use log::info;
@@ -34,32 +37,13 @@ async fn private_transfer_to_owned_account() -> Result<()> {
     let from: AccountId = ctx.existing_private_accounts()[0];
     let to: AccountId = ctx.existing_private_accounts()[1];
 
-    let command = Command::AuthTransfer(AuthTransferSubcommand::Send {
-        from: private_mention(from),
-        to: Some(private_mention(to)),
-        to_npk: None,
-        to_vpk: None,
-        to_keys: None,
-        to_identifier: Some(0),
-        amount: 100,
-    });
-
-    wallet::cli::execute_subcommand(ctx.wallet_mut(), command).await?;
+    send(&mut ctx, private_mention(from), private_mention(to), 100).await?;
 
     info!("Waiting for next block creation");
     tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await;
 
-    let new_commitment1 = ctx
-        .wallet()
-        .get_private_account_commitment(from)
-        .context("Failed to get private account commitment for sender")?;
-    assert!(verify_commitment_is_in_state(new_commitment1, ctx.sequencer_client()).await);
-
-    let new_commitment2 = ctx
-        .wallet()
-        .get_private_account_commitment(to)
-        .context("Failed to get private account commitment for receiver")?;
-    assert!(verify_commitment_is_in_state(new_commitment2, ctx.sequencer_client()).await);
+    assert_private_commitment_in_state(&ctx, from, "sender").await?;
+    assert_private_commitment_in_state(&ctx, to, "receiver").await?;
 
     info!("Successfully transferred privately to owned account");
 
@@ -125,17 +109,7 @@ async fn deshielded_transfer_to_public_account() -> Result<()> {
         .context("Failed to get sender's private account")?;
     assert_eq!(from_acc.balance, 10000);
 
-    let command = Command::AuthTransfer(AuthTransferSubcommand::Send {
-        from: private_mention(from),
-        to: Some(public_mention(to)),
-        to_npk: None,
-        to_vpk: None,
-        to_keys: None,
-        to_identifier: Some(0),
-        amount: 100,
-    });
-
-    wallet::cli::execute_subcommand(ctx.wallet_mut(), command).await?;
+    send(&mut ctx, private_mention(from), public_mention(to), 100).await?;
 
     info!("Waiting for next block creation");
     tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await;
@@ -144,13 +118,9 @@ async fn deshielded_transfer_to_public_account() -> Result<()> {
         .wallet()
         .get_account_private(from)
         .context("Failed to get sender's private account")?;
-    let new_commitment = ctx
-        .wallet()
-        .get_private_account_commitment(from)
-        .context("Failed to get private account commitment")?;
-    assert!(verify_commitment_is_in_state(new_commitment, ctx.sequencer_client()).await);
+    assert_private_commitment_in_state(&ctx, from, "sender").await?;
 
-    let acc_2_balance = ctx.sequencer_client().get_account_balance(to).await?;
+    let acc_2_balance = account_balance(&ctx, to).await?;
 
     assert_eq!(from_acc.balance, 9900);
     assert_eq!(acc_2_balance, 20100);
@@ -167,18 +137,7 @@ async fn private_transfer_to_owned_account_using_claiming_path() -> Result<()> {
     let from: AccountId = ctx.existing_private_accounts()[0];
 
     // Create a new private account
-    let command = Command::Account(AccountSubcommand::New(NewSubcommand::Private {
-        cci: None,
-        label: None,
-    }));
-
-    let sub_ret = wallet::cli::execute_subcommand(ctx.wallet_mut(), command).await?;
-    let SubcommandReturnValue::RegisterAccount {
-        account_id: to_account_id,
-    } = sub_ret
-    else {
-        anyhow::bail!("Expected RegisterAccount return value");
-    };
+    let to_account_id = new_account(&mut ctx, true, None).await?;
 
     // Get the keys for the newly created account
     let to = ctx
@@ -207,14 +166,13 @@ async fn private_transfer_to_owned_account_using_claiming_path() -> Result<()> {
     let tx = fetch_privacy_preserving_tx(ctx.sequencer_client(), tx_hash).await;
 
     // Sync the wallet to claim the new account
-    let command = Command::Account(AccountSubcommand::SyncPrivate {});
-    wallet::cli::execute_subcommand(ctx.wallet_mut(), command).await?;
+    sync_private(&mut ctx).await?;
 
-    let new_commitment1 = ctx
+    let sender_commitment = ctx
         .wallet()
         .get_private_account_commitment(from)
         .context("Failed to get private account commitment for sender")?;
-    assert_eq!(tx.message.new_commitments[0], new_commitment1);
+    assert_eq!(tx.message.new_commitments[0], sender_commitment);
 
     assert_eq!(tx.message.new_commitments.len(), 2);
     for commitment in tx.message.new_commitments {
@@ -239,17 +197,7 @@ async fn shielded_transfer_to_owned_private_account() -> Result<()> {
     let from: AccountId = ctx.existing_public_accounts()[0];
     let to: AccountId = ctx.existing_private_accounts()[1];
 
-    let command = Command::AuthTransfer(AuthTransferSubcommand::Send {
-        from: public_mention(from),
-        to: Some(private_mention(to)),
-        to_npk: None,
-        to_vpk: None,
-        to_keys: None,
-        to_identifier: Some(0),
-        amount: 100,
-    });
-
-    wallet::cli::execute_subcommand(ctx.wallet_mut(), command).await?;
+    send(&mut ctx, public_mention(from), private_mention(to), 100).await?;
 
     info!("Waiting for next block creation");
     tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await;
@@ -258,13 +206,9 @@ async fn shielded_transfer_to_owned_private_account() -> Result<()> {
         .wallet()
         .get_account_private(to)
         .context("Failed to get receiver's private account")?;
-    let new_commitment = ctx
-        .wallet()
-        .get_private_account_commitment(to)
-        .context("Failed to get receiver's commitment")?;
-    assert!(verify_commitment_is_in_state(new_commitment, ctx.sequencer_client()).await);
+    assert_private_commitment_in_state(&ctx, to, "receiver").await?;
 
-    let acc_from_balance = ctx.sequencer_client().get_account_balance(from).await?;
+    let acc_from_balance = account_balance(&ctx, from).await?;
 
     assert_eq!(acc_from_balance, 9900);
     assert_eq!(acc_to.balance, 20100);
@@ -303,7 +247,7 @@ async fn shielded_transfer_to_foreign_account() -> Result<()> {
 
     let tx = fetch_privacy_preserving_tx(ctx.sequencer_client(), tx_hash).await;
 
-    let acc_1_balance = ctx.sequencer_client().get_account_balance(from).await?;
+    let acc_1_balance = account_balance(&ctx, from).await?;
 
     assert!(
         verify_commitment_is_in_state(
@@ -332,18 +276,7 @@ async fn private_transfer_to_owned_account_continuous_run_path() -> Result<()> {
     let from: AccountId = ctx.existing_private_accounts()[0];
 
     // Create a new private account
-    let command = Command::Account(AccountSubcommand::New(NewSubcommand::Private {
-        cci: None,
-        label: None,
-    }));
-    let sub_ret = wallet::cli::execute_subcommand(ctx.wallet_mut(), command).await?;
-
-    let SubcommandReturnValue::RegisterAccount {
-        account_id: to_account_id,
-    } = sub_ret
-    else {
-        anyhow::bail!("Failed to register account");
-    };
+    let to_account_id = new_account(&mut ctx, true, None).await?;
 
     // Get the newly created account's keys
     let to = ctx
@@ -396,14 +329,7 @@ async fn private_transfer_to_owned_account_continuous_run_path() -> Result<()> {
 async fn initialize_private_account() -> Result<()> {
     let mut ctx = TestContext::new().await?;
 
-    let command = Command::Account(AccountSubcommand::New(NewSubcommand::Private {
-        cci: None,
-        label: None,
-    }));
-    let result = wallet::cli::execute_subcommand(ctx.wallet_mut(), command).await?;
-    let SubcommandReturnValue::RegisterAccount { account_id } = result else {
-        anyhow::bail!("Expected RegisterAccount return value");
-    };
+    let account_id = new_account(&mut ctx, true, None).await?;
 
     let command = Command::AuthTransfer(AuthTransferSubcommand::Init {
         account_id: private_mention(account_id),
@@ -413,14 +339,9 @@ async fn initialize_private_account() -> Result<()> {
     tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await;
 
     info!("Syncing private accounts");
-    let command = Command::Account(AccountSubcommand::SyncPrivate {});
-    wallet::cli::execute_subcommand(ctx.wallet_mut(), command).await?;
+    sync_private(&mut ctx).await?;
 
-    let new_commitment = ctx
-        .wallet()
-        .get_private_account_commitment(account_id)
-        .context("Failed to get private account commitment")?;
-    assert!(verify_commitment_is_in_state(new_commitment, ctx.sequencer_client()).await);
+    assert_private_commitment_in_state(&ctx, account_id, "account").await?;
 
     let account = ctx
         .wallet()
@@ -455,32 +376,19 @@ async fn private_transfer_using_from_label() -> Result<()> {
     wallet::cli::execute_subcommand(ctx.wallet_mut(), command).await?;
 
     // Send using the label instead of account ID
-    let command = Command::AuthTransfer(AuthTransferSubcommand::Send {
-        from: CliAccountMention::Label(label),
-        to: Some(private_mention(to)),
-        to_npk: None,
-        to_vpk: None,
-        to_keys: None,
-        to_identifier: Some(0),
-        amount: 100,
-    });
-
-    wallet::cli::execute_subcommand(ctx.wallet_mut(), command).await?;
+    send(
+        &mut ctx,
+        CliAccountMention::Label(label),
+        private_mention(to),
+        100,
+    )
+    .await?;
 
     info!("Waiting for next block creation");
     tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await;
 
-    let new_commitment1 = ctx
-        .wallet()
-        .get_private_account_commitment(from)
-        .context("Failed to get private account commitment for sender")?;
-    assert!(verify_commitment_is_in_state(new_commitment1, ctx.sequencer_client()).await);
-
-    let new_commitment2 = ctx
-        .wallet()
-        .get_private_account_commitment(to)
-        .context("Failed to get private account commitment for receiver")?;
-    assert!(verify_commitment_is_in_state(new_commitment2, ctx.sequencer_client()).await);
+    assert_private_commitment_in_state(&ctx, from, "sender").await?;
+    assert_private_commitment_in_state(&ctx, to, "receiver").await?;
 
     info!("Successfully transferred privately using from_label");
 
@@ -510,14 +418,9 @@ async fn initialize_private_account_using_label() -> Result<()> {
 
     tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await;
 
-    let command = Command::Account(AccountSubcommand::SyncPrivate {});
-    wallet::cli::execute_subcommand(ctx.wallet_mut(), command).await?;
+    sync_private(&mut ctx).await?;
 
-    let new_commitment = ctx
-        .wallet()
-        .get_private_account_commitment(account_id)
-        .context("Failed to get private account commitment")?;
-    assert!(verify_commitment_is_in_state(new_commitment, ctx.sequencer_client()).await);
+    assert_private_commitment_in_state(&ctx, account_id, "account").await?;
 
     let account = ctx
         .wallet()
@@ -593,11 +496,7 @@ async fn shielded_transfers_to_two_identifiers_same_npk() -> Result<()> {
     info!("Waiting for next block creation");
     tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await;
 
-    wallet::cli::execute_subcommand(
-        ctx.wallet_mut(),
-        Command::Account(AccountSubcommand::SyncPrivate {}),
-    )
-    .await?;
+    sync_private(&mut ctx).await?;
 
     // Both accounts must be discovered with the correct balances.
     let account_id_1 = AccountId::for_regular_private_account(&npk, identifier_1);
@@ -664,7 +563,7 @@ async fn ppt_cant_chain_call_faucet() -> Result<()> {
     let npk = NullifierPublicKey::from(&nsk);
     let vpk = ViewingPublicKey::from_bytes(vec![4_u8; 1184]).unwrap();
     let ssk = SharedSecretKey([55_u8; 32]);
-    let epk = EphemeralPublicKey(vec![55_u8; 1088]);
+    let epk = EphemeralPublicKey(vec![55_u8; ML_KEM_768_CIPHERTEXT_LEN]);
     let attacker_vault_id = {
         let seed = vault_core::compute_vault_seed(attacker_id);
         AccountId::for_private_pda(&vault_program_id, &seed, &npk, 1337)
@@ -672,16 +571,12 @@ async fn ppt_cant_chain_call_faucet() -> Result<()> {
     let amount: u128 = 1;
 
     let faucet_pre = AccountWithMetadata::new(
-        ctx.sequencer_client()
-            .get_account(faucet_account_id)
-            .await?,
+        get_account(&ctx, faucet_account_id).await?,
         false,
         faucet_account_id,
     );
     let vault_pda_pre = AccountWithMetadata::new(
-        ctx.sequencer_client()
-            .get_account(attacker_vault_id)
-            .await?,
+        get_account(&ctx, attacker_vault_id).await?,
         false,
         attacker_vault_id,
     );
@@ -710,6 +605,7 @@ async fn ppt_cant_chain_call_faucet() -> Result<()> {
                 npk,
                 ssk,
                 identifier: 1337,
+                commitment_root: DUMMY_COMMITMENT_HASH,
                 seed: None,
             },
         ],
@@ -720,3 +616,100 @@ async fn ppt_cant_chain_call_faucet() -> Result<()> {
 
     Ok(())
 }
+
+async fn prove_init_with_commitment_root(
+    ctx: &TestContext,
+    commitment_root: lee_core::CommitmentSetDigest,
+) -> Result<lee_core::PrivacyPreservingCircuitOutput> {
+    let program = programs::authenticated_transfer();
+    let sender_id = ctx.existing_public_accounts()[0];
+    let sender_pre = AccountWithMetadata::new(
+        ctx.sequencer_client().get_account(sender_id).await?,
+        true,
+        sender_id,
+    );
+
+    let nsk: lee_core::NullifierSecretKey = [7; 32];
+    let npk = NullifierPublicKey::from(&nsk);
+    let vpk = ViewingPublicKey::from_bytes(vec![4_u8; 1184]).unwrap();
+    let ssk = SharedSecretKey([55_u8; 32]);
+    let recipient_account_id = AccountId::for_regular_private_account(&npk, 0);
+    let recipient = AccountWithMetadata::new(Account::default(), false, recipient_account_id);
+
+    let (output, _) = execute_and_prove(
+        vec![sender_pre, recipient],
+        Program::serialize_instruction(authenticated_transfer_core::Instruction::Transfer {
+            amount: 1,
+        })?,
+        vec![
+            InputAccountIdentity::Public,
+            InputAccountIdentity::PrivateUnauthorized {
+                epk: EphemeralPublicKey(Vec::new()),
+                view_tag: EncryptedAccountData::compute_view_tag(&npk, &vpk),
+                npk,
+                ssk,
+                identifier: 0,
+                commitment_root,
+            },
+        ],
+        &program.into(),
+    )?;
+
+    Ok(output)
+}
+
+#[test]
+async fn init_with_dummy_commitment_root_produces_valid_root() -> Result<()> {
+    let ctx = TestContext::new().await?;
+
+    let dummy_proof = ctx
+        .sequencer_client()
+        .get_proof_for_commitment(DUMMY_COMMITMENT)
+        .await?
+        .expect("DUMMY_COMMITMENT must be in genesis commitment set");
+    let expected_digest = compute_digest_for_path(&DUMMY_COMMITMENT, &dummy_proof);
+
+    let nsk: lee_core::NullifierSecretKey = [7; 32];
+    let npk = NullifierPublicKey::from(&nsk);
+    let recipient_account_id = AccountId::for_regular_private_account(&npk, 0);
+
+    let output = prove_init_with_commitment_root(&ctx, expected_digest).await?;
+
+    assert_eq!(output.new_nullifiers.len(), 1);
+    let (nullifier, digest) = &output.new_nullifiers[0];
+    assert_eq!(
+        *nullifier,
+        Nullifier::for_account_initialization(&recipient_account_id)
+    );
+    assert_eq!(*digest, expected_digest);
+    assert_ne!(*digest, DUMMY_COMMITMENT_HASH);
+
+    Ok(())
+}
+
+#[test]
+async fn init_nullifier_digest_is_bound_to_commitment_root() -> Result<()> {
+    let ctx = TestContext::new().await?;
+
+    let dummy_proof = ctx
+        .sequencer_client()
+        .get_proof_for_commitment(DUMMY_COMMITMENT)
+        .await?
+        .expect("DUMMY_COMMITMENT must be in genesis commitment set");
+    let expected_digest = compute_digest_for_path(&DUMMY_COMMITMENT, &dummy_proof);
+
+    let output_with_root = prove_init_with_commitment_root(&ctx, expected_digest).await?;
+    let output_without_root = prove_init_with_commitment_root(&ctx, DUMMY_COMMITMENT_HASH).await?;
+
+    assert_eq!(output_with_root.new_nullifiers[0].1, expected_digest);
+    assert_eq!(
+        output_without_root.new_nullifiers[0].1,
+        DUMMY_COMMITMENT_HASH
+    );
+    assert_ne!(
+        output_with_root.new_nullifiers[0].1,
+        output_without_root.new_nullifiers[0].1,
+    );
+
+    Ok(())
+}

integration_tests/tests/auth_transfer/public.rs

@@ -2,7 +2,10 @@ use std::time::Duration;
 
 use anyhow::Result;
 use common::transaction::LeeTransaction;
-use integration_tests::{TIME_TO_WAIT_FOR_BLOCK_SECONDS, TestContext, public_mention};
+use integration_tests::{
+    TIME_TO_WAIT_FOR_BLOCK_SECONDS, TestContext, account_balance, get_account, new_account,
+    public_mention, send,
+};
 use lee::public_transaction;
 use log::info;
 use sequencer_service_rpc::RpcClient as _;
@@ -10,8 +13,7 @@ use tokio::test;
 use wallet::{
     account::Label,
     cli::{
-        CliAccountMention, Command, SubcommandReturnValue,
-        account::{AccountSubcommand, NewSubcommand},
+        CliAccountMention, Command, account::AccountSubcommand,
         programs::native_token_transfer::AuthTransferSubcommand,
     },
 };
@@ -20,30 +22,22 @@ use wallet::{
 async fn successful_transfer_to_existing_account() -> Result<()> {
     let mut ctx = TestContext::new().await?;
 
-    let command = Command::AuthTransfer(AuthTransferSubcommand::Send {
-        from: public_mention(ctx.existing_public_accounts()[0]),
-        to: Some(public_mention(ctx.existing_public_accounts()[1])),
-        to_npk: None,
-        to_vpk: None,
-        to_keys: None,
-        to_identifier: Some(0),
-        amount: 100,
-    });
-
-    wallet::cli::execute_subcommand(ctx.wallet_mut(), command).await?;
+    let sender = ctx.existing_public_accounts()[0];
+    let receiver = ctx.existing_public_accounts()[1];
+    send(
+        &mut ctx,
+        public_mention(sender),
+        public_mention(receiver),
+        100,
+    )
+    .await?;
 
     info!("Waiting for next block creation");
     tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await;
 
     info!("Checking correct balance move");
-    let acc_1_balance = ctx
-        .sequencer_client()
-        .get_account_balance(ctx.existing_public_accounts()[0])
-        .await?;
-    let acc_2_balance = ctx
-        .sequencer_client()
-        .get_account_balance(ctx.existing_public_accounts()[1])
-        .await?;
+    let acc_1_balance = account_balance(&ctx, sender).await?;
+    let acc_2_balance = account_balance(&ctx, receiver).await?;
 
     info!("Balance of sender: {acc_1_balance:#?}");
     info!("Balance of receiver: {acc_2_balance:#?}");
@@ -58,51 +52,23 @@ async fn successful_transfer_to_existing_account() -> Result<()> {
 pub async fn successful_transfer_to_new_account() -> Result<()> {
     let mut ctx = TestContext::new().await?;
 
-    let command = Command::Account(AccountSubcommand::New(NewSubcommand::Public {
-        cci: None,
-        label: None,
-    }));
+    let new_persistent_account_id = new_account(&mut ctx, false, None).await?;
 
-    wallet::cli::execute_subcommand(ctx.wallet_mut(), command)
-        .await
-        .unwrap();
-
-    let new_persistent_account_id = ctx
-        .wallet()
-        .storage()
-        .key_chain()
-        .public_account_ids()
-        .map(|(account_id, _)| account_id)
-        .find(|acc_id| {
-            *acc_id != ctx.existing_public_accounts()[0]
-                && *acc_id != ctx.existing_public_accounts()[1]
-        })
-        .expect("Failed to find newly created account in the wallet storage");
-
-    let command = Command::AuthTransfer(AuthTransferSubcommand::Send {
-        from: public_mention(ctx.existing_public_accounts()[0]),
-        to: Some(public_mention(new_persistent_account_id)),
-        to_npk: None,
-        to_vpk: None,
-        to_keys: None,
-        to_identifier: Some(0),
-        amount: 100,
-    });
-
-    wallet::cli::execute_subcommand(ctx.wallet_mut(), command).await?;
+    let sender = ctx.existing_public_accounts()[0];
+    send(
+        &mut ctx,
+        public_mention(sender),
+        public_mention(new_persistent_account_id),
+        100,
+    )
+    .await?;
 
     info!("Waiting for next block creation");
     tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await;
 
     info!("Checking correct balance move");
-    let acc_1_balance = ctx
-        .sequencer_client()
-        .get_account_balance(ctx.existing_public_accounts()[0])
-        .await?;
-    let acc_2_balance = ctx
-        .sequencer_client()
-        .get_account_balance(new_persistent_account_id)
-        .await?;
+    let acc_1_balance = account_balance(&ctx, sender).await?;
+    let acc_2_balance = account_balance(&ctx, new_persistent_account_id).await?;
 
     info!("Balance of sender: {acc_1_balance:#?}");
     info!("Balance of receiver: {acc_2_balance:#?}");
@@ -134,14 +100,8 @@ async fn failed_transfer_with_insufficient_balance() -> Result<()> {
     tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await;
 
     info!("Checking balances unchanged");
-    let acc_1_balance = ctx
-        .sequencer_client()
-        .get_account_balance(ctx.existing_public_accounts()[0])
-        .await?;
-    let acc_2_balance = ctx
-        .sequencer_client()
-        .get_account_balance(ctx.existing_public_accounts()[1])
-        .await?;
+    let acc_1_balance = account_balance(&ctx, ctx.existing_public_accounts()[0]).await?;
+    let acc_2_balance = account_balance(&ctx, ctx.existing_public_accounts()[1]).await?;
 
     info!("Balance of sender: {acc_1_balance:#?}");
     info!("Balance of receiver: {acc_2_balance:#?}");
@@ -156,31 +116,24 @@ async fn failed_transfer_with_insufficient_balance() -> Result<()> {
 async fn two_consecutive_successful_transfers() -> Result<()> {
     let mut ctx = TestContext::new().await?;
 
-    // First transfer
-    let command = Command::AuthTransfer(AuthTransferSubcommand::Send {
-        from: public_mention(ctx.existing_public_accounts()[0]),
-        to: Some(public_mention(ctx.existing_public_accounts()[1])),
-        to_npk: None,
-        to_vpk: None,
-        to_keys: None,
-        to_identifier: Some(0),
-        amount: 100,
-    });
+    let sender = ctx.existing_public_accounts()[0];
+    let receiver = ctx.existing_public_accounts()[1];
 
-    wallet::cli::execute_subcommand(ctx.wallet_mut(), command).await?;
+    // First transfer
+    send(
+        &mut ctx,
+        public_mention(sender),
+        public_mention(receiver),
+        100,
+    )
+    .await?;
 
     info!("Waiting for next block creation");
     tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await;
 
     info!("Checking correct balance move after first transfer");
-    let acc_1_balance = ctx
-        .sequencer_client()
-        .get_account_balance(ctx.existing_public_accounts()[0])
-        .await?;
-    let acc_2_balance = ctx
-        .sequencer_client()
-        .get_account_balance(ctx.existing_public_accounts()[1])
-        .await?;
+    let acc_1_balance = account_balance(&ctx, sender).await?;
+    let acc_2_balance = account_balance(&ctx, receiver).await?;
 
     info!("Balance of sender: {acc_1_balance:#?}");
     info!("Balance of receiver: {acc_2_balance:#?}");
@@ -191,30 +144,20 @@ async fn two_consecutive_successful_transfers() -> Result<()> {
     info!("First TX Success!");
 
     // Second transfer
-    let command = Command::AuthTransfer(AuthTransferSubcommand::Send {
-        from: public_mention(ctx.existing_public_accounts()[0]),
-        to: Some(public_mention(ctx.existing_public_accounts()[1])),
-        to_npk: None,
-        to_vpk: None,
-        to_keys: None,
-        to_identifier: Some(0),
-        amount: 100,
-    });
-
-    wallet::cli::execute_subcommand(ctx.wallet_mut(), command).await?;
+    send(
+        &mut ctx,
+        public_mention(sender),
+        public_mention(receiver),
+        100,
+    )
+    .await?;
 
     info!("Waiting for next block creation");
     tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await;
 
     info!("Checking correct balance move after second transfer");
-    let acc_1_balance = ctx
-        .sequencer_client()
-        .get_account_balance(ctx.existing_public_accounts()[0])
-        .await?;
-    let acc_2_balance = ctx
-        .sequencer_client()
-        .get_account_balance(ctx.existing_public_accounts()[1])
-        .await?;
+    let acc_1_balance = account_balance(&ctx, sender).await?;
+    let acc_2_balance = account_balance(&ctx, receiver).await?;
 
     info!("Balance of sender: {acc_1_balance:#?}");
     info!("Balance of receiver: {acc_2_balance:#?}");
@@ -231,14 +174,7 @@ async fn two_consecutive_successful_transfers() -> Result<()> {
 async fn initialize_public_account() -> Result<()> {
     let mut ctx = TestContext::new().await?;
 
-    let command = Command::Account(AccountSubcommand::New(NewSubcommand::Public {
-        cci: None,
-        label: None,
-    }));
-    let result = wallet::cli::execute_subcommand(ctx.wallet_mut(), command).await?;
-    let SubcommandReturnValue::RegisterAccount { account_id } = result else {
-        anyhow::bail!("Expected RegisterAccount return value");
-    };
+    let account_id = new_account(&mut ctx, false, None).await?;
 
     let command = Command::AuthTransfer(AuthTransferSubcommand::Init {
         account_id: public_mention(account_id),
@@ -246,7 +182,7 @@ async fn initialize_public_account() -> Result<()> {
     wallet::cli::execute_subcommand(ctx.wallet_mut(), command).await?;
 
     info!("Checking correct execution");
-    let account = ctx.sequencer_client().get_account(account_id).await?;
+    let account = get_account(&ctx, account_id).await?;
 
     assert_eq!(
         account.program_owner,
@@ -274,30 +210,22 @@ async fn successful_transfer_using_from_label() -> Result<()> {
     wallet::cli::execute_subcommand(ctx.wallet_mut(), command).await?;
 
     // Send using the label instead of account ID
-    let command = Command::AuthTransfer(AuthTransferSubcommand::Send {
-        from: CliAccountMention::Label(label),
-        to: Some(public_mention(ctx.existing_public_accounts()[1])),
-        to_npk: None,
-        to_vpk: None,
-        to_keys: None,
-        to_identifier: Some(0),
-        amount: 100,
-    });
-
-    wallet::cli::execute_subcommand(ctx.wallet_mut(), command).await?;
+    let sender = ctx.existing_public_accounts()[0];
+    let receiver = ctx.existing_public_accounts()[1];
+    send(
+        &mut ctx,
+        CliAccountMention::Label(label),
+        public_mention(receiver),
+        100,
+    )
+    .await?;
 
     info!("Waiting for next block creation");
     tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await;
 
     info!("Checking correct balance move");
-    let acc_1_balance = ctx
-        .sequencer_client()
-        .get_account_balance(ctx.existing_public_accounts()[0])
-        .await?;
-    let acc_2_balance = ctx
-        .sequencer_client()
-        .get_account_balance(ctx.existing_public_accounts()[1])
-        .await?;
+    let acc_1_balance = account_balance(&ctx, sender).await?;
+    let acc_2_balance = account_balance(&ctx, receiver).await?;
 
     assert_eq!(acc_1_balance, 9900);
     assert_eq!(acc_2_balance, 20100);
@@ -320,30 +248,22 @@ async fn successful_transfer_using_to_label() -> Result<()> {
     wallet::cli::execute_subcommand(ctx.wallet_mut(), command).await?;
 
     // Send using the label for the recipient
-    let command = Command::AuthTransfer(AuthTransferSubcommand::Send {
-        from: public_mention(ctx.existing_public_accounts()[0]),
-        to: Some(CliAccountMention::Label(label)),
-        to_npk: None,
-        to_vpk: None,
-        to_keys: None,
-        to_identifier: Some(0),
-        amount: 100,
-    });
-
-    wallet::cli::execute_subcommand(ctx.wallet_mut(), command).await?;
+    let sender = ctx.existing_public_accounts()[0];
+    let receiver = ctx.existing_public_accounts()[1];
+    send(
+        &mut ctx,
+        public_mention(sender),
+        CliAccountMention::Label(label),
+        100,
+    )
+    .await?;
 
     info!("Waiting for next block creation");
     tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await;
 
     info!("Checking correct balance move");
-    let acc_1_balance = ctx
-        .sequencer_client()
-        .get_account_balance(ctx.existing_public_accounts()[0])
-        .await?;
-    let acc_2_balance = ctx
-        .sequencer_client()
-        .get_account_balance(ctx.existing_public_accounts()[1])
-        .await?;
+    let acc_1_balance = account_balance(&ctx, sender).await?;
+    let acc_2_balance = account_balance(&ctx, receiver).await?;
 
     assert_eq!(acc_1_balance, 9900);
     assert_eq!(acc_2_balance, 20100);
@@ -359,14 +279,8 @@ async fn cannot_transfer_funds_from_system_faucet_account() -> Result<()> {
     let faucet_account_id = system_accounts::faucet_account_id();
 
     let recipient = ctx.existing_public_accounts()[0];
-    let recipient_balance_before = ctx
-        .sequencer_client()
-        .get_account_balance(recipient)
-        .await?;
-    let faucet_balance_before = ctx
-        .sequencer_client()
-        .get_account_balance(faucet_account_id)
-        .await?;
+    let recipient_balance_before = account_balance(&ctx, recipient).await?;
+    let faucet_balance_before = account_balance(&ctx, faucet_account_id).await?;
 
     let amount = 1_u128;
     let message = public_transaction::Message::try_new(
@@ -387,14 +301,8 @@ async fn cannot_transfer_funds_from_system_faucet_account() -> Result<()> {
     info!("Waiting for next block creation");
     tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await;
 
-    let recipient_balance_after = ctx
-        .sequencer_client()
-        .get_account_balance(recipient)
-        .await?;
-    let faucet_balance_after = ctx
-        .sequencer_client()
-        .get_account_balance(faucet_account_id)
-        .await?;
+    let recipient_balance_after = account_balance(&ctx, recipient).await?;
+    let faucet_balance_after = account_balance(&ctx, faucet_account_id).await?;
     let tx_on_chain = ctx.sequencer_client().get_transaction(tx_hash).await?;
 
     assert_eq!(recipient_balance_after, recipient_balance_before);
@@ -413,14 +321,8 @@ async fn cannot_execute_faucet_program() -> Result<()> {
     let vault_program_id = programs::vault().id();
     let recipient_vault_id = vault_core::compute_vault_account_id(vault_program_id, recipient);
 
-    let recipient_balance_before = ctx
-        .sequencer_client()
-        .get_account_balance(recipient)
-        .await?;
-    let faucet_balance_before = ctx
-        .sequencer_client()
-        .get_account_balance(faucet_account_id)
-        .await?;
+    let recipient_balance_before = account_balance(&ctx, recipient).await?;
+    let faucet_balance_before = account_balance(&ctx, faucet_account_id).await?;
 
     let amount = 1_u128;
     let message = public_transaction::Message::try_new(
@@ -445,14 +347,8 @@ async fn cannot_execute_faucet_program() -> Result<()> {
     info!("Waiting for next block creation");
     tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await;
 
-    let recipient_balance_after = ctx
-        .sequencer_client()
-        .get_account_balance(recipient)
-        .await?;
-    let faucet_balance_after = ctx
-        .sequencer_client()
-        .get_account_balance(faucet_account_id)
-        .await?;
+    let recipient_balance_after = account_balance(&ctx, recipient).await?;
+    let faucet_balance_after = account_balance(&ctx, faucet_account_id).await?;
     let tx_on_chain = ctx.sequencer_client().get_transaction(tx_hash).await?;
 
     assert_eq!(recipient_balance_after, recipient_balance_before);
@@ -493,28 +389,16 @@ async fn user_tx_that_chain_calls_faucet_is_dropped() -> Result<()> {
         lee::public_transaction::WitnessSet::from_raw_parts(vec![]),
     ));
 
-    let faucet_balance_before = ctx
-        .sequencer_client()
-        .get_account_balance(faucet_account_id)
-        .await?;
-    let vault_balance_before = ctx
-        .sequencer_client()
-        .get_account_balance(attacker_vault_id)
-        .await?;
+    let faucet_balance_before = account_balance(&ctx, faucet_account_id).await?;
+    let vault_balance_before = account_balance(&ctx, attacker_vault_id).await?;
 
     let tx_hash = ctx.sequencer_client().send_transaction(attack_tx).await?;
 
     info!("Waiting for next block creation");
     tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await;
 
-    let faucet_balance_after = ctx
-        .sequencer_client()
-        .get_account_balance(faucet_account_id)
-        .await?;
-    let vault_balance_after = ctx
-        .sequencer_client()
-        .get_account_balance(attacker_vault_id)
-        .await?;
+    let faucet_balance_after = account_balance(&ctx, faucet_account_id).await?;
+    let vault_balance_after = account_balance(&ctx, attacker_vault_id).await?;
     let tx_on_chain = ctx.sequencer_client().get_transaction(tx_hash).await?;
 
     assert_eq!(faucet_balance_after, faucet_balance_before);

integration_tests/tests/bridge.rs

@@ -11,7 +11,8 @@ use borsh::BorshSerialize;
 use common::transaction::LeeTransaction;
 use futures::StreamExt as _;
 use integration_tests::{
-    TIME_TO_WAIT_FOR_BLOCK_SECONDS, TestContext, wait_for_indexer_to_catch_up,
+    TIME_TO_WAIT_FOR_BLOCK_SECONDS, TestContext, account_balance, get_account,
+    wait_for_indexer_to_catch_up,
 };
 use lee::{
     AccountId, execute_and_prove, privacy_preserving_transaction, program::Program,
@@ -65,6 +66,54 @@ async fn public_bridge_deposit_invocation_is_dropped() -> anyhow::Result<()> {
         lee::public_transaction::WitnessSet::from_raw_parts(vec![]),
     ));
 
+    let bridge_balance_before = account_balance(&ctx, bridge_account_id).await?;
+    let vault_balance_before = account_balance(&ctx, recipient_vault_id).await?;
+
+    let tx_hash = ctx.sequencer_client().send_transaction(attack_tx).await?;
+
+    tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await;
+
+    let bridge_balance_after = account_balance(&ctx, bridge_account_id).await?;
+    let vault_balance_after = account_balance(&ctx, recipient_vault_id).await?;
+    let tx_on_chain = ctx.sequencer_client().get_transaction(tx_hash).await?;
+
+    assert_eq!(bridge_balance_after, bridge_balance_before);
+    assert_eq!(vault_balance_after, vault_balance_before);
+    assert!(
+        tx_on_chain.is_none(),
+        "Direct public bridge::Deposit invocation should be rejected"
+    );
+
+    Ok(())
+}
+
+#[test]
+async fn public_bridge_deposit_with_zero_amount_is_rejected() -> anyhow::Result<()> {
+    let ctx = TestContext::new().await?;
+
+    let recipient_id = ctx.existing_public_accounts()[0];
+    let bridge_account_id = system_accounts::bridge_account_id();
+    let vault_program_id = programs::vault().id();
+    let recipient_vault_id = vault_core::compute_vault_account_id(vault_program_id, recipient_id);
+
+    let message = public_transaction::Message::try_new(
+        programs::bridge().id(),
+        vec![bridge_account_id, recipient_vault_id],
+        vec![],
+        bridge_core::Instruction::Deposit {
+            l1_deposit_op_id: [0_u8; 32],
+            vault_program_id,
+            recipient_id,
+            amount: 0,
+        },
+    )
+    .context("Failed to build zero-amount public bridge deposit transaction")?;
+
+    let attack_tx = LeeTransaction::Public(lee::PublicTransaction::new(
+        message,
+        lee::public_transaction::WitnessSet::from_raw_parts(vec![]),
+    ));
+
     let bridge_balance_before = ctx
         .sequencer_client()
         .get_account_balance(bridge_account_id)
@@ -92,7 +141,7 @@ async fn public_bridge_deposit_invocation_is_dropped() -> anyhow::Result<()> {
     assert_eq!(vault_balance_after, vault_balance_before);
     assert!(
         tx_on_chain.is_none(),
-        "Direct public bridge::Deposit invocation should be rejected"
+        "Public bridge::Deposit with zero amount should be rejected"
     );
 
     Ok(())
@@ -109,16 +158,12 @@ async fn private_bridge_deposit_invocation_is_dropped() -> anyhow::Result<()> {
 
     // Get pre-state of bridge and vault accounts
     let bridge_pre = AccountWithMetadata::new(
-        ctx.sequencer_client()
-            .get_account(bridge_account_id)
-            .await?,
+        get_account(&ctx, bridge_account_id).await?,
         false,
         bridge_account_id,
     );
     let vault_pre = AccountWithMetadata::new(
-        ctx.sequencer_client()
-            .get_account(recipient_vault_id)
-            .await?,
+        get_account(&ctx, recipient_vault_id).await?,
         false,
         recipient_vault_id,
     );
@@ -169,27 +214,15 @@ async fn private_bridge_deposit_invocation_is_dropped() -> anyhow::Result<()> {
         witness_set,
     ));
 
-    let bridge_balance_before = ctx
-        .sequencer_client()
-        .get_account_balance(bridge_account_id)
-        .await?;
-    let vault_balance_before = ctx
-        .sequencer_client()
-        .get_account_balance(recipient_vault_id)
-        .await?;
+    let bridge_balance_before = account_balance(&ctx, bridge_account_id).await?;
+    let vault_balance_before = account_balance(&ctx, recipient_vault_id).await?;
 
     let tx_hash = ctx.sequencer_client().send_transaction(attack_tx).await?;
 
     tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await;
 
-    let bridge_balance_after = ctx
-        .sequencer_client()
-        .get_account_balance(bridge_account_id)
-        .await?;
-    let vault_balance_after = ctx
-        .sequencer_client()
-        .get_account_balance(recipient_vault_id)
-        .await?;
+    let bridge_balance_after = account_balance(&ctx, bridge_account_id).await?;
+    let vault_balance_after = account_balance(&ctx, recipient_vault_id).await?;
     let tx_on_chain = ctx.sequencer_client().get_transaction(tx_hash).await?;
 
     assert_eq!(bridge_balance_after, bridge_balance_before);
@@ -361,7 +394,7 @@ async fn wait_for_vault_balance(
         + Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS);
     tokio::time::timeout(timeout, async {
         loop {
-            let balance = ctx.sequencer_client().get_account_balance(vault_id).await?;
+            let balance = account_balance(ctx, vault_id).await?;
             if balance == expected_balance {
                 return Ok(());
             }
@@ -389,14 +422,8 @@ async fn bedrock_deposit_claim_and_withdraw_round_trip_succeeds() -> anyhow::Res
     let vault_program_id = programs::vault().id();
     let recipient_vault_id = vault_core::compute_vault_account_id(vault_program_id, recipient_id);
 
-    let vault_balance_before = ctx
-        .sequencer_client()
-        .get_account_balance(recipient_vault_id)
-        .await?;
-    let recipient_balance_before = ctx
-        .sequencer_client()
-        .get_account_balance(recipient_id)
-        .await?;
+    let vault_balance_before = account_balance(&ctx, recipient_vault_id).await?;
+    let recipient_balance_before = account_balance(&ctx, recipient_id).await?;
 
     // Submit deposit to Bedrock
     submit_bedrock_deposit(ctx.bedrock_addr(), bedrock_account_pk, recipient_id, amount)
@@ -447,14 +474,8 @@ async fn bedrock_deposit_claim_and_withdraw_round_trip_succeeds() -> anyhow::Res
     tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await;
 
     let claim_on_chain = ctx.sequencer_client().get_transaction(claim_hash).await?;
-    let vault_balance_after_claim = ctx
-        .sequencer_client()
-        .get_account_balance(recipient_vault_id)
-        .await?;
-    let recipient_balance_after_claim = ctx
-        .sequencer_client()
-        .get_account_balance(recipient_id)
-        .await?;
+    let vault_balance_after_claim = account_balance(&ctx, recipient_vault_id).await?;
+    let recipient_balance_after_claim = account_balance(&ctx, recipient_id).await?;
 
     assert!(
         claim_on_chain.is_some(),
@@ -483,7 +504,7 @@ async fn bedrock_deposit_claim_and_withdraw_round_trip_succeeds() -> anyhow::Res
             account_id.into(),
         )
         .await?;
-        let sequencer_account = ctx.sequencer_client().get_account(account_id).await?;
+        let sequencer_account = get_account(&ctx, account_id).await?;
         assert_eq!(
             indexer_account,
             sequencer_account.into(),

integration_tests/tests/indexer_state_consistency.rs

@@ -1,51 +1,36 @@
 #![expect(
-    clippy::shadow_unrelated,
     clippy::tests_outside_test_module,
     reason = "We don't care about these in tests"
 )]
 
 use std::time::Duration;
 
-use anyhow::{Context as _, Result};
+use anyhow::Result;
 use indexer_service_rpc::RpcClient as _;
 use integration_tests::{
-    TIME_TO_WAIT_FOR_BLOCK_SECONDS, TestContext, private_mention, public_mention,
-    verify_commitment_is_in_state, wait_for_indexer_to_catch_up,
+    TIME_TO_WAIT_FOR_BLOCK_SECONDS, TestContext, account_balance,
+    assert_private_commitment_in_state, get_account, private_mention, public_mention, send,
+    wait_for_indexer_to_catch_up,
 };
 use lee::AccountId;
 use log::info;
-use wallet::cli::{Command, programs::native_token_transfer::AuthTransferSubcommand};
 
 #[tokio::test]
 async fn indexer_state_consistency() -> Result<()> {
     let mut ctx = TestContext::new().await?;
 
-    let command = Command::AuthTransfer(AuthTransferSubcommand::Send {
-        from: public_mention(ctx.existing_public_accounts()[0]),
-        to: Some(public_mention(ctx.existing_public_accounts()[1])),
-        to_npk: None,
-        to_vpk: None,
-        to_keys: None,
-        to_identifier: Some(0),
-        amount: 100,
-    });
-
-    wallet::cli::execute_subcommand(ctx.wallet_mut(), command).await?;
+    let (acc0, acc1) = (
+        ctx.existing_public_accounts()[0],
+        ctx.existing_public_accounts()[1],
+    );
+    send(&mut ctx, public_mention(acc0), public_mention(acc1), 100).await?;
 
     info!("Waiting for next block creation");
     tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await;
 
     info!("Checking correct balance move");
-    let acc_1_balance = sequencer_service_rpc::RpcClient::get_account_balance(
-        ctx.sequencer_client(),
-        ctx.existing_public_accounts()[0],
-    )
-    .await?;
-    let acc_2_balance = sequencer_service_rpc::RpcClient::get_account_balance(
-        ctx.sequencer_client(),
-        ctx.existing_public_accounts()[1],
-    )
-    .await?;
+    let acc_1_balance = account_balance(&ctx, ctx.existing_public_accounts()[0]).await?;
+    let acc_2_balance = account_balance(&ctx, ctx.existing_public_accounts()[1]).await?;
 
     info!("Balance of sender: {acc_1_balance:#?}");
     info!("Balance of receiver: {acc_2_balance:#?}");
@@ -56,32 +41,13 @@ async fn indexer_state_consistency() -> Result<()> {
     let from: AccountId = ctx.existing_private_accounts()[0];
     let to: AccountId = ctx.existing_private_accounts()[1];
 
-    let command = Command::AuthTransfer(AuthTransferSubcommand::Send {
-        from: private_mention(from),
-        to: Some(private_mention(to)),
-        to_npk: None,
-        to_vpk: None,
-        to_keys: None,
-        to_identifier: Some(0),
-        amount: 100,
-    });
-
-    wallet::cli::execute_subcommand(ctx.wallet_mut(), command).await?;
+    send(&mut ctx, private_mention(from), private_mention(to), 100).await?;
 
     info!("Waiting for next block creation");
     tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await;
 
-    let new_commitment1 = ctx
-        .wallet()
-        .get_private_account_commitment(from)
-        .context("Failed to get private account commitment for sender")?;
-    assert!(verify_commitment_is_in_state(new_commitment1, ctx.sequencer_client()).await);
-
-    let new_commitment2 = ctx
-        .wallet()
-        .get_private_account_commitment(to)
-        .context("Failed to get private account commitment for receiver")?;
-    assert!(verify_commitment_is_in_state(new_commitment2, ctx.sequencer_client()).await);
+    assert_private_commitment_in_state(&ctx, from, "sender").await?;
+    assert_private_commitment_in_state(&ctx, to, "receiver").await?;
 
     info!("Successfully transferred privately to owned account");
 
@@ -100,16 +66,8 @@ async fn indexer_state_consistency() -> Result<()> {
         .unwrap();
 
     info!("Checking correct state transition");
-    let acc1_seq_state = sequencer_service_rpc::RpcClient::get_account(
-        ctx.sequencer_client(),
-        ctx.existing_public_accounts()[0],
-    )
-    .await?;
-    let acc2_seq_state = sequencer_service_rpc::RpcClient::get_account(
-        ctx.sequencer_client(),
-        ctx.existing_public_accounts()[1],
-    )
-    .await?;
+    let acc1_seq_state = get_account(&ctx, ctx.existing_public_accounts()[0]).await?;
+    let acc2_seq_state = get_account(&ctx, ctx.existing_public_accounts()[1]).await?;
 
     assert_eq!(acc1_ind_state, acc1_seq_state.into());
     assert_eq!(acc2_ind_state, acc2_seq_state.into());

integration_tests/tests/indexer_state_consistency_with_labels.rs

@@ -9,12 +9,13 @@ use std::time::Duration;
 use anyhow::Result;
 use indexer_service_rpc::RpcClient as _;
 use integration_tests::{
-    TIME_TO_WAIT_FOR_BLOCK_SECONDS, TestContext, public_mention, wait_for_indexer_to_catch_up,
+    TIME_TO_WAIT_FOR_BLOCK_SECONDS, TestContext, account_balance, get_account, public_mention,
+    send, wait_for_indexer_to_catch_up,
 };
 use log::info;
 use wallet::{
     account::Label,
-    cli::{CliAccountMention, Command, programs::native_token_transfer::AuthTransferSubcommand},
+    cli::{CliAccountMention, Command},
 };
 
 #[tokio::test]
@@ -38,31 +39,19 @@ async fn indexer_state_consistency_with_labels() -> Result<()> {
     wallet::cli::execute_subcommand(ctx.wallet_mut(), label_cmd).await?;
 
     // Send using labels instead of account IDs
-    let command = Command::AuthTransfer(AuthTransferSubcommand::Send {
-        from: CliAccountMention::Label(from_label),
-        to: Some(CliAccountMention::Label(to_label)),
-        to_npk: None,
-        to_vpk: None,
-        to_keys: None,
-        to_identifier: Some(0),
-        amount: 100,
-    });
-
-    wallet::cli::execute_subcommand(ctx.wallet_mut(), command).await?;
+    send(
+        &mut ctx,
+        CliAccountMention::Label(from_label),
+        CliAccountMention::Label(to_label),
+        100,
+    )
+    .await?;
 
     info!("Waiting for next block creation");
     tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await;
 
-    let acc_1_balance = sequencer_service_rpc::RpcClient::get_account_balance(
-        ctx.sequencer_client(),
-        ctx.existing_public_accounts()[0],
-    )
-    .await?;
-    let acc_2_balance = sequencer_service_rpc::RpcClient::get_account_balance(
-        ctx.sequencer_client(),
-        ctx.existing_public_accounts()[1],
-    )
-    .await?;
+    let acc_1_balance = account_balance(&ctx, ctx.existing_public_accounts()[0]).await?;
+    let acc_2_balance = account_balance(&ctx, ctx.existing_public_accounts()[1]).await?;
 
     assert_eq!(acc_1_balance, 9900);
     assert_eq!(acc_2_balance, 20100);
@@ -75,11 +64,7 @@ async fn indexer_state_consistency_with_labels() -> Result<()> {
         .get_account(ctx.existing_public_accounts()[0].into())
         .await
         .unwrap();
-    let acc1_seq_state = sequencer_service_rpc::RpcClient::get_account(
-        ctx.sequencer_client(),
-        ctx.existing_public_accounts()[0],
-    )
-    .await?;
+    let acc1_seq_state = get_account(&ctx, ctx.existing_public_accounts()[0]).await?;
 
     assert_eq!(acc1_ind_state, acc1_seq_state.into());
 

integration_tests/tests/keys.rs

@@ -8,8 +8,9 @@ use std::{str::FromStr as _, time::Duration};
 
 use anyhow::{Context as _, Result};
 use integration_tests::{
-    TIME_TO_WAIT_FOR_BLOCK_SECONDS, TestContext, fetch_privacy_preserving_tx, private_mention,
-    public_mention, verify_commitment_is_in_state,
+    TIME_TO_WAIT_FOR_BLOCK_SECONDS, TestContext, assert_public_account_restored,
+    fetch_privacy_preserving_tx, new_account, private_mention, public_mention,
+    restored_private_account, send, verify_commitment_is_in_state,
 };
 use key_protocol::key_management::key_tree::chain_index::ChainIndex;
 use lee::AccountId;
@@ -17,8 +18,7 @@ use log::info;
 use sequencer_service_rpc::RpcClient as _;
 use tokio::test;
 use wallet::cli::{
-    Command, SubcommandReturnValue,
-    account::{AccountSubcommand, NewSubcommand},
+    Command, SubcommandReturnValue, account::AccountSubcommand,
     programs::native_token_transfer::AuthTransferSubcommand,
 };
 
@@ -28,35 +28,12 @@ async fn sync_private_account_with_non_zero_chain_index() -> Result<()> {
 
     let from: AccountId = ctx.existing_private_accounts()[0];
 
-    // Create a new private account
-    let command = Command::Account(AccountSubcommand::New(NewSubcommand::Private {
-        cci: None,
-        label: None,
-    }));
-
+    // Key Tree shift — create 3 accounts to advance the key index
     for _ in 0..3 {
-        // Key Tree shift
-        // This way we have account with child index > 0.
-        let result = wallet::cli::execute_subcommand(
-            ctx.wallet_mut(),
-            Command::Account(AccountSubcommand::New(NewSubcommand::Private {
-                cci: None,
-                label: None,
-            })),
-        )
-        .await?;
-        let SubcommandReturnValue::RegisterAccount { account_id: _ } = result else {
-            anyhow::bail!("Expected RegisterAccount return value");
-        };
+        new_account(&mut ctx, true, None).await?;
     }
 
-    let sub_ret = wallet::cli::execute_subcommand(ctx.wallet_mut(), command).await?;
-    let SubcommandReturnValue::RegisterAccount {
-        account_id: to_account_id,
-    } = sub_ret
-    else {
-        anyhow::bail!("Expected RegisterAccount return value");
-    };
+    let to_account_id = new_account(&mut ctx, true, None).await?;
 
     // Get the keys for the newly created account
     let to_account = ctx
@@ -118,107 +95,47 @@ async fn restore_keys_from_seed() -> Result<()> {
 
     let from: AccountId = ctx.existing_private_accounts()[0];
 
-    // Create first private account at root
-    let command = Command::Account(AccountSubcommand::New(NewSubcommand::Private {
-        cci: Some(ChainIndex::root()),
-        label: None,
-    }));
-    let result = wallet::cli::execute_subcommand(ctx.wallet_mut(), command).await?;
-    let SubcommandReturnValue::RegisterAccount {
-        account_id: to_account_id1,
-    } = result
-    else {
-        anyhow::bail!("Expected RegisterAccount return value");
-    };
+    // Create private accounts at root and /0
+    let to_account_id1 = new_account(&mut ctx, true, Some(ChainIndex::root())).await?;
+    let to_account_id2 = new_account(&mut ctx, true, Some(ChainIndex::from_str("/0")?)).await?;
 
-    // Create second private account at /0
-    let command = Command::Account(AccountSubcommand::New(NewSubcommand::Private {
-        cci: Some(ChainIndex::from_str("/0")?),
-        label: None,
-    }));
-    let result = wallet::cli::execute_subcommand(ctx.wallet_mut(), command).await?;
-    let SubcommandReturnValue::RegisterAccount {
-        account_id: to_account_id2,
-    } = result
-    else {
-        anyhow::bail!("Expected RegisterAccount return value");
-    };
-
-    // Send to first private account
-    let command = Command::AuthTransfer(AuthTransferSubcommand::Send {
-        from: private_mention(from),
-        to: Some(private_mention(to_account_id1)),
-        to_npk: None,
-        to_vpk: None,
-        to_keys: None,
-        to_identifier: Some(0),
-        amount: 100,
-    });
-    wallet::cli::execute_subcommand(ctx.wallet_mut(), command).await?;
-
-    // Send to second private account
-    let command = Command::AuthTransfer(AuthTransferSubcommand::Send {
-        from: private_mention(from),
-        to: Some(private_mention(to_account_id2)),
-        to_npk: None,
-        to_vpk: None,
-        to_keys: None,
-        to_identifier: Some(0),
-        amount: 101,
-    });
-    wallet::cli::execute_subcommand(ctx.wallet_mut(), command).await?;
+    // Send to both private accounts
+    send(
+        &mut ctx,
+        private_mention(from),
+        private_mention(to_account_id1),
+        100,
+    )
+    .await?;
+    send(
+        &mut ctx,
+        private_mention(from),
+        private_mention(to_account_id2),
+        101,
+    )
+    .await?;
 
     let from: AccountId = ctx.existing_public_accounts()[0];
 
-    // Create first public account at root
-    let command = Command::Account(AccountSubcommand::New(NewSubcommand::Public {
-        cci: Some(ChainIndex::root()),
-        label: None,
-    }));
-    let result = wallet::cli::execute_subcommand(ctx.wallet_mut(), command).await?;
-    let SubcommandReturnValue::RegisterAccount {
-        account_id: to_account_id3,
-    } = result
-    else {
-        anyhow::bail!("Expected RegisterAccount return value");
-    };
+    // Create public accounts at root and /0
+    let to_account_id3 = new_account(&mut ctx, false, Some(ChainIndex::root())).await?;
+    let to_account_id4 = new_account(&mut ctx, false, Some(ChainIndex::from_str("/0")?)).await?;
 
-    // Create second public account at /0
-    let command = Command::Account(AccountSubcommand::New(NewSubcommand::Public {
-        cci: Some(ChainIndex::from_str("/0")?),
-        label: None,
-    }));
-    let result = wallet::cli::execute_subcommand(ctx.wallet_mut(), command).await?;
-    let SubcommandReturnValue::RegisterAccount {
-        account_id: to_account_id4,
-    } = result
-    else {
-        anyhow::bail!("Expected RegisterAccount return value");
-    };
-
-    // Send to first public account
-    let command = Command::AuthTransfer(AuthTransferSubcommand::Send {
-        from: public_mention(from),
-        to: Some(public_mention(to_account_id3)),
-        to_npk: None,
-        to_vpk: None,
-        to_keys: None,
-        to_identifier: Some(0),
-        amount: 102,
-    });
-    wallet::cli::execute_subcommand(ctx.wallet_mut(), command).await?;
-
-    // Send to second public account
-    let command = Command::AuthTransfer(AuthTransferSubcommand::Send {
-        from: public_mention(from),
-        to: Some(public_mention(to_account_id4)),
-        to_npk: None,
-        to_vpk: None,
-        to_keys: None,
-        to_identifier: Some(0),
-        amount: 103,
-    });
-    wallet::cli::execute_subcommand(ctx.wallet_mut(), command).await?;
+    // Send to both public accounts
+    send(
+        &mut ctx,
+        public_mention(from),
+        public_mention(to_account_id3),
+        102,
+    )
+    .await?;
+    send(
+        &mut ctx,
+        public_mention(from),
+        public_mention(to_account_id4),
+        103,
+    )
+    .await?;
 
     info!("Preparation complete, performing keys restoration");
 
@@ -226,34 +143,12 @@ async fn restore_keys_from_seed() -> Result<()> {
     wallet::cli::execute_keys_restoration(ctx.wallet_mut(), 10).await?;
 
     // Verify restored private accounts
-    let acc1 = ctx
-        .wallet()
-        .storage()
-        .key_chain()
-        .private_account(to_account_id1)
-        .expect("Acc 1 should be restored");
-
-    let acc2 = ctx
-        .wallet()
-        .storage()
-        .key_chain()
-        .private_account(to_account_id2)
-        .expect("Acc 2 should be restored");
+    let acc1 = restored_private_account(&ctx, to_account_id1, "Acc 1");
+    let acc2 = restored_private_account(&ctx, to_account_id2, "Acc 2");
 
     // Verify restored public accounts
-    let _acc3 = ctx
-        .wallet()
-        .storage()
-        .key_chain()
-        .pub_account_signing_key(to_account_id3)
-        .expect("Acc 3 should be restored");
-
-    let _acc4 = ctx
-        .wallet()
-        .storage()
-        .key_chain()
-        .pub_account_signing_key(to_account_id4)
-        .expect("Acc 4 should be restored");
+    assert_public_account_restored(&ctx, to_account_id3, "Acc 3");
+    assert_public_account_restored(&ctx, to_account_id4, "Acc 4");
 
     assert_eq!(
         acc1.account.program_owner,
@@ -270,27 +165,20 @@ async fn restore_keys_from_seed() -> Result<()> {
     info!("Tree checks passed, testing restored accounts can transact");
 
     // Test that restored accounts can send transactions
-    let command = Command::AuthTransfer(AuthTransferSubcommand::Send {
-        from: private_mention(to_account_id1),
-        to: Some(private_mention(to_account_id2)),
-        to_npk: None,
-        to_vpk: None,
-        to_keys: None,
-        to_identifier: Some(0),
-        amount: 10,
-    });
-    wallet::cli::execute_subcommand(ctx.wallet_mut(), command).await?;
-
-    let command = Command::AuthTransfer(AuthTransferSubcommand::Send {
-        from: public_mention(to_account_id3),
-        to: Some(public_mention(to_account_id4)),
-        to_npk: None,
-        to_vpk: None,
-        to_keys: None,
-        to_identifier: Some(0),
-        amount: 11,
-    });
-    wallet::cli::execute_subcommand(ctx.wallet_mut(), command).await?;
+    send(
+        &mut ctx,
+        private_mention(to_account_id1),
+        private_mention(to_account_id2),
+        10,
+    )
+    .await?;
+    send(
+        &mut ctx,
+        public_mention(to_account_id3),
+        public_mention(to_account_id4),
+        11,
+    )
+    .await?;
 
     tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await;
 

integration_tests/tests/pinata.rs

@@ -8,15 +8,13 @@ use std::time::Duration;
 
 use anyhow::{Context as _, Result};
 use integration_tests::{
-    TIME_TO_WAIT_FOR_BLOCK_SECONDS, TestContext, private_mention, public_mention,
-    verify_commitment_is_in_state,
+    TIME_TO_WAIT_FOR_BLOCK_SECONDS, TestContext, account_balance, new_account, private_mention,
+    public_mention, sync_private, verify_commitment_is_in_state,
 };
 use log::info;
-use sequencer_service_rpc::RpcClient as _;
 use tokio::test;
 use wallet::cli::{
     Command, SubcommandReturnValue,
-    account::{AccountSubcommand, NewSubcommand},
     programs::{
         native_token_transfer::AuthTransferSubcommand, pinata::PinataProgramAgnosticSubcommand,
     },
@@ -26,25 +24,9 @@ use wallet::cli::{
 async fn claim_pinata_to_uninitialized_public_account_fails_fast() -> Result<()> {
     let mut ctx = TestContext::new().await?;
 
-    let result = wallet::cli::execute_subcommand(
-        ctx.wallet_mut(),
-        Command::Account(AccountSubcommand::New(NewSubcommand::Public {
-            cci: None,
-            label: None,
-        })),
-    )
-    .await?;
-    let SubcommandReturnValue::RegisterAccount {
-        account_id: winner_account_id,
-    } = result
-    else {
-        anyhow::bail!("Expected RegisterAccount return value");
-    };
+    let winner_account_id = new_account(&mut ctx, false, None).await?;
 
-    let pinata_balance_pre = ctx
-        .sequencer_client()
-        .get_account_balance(system_accounts::pinata_account_id())
-        .await?;
+    let pinata_balance_pre = account_balance(&ctx, system_accounts::pinata_account_id()).await?;
 
     let claim_result = wallet::cli::execute_subcommand(
         ctx.wallet_mut(),
@@ -64,10 +46,7 @@ async fn claim_pinata_to_uninitialized_public_account_fails_fast() -> Result<()>
         "Expected init guidance, got: {err}",
     );
 
-    let pinata_balance_post = ctx
-        .sequencer_client()
-        .get_account_balance(system_accounts::pinata_account_id())
-        .await?;
+    let pinata_balance_post = account_balance(&ctx, system_accounts::pinata_account_id()).await?;
 
     assert_eq!(pinata_balance_post, pinata_balance_pre);
 
@@ -78,25 +57,9 @@ async fn claim_pinata_to_uninitialized_public_account_fails_fast() -> Result<()>
 async fn claim_pinata_to_uninitialized_private_account_fails_fast() -> Result<()> {
     let mut ctx = TestContext::new().await?;
 
-    let result = wallet::cli::execute_subcommand(
-        ctx.wallet_mut(),
-        Command::Account(AccountSubcommand::New(NewSubcommand::Private {
-            cci: None,
-            label: None,
-        })),
-    )
-    .await?;
-    let SubcommandReturnValue::RegisterAccount {
-        account_id: winner_account_id,
-    } = result
-    else {
-        anyhow::bail!("Expected RegisterAccount return value");
-    };
+    let winner_account_id = new_account(&mut ctx, true, None).await?;
 
-    let pinata_balance_pre = ctx
-        .sequencer_client()
-        .get_account_balance(system_accounts::pinata_account_id())
-        .await?;
+    let pinata_balance_pre = account_balance(&ctx, system_accounts::pinata_account_id()).await?;
 
     let claim_result = wallet::cli::execute_subcommand(
         ctx.wallet_mut(),
@@ -116,10 +79,7 @@ async fn claim_pinata_to_uninitialized_private_account_fails_fast() -> Result<()
         "Expected init guidance, got: {err}",
     );
 
-    let pinata_balance_post = ctx
-        .sequencer_client()
-        .get_account_balance(system_accounts::pinata_account_id())
-        .await?;
+    let pinata_balance_post = account_balance(&ctx, system_accounts::pinata_account_id()).await?;
 
     assert_eq!(pinata_balance_post, pinata_balance_pre);
 
@@ -135,10 +95,7 @@ async fn claim_pinata_to_existing_public_account() -> Result<()> {
         to: public_mention(ctx.existing_public_accounts()[0]),
     });
 
-    let pinata_balance_pre = ctx
-        .sequencer_client()
-        .get_account_balance(system_accounts::pinata_account_id())
-        .await?;
+    let pinata_balance_pre = account_balance(&ctx, system_accounts::pinata_account_id()).await?;
 
     wallet::cli::execute_subcommand(ctx.wallet_mut(), command).await?;
 
@@ -146,15 +103,9 @@ async fn claim_pinata_to_existing_public_account() -> Result<()> {
     tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await;
 
     info!("Checking correct balance move");
-    let pinata_balance_post = ctx
-        .sequencer_client()
-        .get_account_balance(system_accounts::pinata_account_id())
-        .await?;
+    let pinata_balance_post = account_balance(&ctx, system_accounts::pinata_account_id()).await?;
 
-    let winner_balance_post = ctx
-        .sequencer_client()
-        .get_account_balance(ctx.existing_public_accounts()[0])
-        .await?;
+    let winner_balance_post = account_balance(&ctx, ctx.existing_public_accounts()[0]).await?;
 
     assert_eq!(pinata_balance_post, pinata_balance_pre - pinata_prize);
     assert_eq!(winner_balance_post, 10000 + pinata_prize);
@@ -173,10 +124,7 @@ async fn claim_pinata_to_existing_private_account() -> Result<()> {
         to: private_mention(ctx.existing_private_accounts()[0]),
     });
 
-    let pinata_balance_pre = ctx
-        .sequencer_client()
-        .get_account_balance(system_accounts::pinata_account_id())
-        .await?;
+    let pinata_balance_pre = account_balance(&ctx, system_accounts::pinata_account_id()).await?;
 
     let result = wallet::cli::execute_subcommand(ctx.wallet_mut(), command).await?;
     let SubcommandReturnValue::PrivacyPreservingTransfer { tx_hash: _ } = result else {
@@ -187,8 +135,7 @@ async fn claim_pinata_to_existing_private_account() -> Result<()> {
     tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await;
 
     info!("Syncing private accounts");
-    let command = Command::Account(AccountSubcommand::SyncPrivate {});
-    wallet::cli::execute_subcommand(ctx.wallet_mut(), command).await?;
+    sync_private(&mut ctx).await?;
 
     let new_commitment = ctx
         .wallet()
@@ -196,10 +143,7 @@ async fn claim_pinata_to_existing_private_account() -> Result<()> {
         .context("Failed to get private account commitment")?;
     assert!(verify_commitment_is_in_state(new_commitment, ctx.sequencer_client()).await);
 
-    let pinata_balance_post = ctx
-        .sequencer_client()
-        .get_account_balance(system_accounts::pinata_account_id())
-        .await?;
+    let pinata_balance_post = account_balance(&ctx, system_accounts::pinata_account_id()).await?;
 
     assert_eq!(pinata_balance_post, pinata_balance_pre - pinata_prize);
 
@@ -215,20 +159,7 @@ async fn claim_pinata_to_new_private_account() -> Result<()> {
     let pinata_prize = 150;
 
     // Create new private account
-    let result = wallet::cli::execute_subcommand(
-        ctx.wallet_mut(),
-        Command::Account(AccountSubcommand::New(NewSubcommand::Private {
-            cci: None,
-            label: None,
-        })),
-    )
-    .await?;
-    let SubcommandReturnValue::RegisterAccount {
-        account_id: winner_account_id,
-    } = result
-    else {
-        anyhow::bail!("Expected RegisterAccount return value");
-    };
+    let winner_account_id = new_account(&mut ctx, true, None).await?;
 
     // Initialize account under auth transfer program
     let command = Command::AuthTransfer(AuthTransferSubcommand::Init {
@@ -250,10 +181,7 @@ async fn claim_pinata_to_new_private_account() -> Result<()> {
         to: private_mention(winner_account_id),
     });
 
-    let pinata_balance_pre = ctx
-        .sequencer_client()
-        .get_account_balance(system_accounts::pinata_account_id())
-        .await?;
+    let pinata_balance_pre = account_balance(&ctx, system_accounts::pinata_account_id()).await?;
 
     wallet::cli::execute_subcommand(ctx.wallet_mut(), command).await?;
 
@@ -266,10 +194,7 @@ async fn claim_pinata_to_new_private_account() -> Result<()> {
         .context("Failed to get private account commitment")?;
     assert!(verify_commitment_is_in_state(new_commitment, ctx.sequencer_client()).await);
 
-    let pinata_balance_post = ctx
-        .sequencer_client()
-        .get_account_balance(system_accounts::pinata_account_id())
-        .await?;
+    let pinata_balance_post = account_balance(&ctx, system_accounts::pinata_account_id()).await?;
 
     assert_eq!(pinata_balance_post, pinata_balance_pre - pinata_prize);
 

integration_tests/tests/private_pda.rs

@@ -9,7 +9,7 @@ use anyhow::{Context as _, Result};
 use authenticated_transfer_core::Instruction as AuthTransferInstruction;
 use common::transaction::LeeTransaction;
 use integration_tests::{
-    TIME_TO_WAIT_FOR_BLOCK_SECONDS, TestContext, verify_commitment_is_in_state,
+    TIME_TO_WAIT_FOR_BLOCK_SECONDS, TestContext, sync_private, verify_commitment_is_in_state,
 };
 use key_protocol::key_management::ephemeral_key_holder::EphemeralKeyHolder;
 use lee::{
@@ -22,7 +22,7 @@ use lee::{
     program::Program,
 };
 use lee_core::{
-    EncryptedAccountData, InputAccountIdentity, NullifierPublicKey,
+    DUMMY_COMMITMENT_HASH, EncryptedAccountData, InputAccountIdentity, NullifierPublicKey,
     account::{Account, AccountWithMetadata},
     encryption::ViewingPublicKey,
     program::PdaSeed,
@@ -30,10 +30,7 @@ use lee_core::{
 use log::info;
 use sequencer_service_rpc::RpcClient as _;
 use tokio::test;
-use wallet::{
-    AccountIdentity, WalletCore,
-    cli::{Command, account::AccountSubcommand},
-};
+use wallet::{AccountIdentity, WalletCore};
 
 /// Funds a private PDA by calling `auth_transfer` directly.
 #[expect(
@@ -78,6 +75,7 @@ async fn fund_private_pda(
             npk,
             ssk,
             identifier,
+            commitment_root: DUMMY_COMMITMENT_HASH,
             seed: Some((seed, authority_program_id)),
         },
     ];
@@ -217,11 +215,7 @@ async fn private_pda_family_members_receive_and_spend() -> Result<()> {
     tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await;
 
     // Sync so alice's wallet discovers and stores both PDAs.
-    wallet::cli::execute_subcommand(
-        ctx.wallet_mut(),
-        Command::Account(AccountSubcommand::SyncPrivate {}),
-    )
-    .await?;
+    sync_private(&mut ctx).await?;
 
     // Both PDAs must be discoverable and have the correct balance.
     let pda_0_account = ctx
@@ -300,11 +294,7 @@ async fn private_pda_family_members_receive_and_spend() -> Result<()> {
     info!("Waiting for block");
     tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await;
 
-    wallet::cli::execute_subcommand(
-        ctx.wallet_mut(),
-        Command::Account(AccountSubcommand::SyncPrivate {}),
-    )
-    .await?;
+    sync_private(&mut ctx).await?;
 
     // After spending, PDAs should have the remaining balance.
     let pda_0_spent = ctx

integration_tests/tests/program_deployment.rs

@@ -7,14 +7,11 @@ use std::{io::Write as _, time::Duration};
 
 use anyhow::Result;
 use common::transaction::LeeTransaction;
-use integration_tests::{TIME_TO_WAIT_FOR_BLOCK_SECONDS, TestContext};
+use integration_tests::{TIME_TO_WAIT_FOR_BLOCK_SECONDS, TestContext, get_account, new_account};
 use log::info;
 use sequencer_service_rpc::RpcClient as _;
 use tokio::test;
-use wallet::cli::{
-    Command, SubcommandReturnValue,
-    account::{AccountSubcommand, NewSubcommand},
-};
+use wallet::cli::Command;
 
 #[test]
 async fn deploy_and_execute_program() -> Result<()> {
@@ -35,17 +32,7 @@ async fn deploy_and_execute_program() -> Result<()> {
     info!("Waiting for next block creation");
     tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await;
 
-    let SubcommandReturnValue::RegisterAccount { account_id } = wallet::cli::execute_subcommand(
-        ctx.wallet_mut(),
-        Command::Account(AccountSubcommand::New(NewSubcommand::Public {
-            cci: None,
-            label: None,
-        })),
-    )
-    .await?
-    else {
-        panic!("Expected RegisterAccount return value");
-    };
+    let account_id = new_account(&mut ctx, false, None).await?;
 
     let nonces = ctx.wallet().get_accounts_nonces(vec![account_id]).await?;
     let private_key = ctx
@@ -66,7 +53,7 @@ async fn deploy_and_execute_program() -> Result<()> {
     // block
     tokio::time::sleep(Duration::from_secs(2 * TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await;
 
-    let post_state_account = ctx.sequencer_client().get_account(account_id).await?;
+    let post_state_account = get_account(&ctx, account_id).await?;
 
     let expected_data: &[u8] = &[];
     assert_eq!(post_state_account.program_owner, claimer.id());

integration_tests/tests/shared_accounts.rs

@@ -19,7 +19,7 @@ use std::time::Duration;
 
 use anyhow::{Context as _, Result};
 use integration_tests::{
-    TIME_TO_WAIT_FOR_BLOCK_SECONDS, TestContext, private_mention, public_mention,
+    TIME_TO_WAIT_FOR_BLOCK_SECONDS, TestContext, private_mention, public_mention, sync_private,
 };
 use log::info;
 use tokio::test;
@@ -197,8 +197,7 @@ async fn fund_shared_account_from_public() -> Result<()> {
     tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await;
 
     // Sync private accounts
-    let command = Command::Account(AccountSubcommand::SyncPrivate);
-    wallet::cli::execute_subcommand(ctx.wallet_mut(), command).await?;
+    sync_private(&mut ctx).await?;
 
     // Fund from a public account
     let from_public = ctx.existing_public_accounts()[0];
@@ -216,8 +215,7 @@ async fn fund_shared_account_from_public() -> Result<()> {
     tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await;
 
     // Sync private accounts
-    let command = Command::Account(AccountSubcommand::SyncPrivate);
-    wallet::cli::execute_subcommand(ctx.wallet_mut(), command).await?;
+    sync_private(&mut ctx).await?;
 
     // Verify the shared account was updated
     let entry = ctx

integration_tests/tests/token.rs

@@ -8,12 +8,11 @@ use std::time::Duration;
 
 use anyhow::{Context as _, Result};
 use integration_tests::{
-    TIME_TO_WAIT_FOR_BLOCK_SECONDS, TestContext, private_mention, public_mention,
-    verify_commitment_is_in_state,
+    TIME_TO_WAIT_FOR_BLOCK_SECONDS, TestContext, get_account, new_account, private_mention,
+    public_mention, sync_private, verify_commitment_is_in_state,
 };
 use key_protocol::key_management::key_tree::chain_index::ChainIndex;
 use log::info;
-use sequencer_service_rpc::RpcClient as _;
 use token_core::{TokenDefinition, TokenHolding};
 use tokio::test;
 use wallet::{
@@ -30,52 +29,13 @@ async fn create_and_transfer_public_token() -> Result<()> {
     let mut ctx = TestContext::new().await?;
 
     // Create new account for the token definition
-    let result = wallet::cli::execute_subcommand(
-        ctx.wallet_mut(),
-        Command::Account(AccountSubcommand::New(NewSubcommand::Public {
-            cci: None,
-            label: None,
-        })),
-    )
-    .await?;
-    let SubcommandReturnValue::RegisterAccount {
-        account_id: definition_account_id,
-    } = result
-    else {
-        anyhow::bail!("Expected RegisterAccount return value");
-    };
+    let definition_account_id = new_account(&mut ctx, false, None).await?;
 
     // Create new account for the token supply holder
-    let result = wallet::cli::execute_subcommand(
-        ctx.wallet_mut(),
-        Command::Account(AccountSubcommand::New(NewSubcommand::Public {
-            cci: None,
-            label: None,
-        })),
-    )
-    .await?;
-    let SubcommandReturnValue::RegisterAccount {
-        account_id: supply_account_id,
-    } = result
-    else {
-        anyhow::bail!("Expected RegisterAccount return value");
-    };
+    let supply_account_id = new_account(&mut ctx, false, None).await?;
 
     // Create new account for receiving a token transaction
-    let result = wallet::cli::execute_subcommand(
-        ctx.wallet_mut(),
-        Command::Account(AccountSubcommand::New(NewSubcommand::Public {
-            cci: None,
-            label: None,
-        })),
-    )
-    .await?;
-    let SubcommandReturnValue::RegisterAccount {
-        account_id: recipient_account_id,
-    } = result
-    else {
-        anyhow::bail!("Expected RegisterAccount return value");
-    };
+    let recipient_account_id = new_account(&mut ctx, false, None).await?;
 
     // Create new token
     let name = "A NAME".to_owned();
@@ -92,10 +52,7 @@ async fn create_and_transfer_public_token() -> Result<()> {
     tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await;
 
     // Check the status of the token definition account
-    let definition_acc = ctx
-        .sequencer_client()
-        .get_account(definition_account_id)
-        .await?;
+    let definition_acc = get_account(&ctx, definition_account_id).await?;
     let token_definition = TokenDefinition::try_from(&definition_acc.data)?;
 
     assert_eq!(definition_acc.program_owner, programs::token().id());
@@ -109,10 +66,7 @@ async fn create_and_transfer_public_token() -> Result<()> {
     );
 
     // Check the status of the token holding account with the total supply
-    let supply_acc = ctx
-        .sequencer_client()
-        .get_account(supply_account_id)
-        .await?;
+    let supply_acc = get_account(&ctx, supply_account_id).await?;
 
     // The account must be owned by the token program
     assert_eq!(supply_acc.program_owner, programs::token().id());
@@ -143,10 +97,7 @@ async fn create_and_transfer_public_token() -> Result<()> {
     tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await;
 
     // Check the status of the supply account after transfer
-    let supply_acc = ctx
-        .sequencer_client()
-        .get_account(supply_account_id)
-        .await?;
+    let supply_acc = get_account(&ctx, supply_account_id).await?;
     assert_eq!(supply_acc.program_owner, programs::token().id());
     let token_holding = TokenHolding::try_from(&supply_acc.data)?;
     assert_eq!(
@@ -158,10 +109,7 @@ async fn create_and_transfer_public_token() -> Result<()> {
     );
 
     // Check the status of the recipient account after transfer
-    let recipient_acc = ctx
-        .sequencer_client()
-        .get_account(recipient_account_id)
-        .await?;
+    let recipient_acc = get_account(&ctx, recipient_account_id).await?;
     assert_eq!(recipient_acc.program_owner, programs::token().id());
     let token_holding = TokenHolding::try_from(&recipient_acc.data)?;
     assert_eq!(
@@ -186,10 +134,7 @@ async fn create_and_transfer_public_token() -> Result<()> {
     tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await;
 
     // Check the status of the token definition account after burn
-    let definition_acc = ctx
-        .sequencer_client()
-        .get_account(definition_account_id)
-        .await?;
+    let definition_acc = get_account(&ctx, definition_account_id).await?;
     let token_definition = TokenDefinition::try_from(&definition_acc.data)?;
 
     assert_eq!(
@@ -202,10 +147,7 @@ async fn create_and_transfer_public_token() -> Result<()> {
     );
 
     // Check the status of the recipient account after burn
-    let recipient_acc = ctx
-        .sequencer_client()
-        .get_account(recipient_account_id)
-        .await?;
+    let recipient_acc = get_account(&ctx, recipient_account_id).await?;
     let token_holding = TokenHolding::try_from(&recipient_acc.data)?;
 
     assert_eq!(
@@ -234,10 +176,7 @@ async fn create_and_transfer_public_token() -> Result<()> {
     tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await;
 
     // Check the status of the token definition account after mint
-    let definition_acc = ctx
-        .sequencer_client()
-        .get_account(definition_account_id)
-        .await?;
+    let definition_acc = get_account(&ctx, definition_account_id).await?;
     let token_definition = TokenDefinition::try_from(&definition_acc.data)?;
 
     assert_eq!(
@@ -250,10 +189,7 @@ async fn create_and_transfer_public_token() -> Result<()> {
     );
 
     // Check the status of the recipient account after mint
-    let recipient_acc = ctx
-        .sequencer_client()
-        .get_account(recipient_account_id)
-        .await?;
+    let recipient_acc = get_account(&ctx, recipient_account_id).await?;
     let token_holding = TokenHolding::try_from(&recipient_acc.data)?;
 
     assert_eq!(
@@ -274,52 +210,13 @@ async fn create_and_transfer_token_with_private_supply() -> Result<()> {
     let mut ctx = TestContext::new().await?;
 
     // Create new account for the token definition (public)
-    let result = wallet::cli::execute_subcommand(
-        ctx.wallet_mut(),
-        Command::Account(AccountSubcommand::New(NewSubcommand::Public {
-            cci: None,
-            label: None,
-        })),
-    )
-    .await?;
-    let SubcommandReturnValue::RegisterAccount {
-        account_id: definition_account_id,
-    } = result
-    else {
-        anyhow::bail!("Expected RegisterAccount return value");
-    };
+    let definition_account_id = new_account(&mut ctx, false, None).await?;
 
     // Create new account for the token supply holder (private)
-    let result = wallet::cli::execute_subcommand(
-        ctx.wallet_mut(),
-        Command::Account(AccountSubcommand::New(NewSubcommand::Private {
-            cci: None,
-            label: None,
-        })),
-    )
-    .await?;
-    let SubcommandReturnValue::RegisterAccount {
-        account_id: supply_account_id,
-    } = result
-    else {
-        anyhow::bail!("Expected RegisterAccount return value");
-    };
+    let supply_account_id = new_account(&mut ctx, true, None).await?;
 
     // Create new account for receiving a token transaction (private)
-    let result = wallet::cli::execute_subcommand(
-        ctx.wallet_mut(),
-        Command::Account(AccountSubcommand::New(NewSubcommand::Private {
-            cci: None,
-            label: None,
-        })),
-    )
-    .await?;
-    let SubcommandReturnValue::RegisterAccount {
-        account_id: recipient_account_id,
-    } = result
-    else {
-        anyhow::bail!("Expected RegisterAccount return value");
-    };
+    let recipient_account_id = new_account(&mut ctx, true, None).await?;
 
     // Create new token
     let name = "A NAME".to_owned();
@@ -337,10 +234,7 @@ async fn create_and_transfer_token_with_private_supply() -> Result<()> {
     tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await;
 
     // Check the status of the token definition account
-    let definition_acc = ctx
-        .sequencer_client()
-        .get_account(definition_account_id)
-        .await?;
+    let definition_acc = get_account(&ctx, definition_account_id).await?;
     let token_definition = TokenDefinition::try_from(&definition_acc.data)?;
 
     assert_eq!(definition_acc.program_owner, programs::token().id());
@@ -402,10 +296,7 @@ async fn create_and_transfer_token_with_private_supply() -> Result<()> {
     tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await;
 
     // Check the token definition account after burn
-    let definition_acc = ctx
-        .sequencer_client()
-        .get_account(definition_account_id)
-        .await?;
+    let definition_acc = get_account(&ctx, definition_account_id).await?;
     let token_definition = TokenDefinition::try_from(&definition_acc.data)?;
 
     assert_eq!(
@@ -448,36 +339,10 @@ async fn create_token_with_private_definition() -> Result<()> {
     let mut ctx = TestContext::new().await?;
 
     // Create token definition account (private)
-    let result = wallet::cli::execute_subcommand(
-        ctx.wallet_mut(),
-        Command::Account(AccountSubcommand::New(NewSubcommand::Private {
-            cci: Some(ChainIndex::root()),
-            label: None,
-        })),
-    )
-    .await?;
-    let SubcommandReturnValue::RegisterAccount {
-        account_id: definition_account_id,
-    } = result
-    else {
-        anyhow::bail!("Expected RegisterAccount return value");
-    };
+    let definition_account_id = new_account(&mut ctx, true, Some(ChainIndex::root())).await?;
 
     // Create supply account (public)
-    let result = wallet::cli::execute_subcommand(
-        ctx.wallet_mut(),
-        Command::Account(AccountSubcommand::New(NewSubcommand::Public {
-            cci: Some(ChainIndex::root()),
-            label: None,
-        })),
-    )
-    .await?;
-    let SubcommandReturnValue::RegisterAccount {
-        account_id: supply_account_id,
-    } = result
-    else {
-        anyhow::bail!("Expected RegisterAccount return value");
-    };
+    let supply_account_id = new_account(&mut ctx, false, Some(ChainIndex::root())).await?;
 
     // Create token with private definition
     let name = "A NAME".to_owned();
@@ -502,10 +367,7 @@ async fn create_token_with_private_definition() -> Result<()> {
     assert!(verify_commitment_is_in_state(new_commitment, ctx.sequencer_client()).await);
 
     // Verify supply account
-    let supply_acc = ctx
-        .sequencer_client()
-        .get_account(supply_account_id)
-        .await?;
+    let supply_acc = get_account(&ctx, supply_account_id).await?;
 
     assert_eq!(supply_acc.program_owner, programs::token().id());
     let token_holding = TokenHolding::try_from(&supply_acc.data)?;
@@ -518,36 +380,10 @@ async fn create_token_with_private_definition() -> Result<()> {
     );
 
     // Create private recipient account
-    let result = wallet::cli::execute_subcommand(
-        ctx.wallet_mut(),
-        Command::Account(AccountSubcommand::New(NewSubcommand::Private {
-            cci: None,
-            label: None,
-        })),
-    )
-    .await?;
-    let SubcommandReturnValue::RegisterAccount {
-        account_id: recipient_account_id_private,
-    } = result
-    else {
-        anyhow::bail!("Expected RegisterAccount return value");
-    };
+    let recipient_account_id_private = new_account(&mut ctx, true, None).await?;
 
     // Create public recipient account
-    let result = wallet::cli::execute_subcommand(
-        ctx.wallet_mut(),
-        Command::Account(AccountSubcommand::New(NewSubcommand::Public {
-            cci: None,
-            label: None,
-        })),
-    )
-    .await?;
-    let SubcommandReturnValue::RegisterAccount {
-        account_id: recipient_account_id_public,
-    } = result
-    else {
-        anyhow::bail!("Expected RegisterAccount return value");
-    };
+    let recipient_account_id_public = new_account(&mut ctx, false, None).await?;
 
     // Mint to public account
     let mint_amount_public = 10;
@@ -583,10 +419,7 @@ async fn create_token_with_private_definition() -> Result<()> {
     );
 
     // Verify public recipient received tokens
-    let recipient_acc = ctx
-        .sequencer_client()
-        .get_account(recipient_account_id_public)
-        .await?;
+    let recipient_acc = get_account(&ctx, recipient_account_id_public).await?;
     let token_holding = TokenHolding::try_from(&recipient_acc.data)?;
 
     assert_eq!(
@@ -646,36 +479,10 @@ async fn create_token_with_private_definition_and_supply() -> Result<()> {
     let mut ctx = TestContext::new().await?;
 
     // Create token definition account (private)
-    let result = wallet::cli::execute_subcommand(
-        ctx.wallet_mut(),
-        Command::Account(AccountSubcommand::New(NewSubcommand::Private {
-            cci: None,
-            label: None,
-        })),
-    )
-    .await?;
-    let SubcommandReturnValue::RegisterAccount {
-        account_id: definition_account_id,
-    } = result
-    else {
-        anyhow::bail!("Expected RegisterAccount return value");
-    };
+    let definition_account_id = new_account(&mut ctx, true, None).await?;
 
     // Create supply account (private)
-    let result = wallet::cli::execute_subcommand(
-        ctx.wallet_mut(),
-        Command::Account(AccountSubcommand::New(NewSubcommand::Private {
-            cci: None,
-            label: None,
-        })),
-    )
-    .await?;
-    let SubcommandReturnValue::RegisterAccount {
-        account_id: supply_account_id,
-    } = result
-    else {
-        anyhow::bail!("Expected RegisterAccount return value");
-    };
+    let supply_account_id = new_account(&mut ctx, true, None).await?;
 
     // Create token with both private definition and supply
     let name = "A NAME".to_owned();
@@ -722,20 +529,7 @@ async fn create_token_with_private_definition_and_supply() -> Result<()> {
     );
 
     // Create recipient account
-    let result = wallet::cli::execute_subcommand(
-        ctx.wallet_mut(),
-        Command::Account(AccountSubcommand::New(NewSubcommand::Private {
-            cci: None,
-            label: None,
-        })),
-    )
-    .await?;
-    let SubcommandReturnValue::RegisterAccount {
-        account_id: recipient_account_id,
-    } = result
-    else {
-        anyhow::bail!("Expected RegisterAccount return value");
-    };
+    let recipient_account_id = new_account(&mut ctx, true, None).await?;
 
     // Transfer tokens
     let transfer_amount = 7;
@@ -804,52 +598,13 @@ async fn shielded_token_transfer() -> Result<()> {
     let mut ctx = TestContext::new().await?;
 
     // Create token definition account (public)
-    let result = wallet::cli::execute_subcommand(
-        ctx.wallet_mut(),
-        Command::Account(AccountSubcommand::New(NewSubcommand::Public {
-            cci: None,
-            label: None,
-        })),
-    )
-    .await?;
-    let SubcommandReturnValue::RegisterAccount {
-        account_id: definition_account_id,
-    } = result
-    else {
-        anyhow::bail!("Expected RegisterAccount return value");
-    };
+    let definition_account_id = new_account(&mut ctx, false, None).await?;
 
     // Create supply account (public)
-    let result = wallet::cli::execute_subcommand(
-        ctx.wallet_mut(),
-        Command::Account(AccountSubcommand::New(NewSubcommand::Public {
-            cci: None,
-            label: None,
-        })),
-    )
-    .await?;
-    let SubcommandReturnValue::RegisterAccount {
-        account_id: supply_account_id,
-    } = result
-    else {
-        anyhow::bail!("Expected RegisterAccount return value");
-    };
+    let supply_account_id = new_account(&mut ctx, false, None).await?;
 
     // Create recipient account (private) for shielded transfer
-    let result = wallet::cli::execute_subcommand(
-        ctx.wallet_mut(),
-        Command::Account(AccountSubcommand::New(NewSubcommand::Private {
-            cci: None,
-            label: None,
-        })),
-    )
-    .await?;
-    let SubcommandReturnValue::RegisterAccount {
-        account_id: recipient_account_id,
-    } = result
-    else {
-        anyhow::bail!("Expected RegisterAccount return value");
-    };
+    let recipient_account_id = new_account(&mut ctx, true, None).await?;
 
     // Create token
     let name = "A NAME".to_owned();
@@ -884,10 +639,7 @@ async fn shielded_token_transfer() -> Result<()> {
     tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await;
 
     // Verify supply account balance
-    let supply_acc = ctx
-        .sequencer_client()
-        .get_account(supply_account_id)
-        .await?;
+    let supply_acc = get_account(&ctx, supply_account_id).await?;
     let token_holding = TokenHolding::try_from(&supply_acc.data)?;
     assert_eq!(
         token_holding,
@@ -928,52 +680,13 @@ async fn deshielded_token_transfer() -> Result<()> {
     let mut ctx = TestContext::new().await?;
 
     // Create token definition account (public)
-    let result = wallet::cli::execute_subcommand(
-        ctx.wallet_mut(),
-        Command::Account(AccountSubcommand::New(NewSubcommand::Public {
-            cci: None,
-            label: None,
-        })),
-    )
-    .await?;
-    let SubcommandReturnValue::RegisterAccount {
-        account_id: definition_account_id,
-    } = result
-    else {
-        anyhow::bail!("Expected RegisterAccount return value");
-    };
+    let definition_account_id = new_account(&mut ctx, false, None).await?;
 
     // Create supply account (private)
-    let result = wallet::cli::execute_subcommand(
-        ctx.wallet_mut(),
-        Command::Account(AccountSubcommand::New(NewSubcommand::Private {
-            cci: None,
-            label: None,
-        })),
-    )
-    .await?;
-    let SubcommandReturnValue::RegisterAccount {
-        account_id: supply_account_id,
-    } = result
-    else {
-        anyhow::bail!("Expected RegisterAccount return value");
-    };
+    let supply_account_id = new_account(&mut ctx, true, None).await?;
 
     // Create recipient account (public) for deshielded transfer
-    let result = wallet::cli::execute_subcommand(
-        ctx.wallet_mut(),
-        Command::Account(AccountSubcommand::New(NewSubcommand::Public {
-            cci: None,
-            label: None,
-        })),
-    )
-    .await?;
-    let SubcommandReturnValue::RegisterAccount {
-        account_id: recipient_account_id,
-    } = result
-    else {
-        anyhow::bail!("Expected RegisterAccount return value");
-    };
+    let recipient_account_id = new_account(&mut ctx, false, None).await?;
 
     // Create token with private supply
     let name = "A NAME".to_owned();
@@ -1029,10 +742,7 @@ async fn deshielded_token_transfer() -> Result<()> {
     );
 
     // Verify recipient balance
-    let recipient_acc = ctx
-        .sequencer_client()
-        .get_account(recipient_account_id)
-        .await?;
+    let recipient_acc = get_account(&ctx, recipient_account_id).await?;
     let token_holding = TokenHolding::try_from(&recipient_acc.data)?;
     assert_eq!(
         token_holding,
@@ -1052,36 +762,10 @@ async fn token_claiming_path_with_private_accounts() -> Result<()> {
     let mut ctx = TestContext::new().await?;
 
     // Create token definition account (private)
-    let result = wallet::cli::execute_subcommand(
-        ctx.wallet_mut(),
-        Command::Account(AccountSubcommand::New(NewSubcommand::Private {
-            cci: None,
-            label: None,
-        })),
-    )
-    .await?;
-    let SubcommandReturnValue::RegisterAccount {
-        account_id: definition_account_id,
-    } = result
-    else {
-        anyhow::bail!("Expected RegisterAccount return value");
-    };
+    let definition_account_id = new_account(&mut ctx, true, None).await?;
 
     // Create supply account (private)
-    let result = wallet::cli::execute_subcommand(
-        ctx.wallet_mut(),
-        Command::Account(AccountSubcommand::New(NewSubcommand::Private {
-            cci: None,
-            label: None,
-        })),
-    )
-    .await?;
-    let SubcommandReturnValue::RegisterAccount {
-        account_id: supply_account_id,
-    } = result
-    else {
-        anyhow::bail!("Expected RegisterAccount return value");
-    };
+    let supply_account_id = new_account(&mut ctx, true, None).await?;
 
     // Create token
     let name = "A NAME".to_owned();
@@ -1099,20 +783,7 @@ async fn token_claiming_path_with_private_accounts() -> Result<()> {
     tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await;
 
     // Create new private account for claiming path
-    let result = wallet::cli::execute_subcommand(
-        ctx.wallet_mut(),
-        Command::Account(AccountSubcommand::New(NewSubcommand::Private {
-            cci: None,
-            label: None,
-        })),
-    )
-    .await?;
-    let SubcommandReturnValue::RegisterAccount {
-        account_id: recipient_account_id,
-    } = result
-    else {
-        anyhow::bail!("Expected RegisterAccount return value");
-    };
+    let recipient_account_id = new_account(&mut ctx, true, None).await?;
 
     // Get keys for foreign mint (claiming path)
     let holder = ctx
@@ -1143,8 +814,7 @@ async fn token_claiming_path_with_private_accounts() -> Result<()> {
     tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await;
 
     // Sync to claim the account
-    let command = Command::Account(AccountSubcommand::SyncPrivate {});
-    wallet::cli::execute_subcommand(ctx.wallet_mut(), command).await?;
+    sync_private(&mut ctx).await?;
 
     // Verify commitment exists
     let recipient_commitment = ctx
@@ -1224,10 +894,7 @@ async fn create_token_using_labels() -> Result<()> {
     info!("Waiting for next block creation");
     tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await;
 
-    let definition_acc = ctx
-        .sequencer_client()
-        .get_account(definition_account_id)
-        .await?;
+    let definition_acc = get_account(&ctx, definition_account_id).await?;
     let token_definition = TokenDefinition::try_from(&definition_acc.data)?;
 
     assert_eq!(definition_acc.program_owner, programs::token().id());
@@ -1240,10 +907,7 @@ async fn create_token_using_labels() -> Result<()> {
         }
     );
 
-    let supply_acc = ctx
-        .sequencer_client()
-        .get_account(supply_account_id)
-        .await?;
+    let supply_acc = get_account(&ctx, supply_account_id).await?;
     let token_holding = TokenHolding::try_from(&supply_acc.data)?;
     assert_eq!(
         token_holding,
@@ -1263,20 +927,7 @@ async fn transfer_token_using_from_label() -> Result<()> {
     let mut ctx = TestContext::new().await?;
 
     // Create definition account
-    let result = wallet::cli::execute_subcommand(
-        ctx.wallet_mut(),
-        Command::Account(AccountSubcommand::New(NewSubcommand::Public {
-            cci: None,
-            label: None,
-        })),
-    )
-    .await?;
-    let SubcommandReturnValue::RegisterAccount {
-        account_id: definition_account_id,
-    } = result
-    else {
-        anyhow::bail!("Expected RegisterAccount return value");
-    };
+    let definition_account_id = new_account(&mut ctx, false, None).await?;
 
     // Create supply account with a label
     let supply_label = Label::new("token-supply-sender");
@@ -1296,20 +947,7 @@ async fn transfer_token_using_from_label() -> Result<()> {
     };
 
     // Create recipient account
-    let result = wallet::cli::execute_subcommand(
-        ctx.wallet_mut(),
-        Command::Account(AccountSubcommand::New(NewSubcommand::Public {
-            cci: None,
-            label: None,
-        })),
-    )
-    .await?;
-    let SubcommandReturnValue::RegisterAccount {
-        account_id: recipient_account_id,
-    } = result
-    else {
-        anyhow::bail!("Expected RegisterAccount return value");
-    };
+    let recipient_account_id = new_account(&mut ctx, false, None).await?;
 
     // Create token
     let total_supply = 50;
@@ -1340,10 +978,7 @@ async fn transfer_token_using_from_label() -> Result<()> {
     info!("Waiting for next block creation");
     tokio::time::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS)).await;
 
-    let recipient_acc = ctx
-        .sequencer_client()
-        .get_account(recipient_account_id)
-        .await?;
+    let recipient_acc = get_account(&ctx, recipient_account_id).await?;
     let token_holding = TokenHolding::try_from(&recipient_acc.data)?;
     assert_eq!(
         token_holding,

integration_tests/tests/tps.rs

@@ -23,7 +23,8 @@ use lee::{
     public_transaction as putx,
 };
 use lee_core::{
-    EncryptedAccountData, InputAccountIdentity, MembershipProof, NullifierPublicKey,
+    DUMMY_COMMITMENT_HASH, EncryptedAccountData, InputAccountIdentity, MembershipProof,
+    NullifierPublicKey,
     account::{AccountWithMetadata, Nonce, data::Data},
     encryption::ViewingPublicKey,
 };
@@ -314,6 +315,7 @@ fn build_privacy_transaction() -> PrivacyPreservingTransaction {
                 npk: recipient_npk,
                 ssk: recipient_ss,
                 identifier: 0,
+                commitment_root: DUMMY_COMMITMENT_HASH,
             },
         ],
         &program.into(),

integration_tests/tests/wallet_ffi.rs

@@ -16,6 +16,7 @@ use std::{
     ffi::{CStr, CString, c_char},
     io::Write as _,
     path::Path,
+    str::FromStr as _,
     time::Duration,
 };
 
@@ -30,9 +31,11 @@ use log::info;
 use tempfile::tempdir;
 use wallet::{account::HumanReadableAccount, program_facades::vault::Vault};
 use wallet_ffi::{
-    FfiAccount, FfiAccountIdentity, FfiAccountList, FfiBytes32, FfiPrivateAccountKeys,
-    FfiProgramId, FfiPublicAccountKey, FfiTransferResult, FfiU128, WalletHandle, error,
+    FfiAccount, FfiAccountIdWithPrivacy, FfiAccountIdentity, FfiAccountList, FfiBytes32,
+    FfiPrivateAccountKeys, FfiProgramId, FfiPublicAccountKey, FfiTransferResult, FfiU128,
+    WalletHandle, error,
     generic_transaction::{FfiProgramWithDependencies, FfiTransactionResult},
+    label::{AccountIdResolvedFromLabel, LabelAvailability, LabelList},
     wallet::FfiCreateWalletOutput,
 };
 
@@ -257,6 +260,29 @@ unsafe extern "C" {
     fn wallet_ffi_free_transaction_result(result: *mut FfiTransactionResult);
 
     fn wallet_ffi_free_account_identity(account_identity: *mut FfiAccountIdentity);
+
+    fn wallet_ffi_check_label_available(
+        handle: *mut WalletHandle,
+        label: *const c_char,
+    ) -> LabelAvailability;
+
+    fn wallet_ffi_add_label(
+        handle: *mut WalletHandle,
+        label: *const c_char,
+        account_id_with_privacy: FfiAccountIdWithPrivacy,
+    ) -> error::WalletFfiError;
+
+    fn wallet_ffi_resolve_label(
+        handle: *mut WalletHandle,
+        label: *const c_char,
+    ) -> AccountIdResolvedFromLabel;
+
+    fn wallet_ffi_get_all_labels_for_account(
+        handle: *mut WalletHandle,
+        account_id_with_privacy: FfiAccountIdWithPrivacy,
+    ) -> LabelList;
+
+    fn wallet_ffi_free_label_list(label_list: *mut LabelList) -> error::WalletFfiError;
 }
 
 fn new_wallet_ffi_with_test_context_config(
@@ -1926,3 +1952,125 @@ fn test_wallet_ffi_vault_balance_and_claim_private() -> Result<()> {
 
     Ok(())
 }
+
+#[test]
+fn test_wallet_ffi_single_label() -> Result<()> {
+    let ctx = BlockingTestContext::new()?;
+    let home = tempfile::tempdir()?;
+    let FfiCreateWalletOutput {
+        wallet: wallet_ffi_handle,
+        mnemonic: _,
+    } = new_wallet_ffi_with_test_context_config(&ctx, home.path())?;
+
+    let mut out_account_id_1 = FfiBytes32::from_bytes([0; 32]);
+    unsafe {
+        wallet_ffi_create_account_public(wallet_ffi_handle, &raw mut out_account_id_1).unwrap();
+    }
+
+    info!("Waiting for next block creation");
+    std::thread::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS));
+
+    let lab_1 = CString::from_str("LABEL1").unwrap().into_raw();
+
+    let lab_1_availability = unsafe { wallet_ffi_check_label_available(wallet_ffi_handle, lab_1) };
+
+    assert_eq!(lab_1_availability.error, error::WalletFfiError::Success);
+    assert!(lab_1_availability.is_available);
+
+    let acc_1_id_with_privacy = FfiAccountIdWithPrivacy {
+        account_id: out_account_id_1,
+        is_private: false,
+    };
+
+    let err = unsafe { wallet_ffi_add_label(wallet_ffi_handle, lab_1, acc_1_id_with_privacy) };
+
+    assert_eq!(err, error::WalletFfiError::Success);
+
+    let lab_1_availability = unsafe { wallet_ffi_check_label_available(wallet_ffi_handle, lab_1) };
+
+    assert!(!lab_1_availability.is_available);
+
+    let acc_resolved = unsafe { wallet_ffi_resolve_label(wallet_ffi_handle, lab_1) };
+
+    assert_eq!(acc_resolved.account_id, acc_1_id_with_privacy);
+
+    unsafe {
+        wallet_ffi_free_string(lab_1);
+        wallet_ffi_destroy(wallet_ffi_handle);
+    }
+
+    Ok(())
+}
+
+#[test]
+fn test_wallet_ffi_more_labels() -> Result<()> {
+    let ctx = BlockingTestContext::new()?;
+    let home = tempfile::tempdir()?;
+    let FfiCreateWalletOutput {
+        wallet: wallet_ffi_handle,
+        mnemonic: _,
+    } = new_wallet_ffi_with_test_context_config(&ctx, home.path())?;
+
+    let mut out_account_id_1 = FfiBytes32::from_bytes([0; 32]);
+    unsafe {
+        wallet_ffi_create_account_public(wallet_ffi_handle, &raw mut out_account_id_1).unwrap();
+    }
+
+    info!("Waiting for next block creation");
+    std::thread::sleep(Duration::from_secs(TIME_TO_WAIT_FOR_BLOCK_SECONDS));
+
+    let lab_1 = CString::from_str("LABEL1").unwrap().into_raw();
+    let lab_2 = CString::from_str("LABEL2").unwrap().into_raw();
+    let lab_3 = CString::from_str("LABEL3").unwrap().into_raw();
+
+    let acc_1_id_with_privacy = FfiAccountIdWithPrivacy {
+        account_id: out_account_id_1,
+        is_private: false,
+    };
+
+    let err = unsafe { wallet_ffi_add_label(wallet_ffi_handle, lab_1, acc_1_id_with_privacy) };
+
+    assert_eq!(err, error::WalletFfiError::Success);
+
+    let err = unsafe { wallet_ffi_add_label(wallet_ffi_handle, lab_2, acc_1_id_with_privacy) };
+
+    assert_eq!(err, error::WalletFfiError::Success);
+
+    let err = unsafe { wallet_ffi_add_label(wallet_ffi_handle, lab_3, acc_1_id_with_privacy) };
+
+    assert_eq!(err, error::WalletFfiError::Success);
+
+    let mut label_list_for_out_acc =
+        unsafe { wallet_ffi_get_all_labels_for_account(wallet_ffi_handle, acc_1_id_with_privacy) };
+
+    assert_eq!(label_list_for_out_acc.error, error::WalletFfiError::Success);
+    assert_eq!(label_list_for_out_acc.labels_size, 3);
+
+    let lab_ref_1 = unsafe { &*label_list_for_out_acc.labels_data.add(0) };
+    let lab_ref_c_str_1 = unsafe { CStr::from_ptr(*lab_ref_1) };
+
+    assert_eq!(lab_ref_c_str_1.to_str().unwrap(), "LABEL1");
+
+    let lab_ref_2 = unsafe { &*label_list_for_out_acc.labels_data.add(1) };
+    let lab_ref_c_str_2 = unsafe { CStr::from_ptr(*lab_ref_2) };
+
+    assert_eq!(lab_ref_c_str_2.to_str().unwrap(), "LABEL2");
+
+    let lab_ref_3 = unsafe { &*label_list_for_out_acc.labels_data.add(2) };
+    let lab_ref_c_str_3 = unsafe { CStr::from_ptr(*lab_ref_3) };
+
+    assert_eq!(lab_ref_c_str_3.to_str().unwrap(), "LABEL3");
+
+    let err = unsafe { wallet_ffi_free_label_list(&raw mut label_list_for_out_acc) };
+
+    assert_eq!(err, error::WalletFfiError::Success);
+
+    unsafe {
+        wallet_ffi_free_string(lab_1);
+        wallet_ffi_free_string(lab_2);
+        wallet_ffi_free_string(lab_3);
+        wallet_ffi_destroy(wallet_ffi_handle);
+    }
+
+    Ok(())
+}

lee/key_protocol/src/key_management/ephemeral_key_holder.rs

@@ -48,14 +48,3 @@ impl EphemeralKeyHolder {
         self.shared_secret
     }
 }
-
-/// Encapsulates a fresh shared secret toward `vpk` and returns `(shared_secret, ciphertext)`.
-///
-/// Used when the local side is acting as an "ephemeral receiver" — i.e. generating a
-/// one-sided encryption that only the holder of the VSK can decrypt.
-#[must_use]
-pub fn produce_one_sided_shared_secret_receiver(
-    vpk: &ViewingPublicKey,
-) -> (SharedSecretKey, EphemeralPublicKey) {
-    SharedSecretKey::encapsulate(vpk)
-}

lee/key_protocol/src/key_management/group_key_holder.rs

@@ -1,7 +1,7 @@
 use aes_gcm::{Aes256Gcm, KeyInit as _, aead::Aead as _};
 use lee_core::{
     SharedSecretKey,
-    encryption::{EphemeralPublicKey, ViewingPublicKey},
+    encryption::{EphemeralPublicKey, ML_KEM_768_CIPHERTEXT_LEN, ViewingPublicKey},
     program::{PdaSeed, ProgramId},
 };
 use rand::{RngCore as _, rngs::OsRng};
@@ -150,7 +150,7 @@ impl GroupKeyHolder {
     ///
     /// Uses ML-KEM-768 encapsulation to derive a shared secret, then AES-256-GCM to encrypt
     /// the payload. The returned bytes are
-    /// `kem_ciphertext (1088) || nonce (12) || ciphertext+tag (48)` = 1148 bytes.
+    /// `kem_ciphertext (ML_KEM_768_CIPHERTEXT_LEN) || nonce (12) || ciphertext+tag (48)`.
     ///
     /// Each call generates a fresh KEM encapsulation, so two seals of the same holder produce
     /// different ciphertexts.
@@ -170,7 +170,7 @@ impl GroupKeyHolder {
             .encrypt(&nonce, self.gms.as_ref())
             .expect("AES-GCM encryption should not fail with valid key/nonce");
 
-        let capacity = 1088_usize
+        let capacity = ML_KEM_768_CIPHERTEXT_LEN
             .checked_add(12)
             .and_then(|n| n.checked_add(ciphertext.len()))
             .expect("seal capacity overflow");
@@ -186,21 +186,21 @@ impl GroupKeyHolder {
     /// Returns `Err` if the ciphertext is too short or the AES-GCM authentication tag
     /// doesn't verify (wrong key or tampered data).
     pub fn unseal(sealed: &[u8], own_key: &SealingSecretKey) -> Result<Self, SealError> {
-        // kem_ciphertext (1088) + nonce (12) = header, then AES-GCM tag (16) minimum.
-        const KEM_CT_LEN: usize = 1088;
-        const HEADER_LEN: usize = KEM_CT_LEN + 12;
+        // kem_ciphertext (ML_KEM_768_CIPHERTEXT_LEN) + nonce (12) = header, then AES-GCM tag (16)
+        // minimum.
+        const HEADER_LEN: usize = ML_KEM_768_CIPHERTEXT_LEN + 12;
         const MIN_LEN: usize = HEADER_LEN + 16;
 
         if sealed.len() < MIN_LEN {
             return Err(SealError::TooShort);
         }
 
-        let kem_ct = EphemeralPublicKey(sealed[..KEM_CT_LEN].to_vec());
-        let nonce = aes_gcm::Nonce::from_slice(&sealed[KEM_CT_LEN..HEADER_LEN]);
+        let kem_ct = EphemeralPublicKey(sealed[..ML_KEM_768_CIPHERTEXT_LEN].to_vec());
+        let nonce = aes_gcm::Nonce::from_slice(&sealed[ML_KEM_768_CIPHERTEXT_LEN..HEADER_LEN]);
         let ciphertext = &sealed[HEADER_LEN..];
 
         let shared = SharedSecretKey::decapsulate(&kem_ct, &own_key.d, &own_key.z)
-            .expect("key_protocol::group_key_holder::GroupKeyHolder::unseal: KEM_CT_LEN guarantees exactly 1088 bytes");
+            .expect("key_protocol::group_key_holder::GroupKeyHolder::unseal: ML_KEM_768_CIPHERTEXT_LEN guarantees exactly 1088 bytes");
         let aes_key = Self::seal_kdf(&shared);
         let cipher = Aes256Gcm::new(&aes_key.into());
 

lee/key_protocol/src/key_management/key_tree/chain_index.rs

@@ -139,7 +139,7 @@ impl ChainIndex {
             .map(|item| Self(item.into_iter().copied().collect()))
     }
 
-    pub fn chain_ids_at_depth(depth: usize) -> impl Iterator<Item = Self> {
+    fn collect_chain_ids_at_depth(depth: usize) -> Vec<Self> {
         let mut stack = vec![Self(vec![0; depth])];
         let mut cumulative_stack = vec![Self(vec![0; depth])];
 
@@ -152,23 +152,18 @@ impl ChainIndex {
             }
         }
 
-        cumulative_stack.into_iter().unique()
+        cumulative_stack
+    }
+
+    pub fn chain_ids_at_depth(depth: usize) -> impl Iterator<Item = Self> {
+        Self::collect_chain_ids_at_depth(depth).into_iter().unique()
     }
 
     pub fn chain_ids_at_depth_rev(depth: usize) -> impl Iterator<Item = Self> {
-        let mut stack = vec![Self(vec![0; depth])];
-        let mut cumulative_stack = vec![Self(vec![0; depth])];
-
-        while let Some(top_id) = stack.pop() {
-            if let Some(collapsed_id) = top_id.collapse_back() {
-                for id in collapsed_id.shuffle_iter() {
-                    stack.push(id.clone());
-                    cumulative_stack.push(id);
-                }
-            }
-        }
-
-        cumulative_stack.into_iter().rev().unique()
+        Self::collect_chain_ids_at_depth(depth)
+            .into_iter()
+            .rev()
+            .unique()
     }
 }
 

lee/key_protocol/src/key_management/key_tree/keys_private.rs

@@ -6,7 +6,7 @@ use sha2::Digest as _;
 
 use crate::key_management::{
     KeyChain,
-    key_tree::traits::KeyTreeNode,
+    key_tree::{split_hash, traits::KeyTreeNode},
     secret_holders::{PrivateKeyHolder, SecretSpendingKey},
 };
 
@@ -23,38 +23,11 @@ impl ChildKeysPrivate {
     #[must_use]
     pub fn root(seed: [u8; 64]) -> Self {
         let hash_value = hmac_sha512::HMAC::mac(seed, b"LEE_master_priv");
+        let (first, ccc) = split_hash(&hash_value);
 
-        let ssk = SecretSpendingKey(
-            *hash_value
-                .first_chunk::<32>()
-                .expect("hash_value is 64 bytes, must be safe to get first 32"),
-        );
-        let ccc = *hash_value
-            .last_chunk::<32>()
-            .expect("hash_value is 64 bytes, must be safe to get last 32");
+        let ssk = SecretSpendingKey(first);
 
-        let nsk = ssk.generate_nullifier_secret_key(None);
-        let vsk = ssk.generate_viewing_secret_seed_key(None);
-
-        let npk = NullifierPublicKey::from(&nsk);
-        let vpk = ViewingPublicKey::from(&vsk);
-
-        Self {
-            value: (
-                KeyChain {
-                    secret_spending_key: ssk,
-                    nullifier_public_key: npk,
-                    viewing_public_key: vpk,
-                    private_key_holder: PrivateKeyHolder {
-                        nullifier_secret_key: nsk,
-                        viewing_secret_key: vsk,
-                    },
-                },
-                BTreeMap::from_iter([(PrivateAccountKind::Regular(0), lee::Account::default())]),
-            ),
-            ccc,
-            cci: None,
-        }
+        Self::from_ssk_and_ccc(ssk, ccc, None)
     }
 
     #[must_use]
@@ -77,18 +50,16 @@ impl ChildKeysPrivate {
         input.extend_from_slice(&cci.to_be_bytes());
 
         let hash_value = hmac_sha512::HMAC::mac(input, self.ccc);
+        let (first, ccc) = split_hash(&hash_value);
 
-        let ssk = SecretSpendingKey(
-            *hash_value
-                .first_chunk::<32>()
-                .expect("hash_value is 64 bytes, must be safe to get first 32"),
-        );
-        let ccc = *hash_value
-            .last_chunk::<32>()
-            .expect("hash_value is 64 bytes, must be safe to get last 32");
+        let ssk = SecretSpendingKey(first);
 
-        let nsk = ssk.generate_nullifier_secret_key(Some(cci));
-        let vsk = ssk.generate_viewing_secret_seed_key(Some(cci));
+        Self::from_ssk_and_ccc(ssk, ccc, Some(cci))
+    }
+
+    fn from_ssk_and_ccc(ssk: SecretSpendingKey, ccc: [u8; 32], cci: Option<u32>) -> Self {
+        let nsk = ssk.generate_nullifier_secret_key(cci);
+        let vsk = ssk.generate_viewing_secret_seed_key(cci);
 
         let npk = NullifierPublicKey::from(&nsk);
         let vpk = ViewingPublicKey::from(&vsk);
@@ -107,7 +78,7 @@ impl ChildKeysPrivate {
                 BTreeMap::from_iter([(PrivateAccountKind::Regular(0), lee::Account::default())]),
             ),
             ccc,
-            cci: Some(cci),
+            cci,
         }
     }
 }
@@ -137,16 +108,16 @@ mod tests {
     use super::*;
     use crate::key_management::{self, secret_holders::ViewingSecretKey};
 
+    const SEED: [u8; 64] = [
+        252, 56, 204, 83, 232, 123, 209, 188, 187, 167, 39, 213, 71, 39, 58, 65, 125, 134, 255, 49,
+        43, 108, 92, 53, 173, 164, 94, 142, 150, 74, 21, 163, 43, 144, 226, 87, 199, 18, 129, 223,
+        176, 198, 5, 150, 157, 70, 210, 254, 14, 105, 89, 191, 246, 27, 52, 170, 56, 114, 39, 38,
+        118, 197, 205, 225,
+    ];
+
     #[test]
     fn master_key_generation() {
-        let seed: [u8; 64] = [
-            252, 56, 204, 83, 232, 123, 209, 188, 187, 167, 39, 213, 71, 39, 58, 65, 125, 134, 255,
-            49, 43, 108, 92, 53, 173, 164, 94, 142, 150, 74, 21, 163, 43, 144, 226, 87, 199, 18,
-            129, 223, 176, 198, 5, 150, 157, 70, 210, 254, 14, 105, 89, 191, 246, 27, 52, 170, 56,
-            114, 39, 38, 118, 197, 205, 225,
-        ];
-
-        let keys = ChildKeysPrivate::root(seed);
+        let keys = ChildKeysPrivate::root(SEED);
 
         let expected_ssk = key_management::secret_holders::SecretSpendingKey([
             246, 79, 26, 124, 135, 95, 52, 51, 201, 27, 48, 194, 2, 144, 51, 219, 245, 128, 139,
@@ -179,6 +150,7 @@ mod tests {
             ],
         );
 
+        // Length matches MlKem768EncapsulationKey::LEN.
         let expected_vpk: [u8; 1184] = [
             127, 229, 162, 212, 104, 117, 4, 150, 192, 103, 122, 195, 14, 35, 12, 60, 52, 23, 220,
             150, 100, 203, 34, 34, 127, 232, 156, 43, 218, 109, 6, 160, 67, 35, 210, 194, 25, 181,
@@ -253,14 +225,7 @@ mod tests {
 
     #[test]
     fn child_keys_generation() {
-        let seed: [u8; 64] = [
-            252, 56, 204, 83, 232, 123, 209, 188, 187, 167, 39, 213, 71, 39, 58, 65, 125, 134, 255,
-            49, 43, 108, 92, 53, 173, 164, 94, 142, 150, 74, 21, 163, 43, 144, 226, 87, 199, 18,
-            129, 223, 176, 198, 5, 150, 157, 70, 210, 254, 14, 105, 89, 191, 246, 27, 52, 170, 56,
-            114, 39, 38, 118, 197, 205, 225,
-        ];
-
-        let root_node = ChildKeysPrivate::root(seed);
+        let root_node = ChildKeysPrivate::root(SEED);
         let child_node = ChildKeysPrivate::nth_child(&root_node, 42_u32);
 
         let expected_ssk = key_management::secret_holders::SecretSpendingKey([
@@ -293,6 +258,7 @@ mod tests {
             ],
         );
 
+        // Length matches MlKem768EncapsulationKey::LEN.
         let expected_vpk: [u8; 1184] = [
             215, 229, 207, 120, 148, 177, 148, 197, 72, 222, 134, 3, 231, 146, 123, 226, 36, 84,
             232, 179, 205, 16, 241, 142, 9, 81, 58, 54, 12, 115, 148, 182, 19, 245, 22, 203, 57,

lee/key_protocol/src/key_management/key_tree/keys_public.rs

@@ -1,7 +1,7 @@
 use k256::elliptic_curve::PrimeField as _;
 use serde::{Deserialize, Serialize};
 
-use crate::key_management::key_tree::traits::KeyTreeNode;
+use crate::key_management::key_tree::{split_hash, traits::KeyTreeNode};
 
 #[derive(Debug, Serialize, Deserialize, Clone)]
 #[cfg_attr(any(test, feature = "test_utils"), derive(PartialEq, Eq))]
@@ -21,52 +21,32 @@ impl ChildKeysPublic {
     #[must_use]
     pub fn root(seed: [u8; 64]) -> Self {
         let hash_value = hmac_sha512::HMAC::mac(seed, "LEE_master_pub");
+        let (first, cc) = split_hash(&hash_value);
 
-        let sk = lee::PrivateKey::try_new(
-            *hash_value
-                .first_chunk::<32>()
-                .expect("hash_value is 64 bytes, must be safe to get first 32"),
-        )
-        .expect("Expect a valid Private Key");
-        let ssk = lee::PrivateKey::tweak(sk.value()).expect("`key_protocol::key_management::keys_public::root()`: Invalid private key produced from `tweak`");
+        let sk = lee::PrivateKey::try_new(first).expect("Expect a valid Private Key");
 
-        let cc = *hash_value
-            .last_chunk::<32>()
-            .expect("hash_value is 64 bytes, must be safe to get last 32");
-        let pk = lee::PublicKey::new_from_private_key(&ssk);
-
-        Self {
-            sk,
-            ssk,
-            pk,
-            cc,
-            cci: None,
-        }
+        Self::from_sk_and_cc(sk, cc, None)
     }
 
     #[must_use]
     pub fn nth_child(&self, cci: u32) -> Self {
         let hash_value = self.compute_hash_value(cci);
+        let (first, cc) = split_hash(&hash_value);
 
-        let lhs = k256::Scalar::from_repr(
-            (*hash_value
-                .first_chunk::<32>()
-                .expect("hash_value is 64 bytes, must be safe to get first 32"))
-            .into(),
-        )
-        .expect("Expect a valid k256 scalar");
+        let lhs = k256::Scalar::from_repr(first.into()).expect("Expect a valid k256 scalar");
         let rhs =
             k256::Scalar::from_repr((*self.sk.value()).into()).expect("Expect a valid k256 scalar");
 
         let sk = lee::PrivateKey::try_new(lhs.add(&rhs).to_bytes().into())
             .expect("Expect a valid private key");
 
-        let ssk = lee::PrivateKey::tweak(sk.value()).expect("`key_protocol::key_management::keys_public::nth_child()`: Invalid private key produced from `tweak`");
-
-        let cc = *hash_value
-            .last_chunk::<32>()
-            .expect("hash_value is 64 bytes, must be safe to get last 32");
+        Self::from_sk_and_cc(sk, cc, Some(cci))
+    }
 
+    fn from_sk_and_cc(sk: lee::PrivateKey, cc: [u8; 32], cci: Option<u32>) -> Self {
+        let ssk = lee::PrivateKey::tweak(sk.value()).expect(
+            "`key_protocol::key_management::keys_public::ChildKeysPublic`: Invalid private key produced from `tweak`",
+        );
         let pk = lee::PublicKey::new_from_private_key(&ssk);
 
         Self {
@@ -74,7 +54,7 @@ impl ChildKeysPublic {
             ssk,
             pk,
             cc,
-            cci: Some(cci),
+            cci,
         }
     }
 
@@ -128,15 +108,16 @@ mod tests {
 
     use super::*;
 
+    const SEED: [u8; 64] = [
+        88, 189, 37, 237, 199, 125, 151, 226, 69, 153, 165, 113, 191, 69, 188, 221, 9, 34, 173,
+        134, 61, 109, 34, 103, 121, 39, 237, 14, 107, 194, 24, 194, 191, 14, 237, 185, 12, 87, 22,
+        227, 38, 71, 17, 144, 251, 118, 217, 115, 33, 222, 201, 61, 203, 246, 121, 214, 6, 187,
+        148, 92, 44, 253, 210, 37,
+    ];
+
     #[test]
     fn master_keys_generation() {
-        let seed = [
-            88, 189, 37, 237, 199, 125, 151, 226, 69, 153, 165, 113, 191, 69, 188, 221, 9, 34, 173,
-            134, 61, 109, 34, 103, 121, 39, 237, 14, 107, 194, 24, 194, 191, 14, 237, 185, 12, 87,
-            22, 227, 38, 71, 17, 144, 251, 118, 217, 115, 33, 222, 201, 61, 203, 246, 121, 214, 6,
-            187, 148, 92, 44, 253, 210, 37,
-        ];
-        let keys = ChildKeysPublic::root(seed);
+        let keys = ChildKeysPublic::root(SEED);
 
         let expected_cc = [
             238, 94, 84, 154, 56, 224, 80, 218, 133, 249, 179, 222, 9, 24, 17, 252, 120, 127, 222,
@@ -169,13 +150,7 @@ mod tests {
 
     #[test]
     fn child_keys_generation() {
-        let seed = [
-            88, 189, 37, 237, 199, 125, 151, 226, 69, 153, 165, 113, 191, 69, 188, 221, 9, 34, 173,
-            134, 61, 109, 34, 103, 121, 39, 237, 14, 107, 194, 24, 194, 191, 14, 237, 185, 12, 87,
-            22, 227, 38, 71, 17, 144, 251, 118, 217, 115, 33, 222, 201, 61, 203, 246, 121, 214, 6,
-            187, 148, 92, 44, 253, 210, 37,
-        ];
-        let root_keys = ChildKeysPublic::root(seed);
+        let root_keys = ChildKeysPublic::root(SEED);
         let cci = (2_u32).pow(31) + 13;
         let child_keys = ChildKeysPublic::nth_child(&root_keys, cci);
 

lee/key_protocol/src/key_management/key_tree/mod.rs

@@ -39,16 +39,7 @@ impl<N: KeyTreeNode> KeyTree<N> {
             .try_into()
             .expect("SeedHolder seed is 64 bytes long");
 
-        let root_keys = N::from_seed(seed_fit);
-        let account_id_map = root_keys
-            .account_ids()
-            .map(|id| (id, ChainIndex::root()))
-            .collect();
-
-        Self {
-            key_map: BTreeMap::from_iter([(ChainIndex::root(), root_keys)]),
-            account_id_map,
-        }
+        Self::new_from_root(N::from_seed(seed_fit))
     }
 
     pub fn new_from_root(root: N) -> Self {
@@ -63,6 +54,15 @@ impl<N: KeyTreeNode> KeyTree<N> {
         }
     }
 
+    fn insert_child(&mut self, child_keys: N, chain_index: ChainIndex) -> ChainIndex {
+        for account_id in child_keys.account_ids() {
+            self.account_id_map.insert(account_id, chain_index.clone());
+        }
+        self.key_map.insert(chain_index.clone(), child_keys);
+
+        chain_index
+    }
+
     pub fn generate_new_node(&mut self, parent_cci: &ChainIndex) -> Option<ChainIndex> {
         let parent_keys = self.key_map.get(parent_cci)?;
         let next_child_id = self
@@ -71,14 +71,8 @@ impl<N: KeyTreeNode> KeyTree<N> {
         let next_cci = parent_cci.nth_child(next_child_id);
 
         let child_keys = parent_keys.derive_child(next_child_id);
-        let account_ids = child_keys.account_ids();
 
-        for account_id in account_ids {
-            self.account_id_map.insert(account_id, next_cci.clone());
-        }
-        self.key_map.insert(next_cci.clone(), child_keys);
-
-        Some(next_cci)
+        Some(self.insert_child(child_keys, next_cci))
     }
 
     pub fn fill_node(&mut self, chain_index: &ChainIndex) -> Option<ChainIndex> {
@@ -86,14 +80,8 @@ impl<N: KeyTreeNode> KeyTree<N> {
         let child_id = *chain_index.chain().last()?;
 
         let child_keys = parent_keys.derive_child(child_id);
-        let account_ids = child_keys.account_ids();
 
-        for account_id in account_ids {
-            self.account_id_map.insert(account_id, chain_index.clone());
-        }
-        self.key_map.insert(chain_index.clone(), child_keys);
-
-        Some(chain_index.clone())
+        Some(self.insert_child(child_keys, chain_index.clone()))
     }
 
     #[must_use]
@@ -200,24 +188,27 @@ impl<N: KeyTreeNode> KeyTree<N> {
 }
 
 impl KeyTree<ChildKeysPublic> {
+    /// Pairs `cci` with the account ID of the node stored at it.
+    fn account_id_for_cci(&self, cci: ChainIndex) -> Option<(lee::AccountId, ChainIndex)> {
+        let node = self.key_map.get(&cci)?;
+        let account_id = node.account_ids().next()?;
+        Some((account_id, cci))
+    }
+
     /// Generate a new public key node, returning the account ID and chain index.
     pub fn generate_new_public_node(
         &mut self,
         parent_cci: &ChainIndex,
     ) -> Option<(lee::AccountId, ChainIndex)> {
         let cci = self.generate_new_node(parent_cci)?;
-        let node = self.key_map.get(&cci)?;
-        let account_id = node.account_ids().next()?;
-        Some((account_id, cci))
+        self.account_id_for_cci(cci)
     }
 
     /// Generate a new public key node using layered placement, returning the account ID and chain
     /// index.
     pub fn generate_new_public_node_layered(&mut self) -> Option<(lee::AccountId, ChainIndex)> {
         let cci = self.generate_new_node_layered()?;
-        let node = self.key_map.get(&cci)?;
-        let account_id = node.account_ids().next()?;
-        Some((account_id, cci))
+        self.account_id_for_cci(cci)
     }
 
     /// Cleanup of non-initialized accounts in a public tree.
@@ -322,6 +313,16 @@ impl KeyTree<ChildKeysPrivate> {
     }
 }
 
+const fn split_hash(hash_value: &[u8; 64]) -> ([u8; 32], [u8; 32]) {
+    let first = *hash_value
+        .first_chunk::<32>()
+        .expect("hash_value is 64 bytes, must be safe to get first 32");
+    let last = *hash_value
+        .last_chunk::<32>()
+        .expect("hash_value is 64 bytes, must be safe to get last 32");
+    (first, last)
+}
+
 #[cfg(test)]
 mod tests {
     #![expect(clippy::shadow_unrelated, reason = "We don't care about this in tests")]
@@ -339,6 +340,18 @@ mod tests {
         }
     }
 
+    #[test]
+    fn split_hash_splits_into_first_and_last_32_bytes() {
+        let mut hash_value = [0_u8; 64];
+        hash_value[..32].fill(0xAA);
+        hash_value[32..].fill(0xBB);
+
+        let (first, last) = split_hash(&hash_value);
+
+        assert_eq!(first, [0xAA; 32]);
+        assert_eq!(last, [0xBB; 32]);
+    }
+
     #[test]
     fn simple_key_tree() {
         let seed_holder = seed_holder_for_tests();

lee/key_protocol/src/key_management/mod.rs

@@ -25,8 +25,22 @@ impl KeyChain {
     #[must_use]
     pub fn new_os_random() -> Self {
         // Currently dropping SeedHolder at the end of initialization.
-        // Now entirely sure if we need it in the future.
+        // Not entirely sure if we need it in the future.
         let seed_holder = SeedHolder::new_os_random();
+
+        Self::from_seed_holder(&seed_holder)
+    }
+
+    #[must_use]
+    pub fn new_mnemonic(passphrase: &str) -> (Self, bip39::Mnemonic) {
+        // Currently dropping SeedHolder at the end of initialization.
+        // Not entirely sure if we need it in the future.
+        let (seed_holder, mnemonic) = SeedHolder::new_mnemonic(passphrase);
+
+        (Self::from_seed_holder(&seed_holder), mnemonic)
+    }
+
+    fn from_seed_holder(seed_holder: &SeedHolder) -> Self {
         let secret_spending_key = seed_holder.produce_top_secret_key_holder();
 
         let private_key_holder = secret_spending_key.produce_private_key_holder(None);
@@ -42,29 +56,6 @@ impl KeyChain {
         }
     }
 
-    #[must_use]
-    pub fn new_mnemonic(passphrase: &str) -> (Self, bip39::Mnemonic) {
-        // Currently dropping SeedHolder at the end of initialization.
-        // Not entirely sure if we need it in the future.
-        let (seed_holder, mnemonic) = SeedHolder::new_mnemonic(passphrase);
-        let secret_spending_key = seed_holder.produce_top_secret_key_holder();
-
-        let private_key_holder = secret_spending_key.produce_private_key_holder(None);
-
-        let nullifier_public_key = private_key_holder.generate_nullifier_public_key();
-        let viewing_public_key = private_key_holder.generate_viewing_public_key();
-
-        (
-            Self {
-                secret_spending_key,
-                private_key_holder,
-                nullifier_public_key,
-                viewing_public_key,
-            },
-            mnemonic,
-        )
-    }
-
     #[must_use]
     pub fn calculate_shared_secret_receiver(
         &self,

lee/key_protocol/src/key_management/secret_holders.rs

@@ -43,16 +43,7 @@ pub struct PrivateKeyHolder {
 impl SeedHolder {
     #[must_use]
     pub fn new_os_random() -> Self {
-        let mut enthopy_bytes: [u8; 32] = [0; 32];
-        OsRng.fill_bytes(&mut enthopy_bytes);
-
-        let mnemonic = Mnemonic::from_entropy(&enthopy_bytes)
-            .expect("Enthropy must be a multiple of 32 bytes");
-        let seed_wide = mnemonic.to_seed("mnemonic");
-
-        Self {
-            seed: seed_wide.to_vec(),
-        }
+        Self::new_mnemonic("mnemonic").0
     }
 
     #[must_use]
@@ -62,14 +53,8 @@ impl SeedHolder {
 
         let mnemonic =
             Mnemonic::from_entropy(&entropy_bytes).expect("Entropy must be a multiple of 32 bytes");
-        let seed_wide = mnemonic.to_seed(passphrase);
 
-        (
-            Self {
-                seed: seed_wide.to_vec(),
-            },
-            mnemonic,
-        )
+        (Self::from_mnemonic(&mnemonic, passphrase), mnemonic)
     }
 
     #[must_use]
@@ -107,10 +92,7 @@ impl SecretSpendingKey {
         const SUFFIX_1: &[u8; 1] = &[1];
         const SUFFIX_2: &[u8; 19] = &[0; 19];
 
-        let index = match index {
-            None => 0_u32,
-            _ => index.expect("Expect a valid u32"),
-        };
+        let index = index.unwrap_or(0);
 
         let mut hasher = sha2::Sha256::new();
         hasher.update(PREFIX);
@@ -129,10 +111,7 @@ impl SecretSpendingKey {
         const SUFFIX_1: &[u8; 1] = &[2];
         const SUFFIX_2: &[u8; 19] = &[0; 19];
 
-        let index = match index {
-            None => 0_u32,
-            _ => index.expect("Expect a valid u32"),
-        };
+        let index = index.unwrap_or(0);
 
         let mut bytes: Vec<u8> = Vec::with_capacity(64);
         bytes.extend_from_slice(PREFIX);
@@ -146,14 +125,7 @@ impl SecretSpendingKey {
 
         let full_seed = hmac_sha512::HMAC::mac(bytes, b"LEE_viewing_seed");
 
-        ViewingSecretKey::new(
-            *full_seed
-                .first_chunk::<32>()
-                .expect("hash_value is 64 bytes, must be safe to get first 32"),
-            *full_seed
-                .last_chunk::<32>()
-                .expect("hash_value is 64 bytes, must be safe to get last 32"),
-        )
+        Self::generate_viewing_secret_key(full_seed)
     }
 
     #[must_use]
@@ -181,7 +153,7 @@ impl From<&ViewingSecretKey> for ViewingPublicKey {
         seed_bytes[32..].copy_from_slice(&sk.z);
         let dk = <MlKem768 as Kem>::DecapsulationKey::from_seed(Seed::from(seed_bytes));
         Self::from_bytes(dk.encapsulation_key().to_bytes().to_vec())
-            .expect("key_protocol::secret_holders::From<&ViewingSecretKey>: ML-KEM-768 encapsulation key is always 1184 bytes")
+            .expect("key_protocol::secret_holders::From<&ViewingSecretKey>: ML-KEM-768 encapsulation key is always ViewingPublicKey::LEN bytes")
     }
 }
 
@@ -201,7 +173,6 @@ impl PrivateKeyHolder {
 mod tests {
     use super::*;
 
-    // TODO? are these necessary?
     #[test]
     fn seed_generation_test() {
         let seed_holder = SeedHolder::new_os_random();

lee/privacy_preserving_circuit/src/execution_state.rs

@@ -59,46 +59,9 @@ impl ExecutionState {
         program_id: ProgramId,
         program_outputs: Vec<ProgramOutput>,
     ) -> Self {
-        // Build position → (npk, identifier) map for private-PDA pre_states, indexed by position
-        // in `account_identities`. The vec is documented as 1:1 with the program's pre_state
-        // order, so position here matches `pre_state_position` used downstream in
-        // `validate_and_sync_states`.
-        let mut private_pda_npk_by_position: HashMap<usize, (NullifierPublicKey, Identifier)> =
-            HashMap::new();
-        for (pos, account_identity) in account_identities.iter().enumerate() {
-            if let Some((npk, identifier)) = account_identity.npk_if_private_pda() {
-                private_pda_npk_by_position.insert(pos, (npk, identifier));
-            }
-        }
-
-        let block_valid_from = program_outputs
-            .iter()
-            .filter_map(|output| output.block_validity_window.start())
-            .max();
-        let block_valid_until = program_outputs
-            .iter()
-            .filter_map(|output| output.block_validity_window.end())
-            .min();
-        let ts_valid_from = program_outputs
-            .iter()
-            .filter_map(|output| output.timestamp_validity_window.start())
-            .max();
-        let ts_valid_until = program_outputs
-            .iter()
-            .filter_map(|output| output.timestamp_validity_window.end())
-            .min();
-
-        let block_validity_window: BlockValidityWindow = (block_valid_from, block_valid_until)
-            .try_into()
-            .expect(
-                "There should be non empty intersection in the program output block validity windows",
-            );
-        let timestamp_validity_window: TimestampValidityWindow =
-            (ts_valid_from, ts_valid_until)
-                .try_into()
-                .expect(
-                    "There should be non empty intersection in the program output timestamp validity windows",
-                );
+        let private_pda_npk_by_position = build_private_pda_npk_map(account_identities);
+        let (block_validity_window, timestamp_validity_window) =
+            intersect_validity_windows(&program_outputs);
 
         let mut execution_state = Self {
             pre_states: Vec::new(),
@@ -136,49 +99,7 @@ impl ExecutionState {
                 panic!("Insufficient program outputs for chained calls");
             };
 
-            // Check that instruction data in chained call is the instruction data in program output
-            assert_eq!(
-                chained_call.instruction_data, program_output.instruction_data,
-                "Mismatched instruction data between chained call and program output"
-            );
-
-            // Check that `program_output` is consistent with the execution of the corresponding
-            // program.
-            let program_output_words =
-                &to_vec(&program_output).expect("program_output must be serializable");
-            env::verify(chained_call.program_id, program_output_words).unwrap_or_else(
-                |_: Infallible| unreachable!("Infallible error is never constructed"),
-            );
-
-            // Verify that the program output's self_program_id matches the expected program ID.
-            // This ensures the proof commits to which program produced the output.
-            assert_eq!(
-                program_output.self_program_id, chained_call.program_id,
-                "Program output self_program_id does not match chained call program_id"
-            );
-
-            // Verify that the program output's caller_program_id matches the actual caller.
-            // This prevents a malicious user from privately executing an internal function
-            // by spoofing caller_program_id (e.g. passing caller_program_id = self_program_id
-            // to bypass access control checks).
-            assert_eq!(
-                program_output.caller_program_id, caller_program_id,
-                "Program output caller_program_id does not match actual caller"
-            );
-
-            // Check that the program is well behaved.
-            // See the # Programs section for the definition of the `validate_execution` method.
-            let validated_execution = validate_execution(
-                &program_output.pre_states,
-                &program_output.post_states,
-                chained_call.program_id,
-            );
-            if let Err(err) = validated_execution {
-                panic!(
-                    "Invalid program behavior in program {:?}: {err}",
-                    chained_call.program_id
-                );
-            }
+            verify_program_output(&chained_call, caller_program_id, &program_output);
 
             for next_call in program_output.chained_calls.iter().rev() {
                 chained_calls.push_front((next_call.clone(), Some(chained_call.program_id)));
@@ -202,41 +123,8 @@ impl ExecutionState {
             "Inner call without a chained call found",
         );
 
-        // Every private-PDA pre_state must have had its npk bound to its account_id, either via
-        // a `Claim::Pda(seed)` in some program's post_state or via a caller's `pda_seeds`
-        // matching the private derivation. An unbound private-PDA pre_state has no
-        // cryptographic link between the supplied npk and the account_id, and must be rejected.
-        for (pos, account_identity) in account_identities.iter().enumerate() {
-            if account_identity.is_private_pda() {
-                assert!(
-                    execution_state
-                        .private_pda_bound_positions
-                        .contains_key(&pos),
-                    "private PDA pre_state at position {pos} has no proven (seed, npk) binding via Claim::Pda or caller pda_seeds"
-                );
-            }
-        }
-
-        // Check that all modified uninitialized accounts were claimed
-        for (account_id, post) in execution_state
-            .pre_states
-            .iter()
-            .filter(|a| a.account.program_owner == DEFAULT_PROGRAM_ID)
-            .map(|a| {
-                let post = execution_state
-                    .post_states
-                    .get(&a.account_id)
-                    .expect("Post state must exist for pre state");
-                (a, post)
-            })
-            .filter(|(pre_default, post)| pre_default.account != **post)
-            .map(|(pre, post)| (pre.account_id, post))
-        {
-            assert_ne!(
-                post.program_owner, DEFAULT_PROGRAM_ID,
-                "Account {account_id} was modified but not claimed"
-            );
-        }
+        execution_state.assert_all_pda_positions_bound(account_identities);
+        execution_state.assert_modified_accounts_claimed();
 
         execution_state
     }
@@ -254,200 +142,185 @@ impl ExecutionState {
         for (pre, mut post) in output_pre_states.into_iter().zip(output_post_states) {
             let pre_account_id = pre.account_id;
             let pre_is_authorized = pre.is_authorized;
-            let post_states_entry = self.post_states.entry(pre.account_id);
-            match &post_states_entry {
-                Entry::Occupied(occupied) => {
-                    #[expect(
-                        clippy::shadow_unrelated,
-                        reason = "Shadowing is intentional to use all fields"
-                    )]
-                    let AccountWithMetadata {
-                        account: pre_account,
-                        account_id: pre_account_id,
-                        is_authorized: pre_is_authorized,
-                    } = pre;
 
-                    // Ensure that new pre state is the same as known post state
-                    assert_eq!(
-                        occupied.get(),
-                        &pre_account,
-                        "Inconsistent pre state for account {pre_account_id}",
+            if let Some(existing) = self.post_states.get(&pre.account_id) {
+                #[expect(
+                    clippy::shadow_unrelated,
+                    reason = "Shadowing is intentional to use all fields"
+                )]
+                let AccountWithMetadata {
+                    account: pre_account,
+                    account_id: pre_account_id,
+                    is_authorized: pre_is_authorized,
+                } = pre;
+
+                assert_eq!(
+                    existing, &pre_account,
+                    "Inconsistent pre state for account {pre_account_id}",
+                );
+
+                let (previous_is_authorized, pre_state_position) = self
+                    .pre_states
+                    .iter()
+                    .enumerate()
+                    .find(|(_, acc)| acc.account_id == pre_account_id)
+                    .map_or_else(
+                        || panic!(
+                            "Pre state must exist in execution state for account {pre_account_id}",
+                        ),
+                        |(pos, acc)| (acc.is_authorized, pos),
                     );
 
-                    let (previous_is_authorized, pre_state_position) = self
-                        .pre_states
-                        .iter()
-                        .enumerate()
-                        .find(|(_, acc)| acc.account_id == pre_account_id)
-                        .map_or_else(
-                            || panic!(
-                                "Pre state must exist in execution state for account {pre_account_id}",
-                            ),
-                            |(pos, acc)| (acc.is_authorized, pos)
-                        );
+                let is_authorized = resolve_authorization_and_record_bindings(
+                    &mut self.pda_family_binding,
+                    &mut self.private_pda_bound_positions,
+                    &self.private_pda_npk_by_position,
+                    &mut self.authorized_accounts,
+                    pre_account_id,
+                    pre_state_position,
+                    caller_program_id,
+                    caller_pda_seeds,
+                    previous_is_authorized,
+                );
 
-                    let is_authorized = resolve_authorization_and_record_bindings(
-                        &mut self.pda_family_binding,
-                        &mut self.private_pda_bound_positions,
-                        &self.private_pda_npk_by_position,
-                        &mut self.authorized_accounts,
-                        pre_account_id,
-                        pre_state_position,
-                        caller_program_id,
-                        caller_pda_seeds,
-                        previous_is_authorized,
-                    );
-
-                    assert_eq!(
-                        pre_is_authorized, is_authorized,
-                        "Inconsistent authorization for account {pre_account_id}",
-                    );
-                }
-                Entry::Vacant(_) => {
-                    // Pre state for the initial call
-                    let pre_state_position = self.pre_states.len();
-                    let external_seed = match account_identities.get(pre_state_position) {
-                        Some(InputAccountIdentity::PrivatePdaInit {
-                            npk,
-                            identifier,
-                            seed: Some((seed, authority_program_id)),
-                            ..
-                        }) => {
-                            let expected = AccountId::for_private_pda(
-                                authority_program_id,
-                                seed,
-                                npk,
-                                *identifier,
-                            );
-                            assert_eq!(
-                                pre_account_id, expected,
-                                "External seed mismatch for PrivatePdaInit at position {pre_state_position}"
-                            );
-                            Some((*seed, *authority_program_id))
-                        }
-                        Some(InputAccountIdentity::PrivatePdaUpdate {
-                            nsk,
-                            identifier,
-                            seed: Some((seed, authority_program_id)),
-                            ..
-                        }) => {
-                            let npk = NullifierPublicKey::from(nsk);
-                            let expected = AccountId::for_private_pda(
-                                authority_program_id,
-                                seed,
-                                &npk,
-                                *identifier,
-                            );
-                            assert_eq!(
-                                pre_account_id, expected,
-                                "External seed mismatch for PrivatePdaUpdate at position {pre_state_position}"
-                            );
-                            Some((*seed, *authority_program_id))
-                        }
-                        _ => None,
-                    };
-                    // External seed is only consulted the first time the account is seen.
-                    // Subsequent calls need no re-check because the entry is already recorded on
-                    // private_pda_bound_positions.
-                    if let Some((seed, authority_program_id)) = external_seed {
-                        assert!(
-                            !pre.is_authorized,
-                            "Private PDA with externally-provided seed must not be authorized at position {pre_state_position}"
-                        );
-                        bind_private_pda_position(
-                            &mut self.private_pda_bound_positions,
-                            pre_state_position,
-                            authority_program_id,
-                            seed,
-                        );
-                        assert_family_binding(
-                            &mut self.pda_family_binding,
-                            authority_program_id,
-                            seed,
-                            pre_account_id,
-                        );
-                    }
-                    self.pre_states.push(pre);
-                }
+                assert_eq!(
+                    pre_is_authorized, is_authorized,
+                    "Inconsistent authorization for account {pre_account_id}",
+                );
+            } else {
+                let pre_state_position = self.pre_states.len();
+                resolve_external_seed(
+                    account_identities,
+                    pre_state_position,
+                    pre_account_id,
+                    pre.is_authorized,
+                    &mut self.private_pda_bound_positions,
+                    &mut self.pda_family_binding,
+                );
+                self.pre_states.push(pre);
             }
 
             if let Some(claim) = post.required_claim() {
-                // The invoked program can only claim accounts with default program id.
-                assert_eq!(
-                    post.account().program_owner,
-                    DEFAULT_PROGRAM_ID,
-                    "Cannot claim an initialized account {pre_account_id}"
+                self.process_claim(
+                    account_identities,
+                    &mut post,
+                    pre_account_id,
+                    pre_is_authorized,
+                    program_id,
+                    claim,
                 );
-
-                let pre_state_position = self
-                    .pre_states
-                    .iter()
-                    .position(|acc| acc.account_id == pre_account_id)
-                    .expect("Pre state must exist at this point");
-
-                let account_identity = &account_identities[pre_state_position];
-                if account_identity.is_public() {
-                    match claim {
-                        Claim::Authorized => {
-                            // Note: no need to check authorized pdas because we have already
-                            // checked consistency of authorization above.
-                            assert!(
-                                pre_is_authorized,
-                                "Cannot claim unauthorized account {pre_account_id}"
-                            );
-                        }
-                        Claim::Pda(seed) => {
-                            let pda = AccountId::for_public_pda(&program_id, &seed);
-                            assert_eq!(
-                                pre_account_id, pda,
-                                "Invalid PDA claim for account {pre_account_id} which does not match derived PDA {pda}"
-                            );
-                            assert_family_binding(
-                                &mut self.pda_family_binding,
-                                program_id,
-                                seed,
-                                pre_account_id,
-                            );
-                        }
-                    }
-                } else {
-                    // Private accounts: don't enforce the claim semantics. Unauthorized private
-                    // claiming is intentionally allowed
-                    match claim {
-                        Claim::Authorized => {}
-                        Claim::Pda(seed) => {
-                            let (npk, identifier) = self
-                                .private_pda_npk_by_position
-                                .get(&pre_state_position)
-                                .expect(
-                                    "private PDA pre_state must have an npk in the position map",
-                                );
-                            let pda =
-                                AccountId::for_private_pda(&program_id, &seed, npk, *identifier);
-                            assert_eq!(
-                                pre_account_id, pda,
-                                "Invalid private PDA claim for account {pre_account_id}"
-                            );
-                            bind_private_pda_position(
-                                &mut self.private_pda_bound_positions,
-                                pre_state_position,
-                                program_id,
-                                seed,
-                            );
-                            assert_family_binding(
-                                &mut self.pda_family_binding,
-                                program_id,
-                                seed,
-                                pre_account_id,
-                            );
-                        }
-                    }
-                }
-
-                post.account_mut().program_owner = program_id;
             }
 
-            post_states_entry.insert_entry(post.into_account());
+            self.post_states.insert(pre_account_id, post.into_account());
+        }
+    }
+
+    fn process_claim(
+        &mut self,
+        account_identities: &[InputAccountIdentity],
+        post: &mut AccountPostState,
+        pre_account_id: AccountId,
+        pre_is_authorized: bool,
+        program_id: ProgramId,
+        claim: Claim,
+    ) {
+        assert_eq!(
+            post.account().program_owner,
+            DEFAULT_PROGRAM_ID,
+            "Cannot claim an initialized account {pre_account_id}"
+        );
+
+        let pre_state_position = self
+            .pre_states
+            .iter()
+            .position(|acc| acc.account_id == pre_account_id)
+            .expect("Pre state must exist at this point");
+
+        let account_identity = &account_identities[pre_state_position];
+        if account_identity.is_public() {
+            match claim {
+                Claim::Authorized => {
+                    assert!(
+                        pre_is_authorized,
+                        "Cannot claim unauthorized account {pre_account_id}"
+                    );
+                }
+                Claim::Pda(seed) => {
+                    let pda = AccountId::for_public_pda(&program_id, &seed);
+                    assert_eq!(
+                        pre_account_id, pda,
+                        "Invalid PDA claim for account {pre_account_id} which does not match derived PDA {pda}"
+                    );
+                    assert_family_binding(
+                        &mut self.pda_family_binding,
+                        program_id,
+                        seed,
+                        pre_account_id,
+                    );
+                }
+            }
+        } else {
+            match claim {
+                Claim::Authorized => {}
+                Claim::Pda(seed) => {
+                    let (npk, identifier) = self
+                        .private_pda_npk_by_position
+                        .get(&pre_state_position)
+                        .expect("private PDA pre_state must have an npk in the position map");
+                    let pda = AccountId::for_private_pda(&program_id, &seed, npk, *identifier);
+                    assert_eq!(
+                        pre_account_id, pda,
+                        "Invalid private PDA claim for account {pre_account_id}"
+                    );
+                    bind_private_pda_position(
+                        &mut self.private_pda_bound_positions,
+                        pre_state_position,
+                        program_id,
+                        seed,
+                    );
+                    assert_family_binding(
+                        &mut self.pda_family_binding,
+                        program_id,
+                        seed,
+                        pre_account_id,
+                    );
+                }
+            }
+        }
+
+        post.account_mut().program_owner = program_id;
+    }
+
+    fn assert_all_pda_positions_bound(&self, account_identities: &[InputAccountIdentity]) {
+        for (pos, account_identity) in account_identities.iter().enumerate() {
+            if account_identity.is_private_pda() {
+                assert!(
+                    self.private_pda_bound_positions.contains_key(&pos),
+                    "private PDA pre_state at position {pos} has no proven (seed, npk) binding via Claim::Pda or caller pda_seeds"
+                );
+            }
+        }
+    }
+
+    fn assert_modified_accounts_claimed(&self) {
+        for (account_id, post) in self
+            .pre_states
+            .iter()
+            .filter(|a| a.account.program_owner == DEFAULT_PROGRAM_ID)
+            .map(|a| {
+                let post = self
+                    .post_states
+                    .get(&a.account_id)
+                    .expect("Post state must exist for pre state");
+                (a, post)
+            })
+            .filter(|(pre_default, post)| pre_default.account != **post)
+            .map(|(pre, post)| (pre.account_id, post))
+        {
+            assert_ne!(
+                post.program_owner, DEFAULT_PROGRAM_ID,
+                "Account {account_id} was modified but not claimed"
+            );
         }
     }
 
@@ -486,6 +359,150 @@ impl ExecutionState {
     }
 }
 
+fn verify_program_output(
+    chained_call: &ChainedCall,
+    caller_program_id: Option<ProgramId>,
+    program_output: &ProgramOutput,
+) {
+    assert_eq!(
+        chained_call.instruction_data, program_output.instruction_data,
+        "Mismatched instruction data between chained call and program output"
+    );
+
+    let program_output_words =
+        &to_vec(program_output).expect("program_output must be serializable");
+    env::verify(chained_call.program_id, program_output_words)
+        .unwrap_or_else(|_: Infallible| unreachable!("Infallible error is never constructed"));
+
+    assert_eq!(
+        program_output.self_program_id, chained_call.program_id,
+        "Program output self_program_id does not match chained call program_id"
+    );
+
+    assert_eq!(
+        program_output.caller_program_id, caller_program_id,
+        "Program output caller_program_id does not match actual caller"
+    );
+
+    if let Err(err) = validate_execution(
+        &program_output.pre_states,
+        &program_output.post_states,
+        chained_call.program_id,
+    ) {
+        panic!(
+            "Invalid program behavior in program {:?}: {err}",
+            chained_call.program_id
+        );
+    }
+}
+
+fn build_private_pda_npk_map(
+    account_identities: &[InputAccountIdentity],
+) -> HashMap<usize, (NullifierPublicKey, Identifier)> {
+    account_identities
+        .iter()
+        .enumerate()
+        .filter_map(|(pos, identity)| {
+            identity
+                .npk_if_private_pda()
+                .map(|(npk, identifier)| (pos, (npk, identifier)))
+        })
+        .collect()
+}
+
+fn intersect_validity_windows(
+    program_outputs: &[ProgramOutput],
+) -> (BlockValidityWindow, TimestampValidityWindow) {
+    let block_valid_from = program_outputs
+        .iter()
+        .filter_map(|output| output.block_validity_window.start())
+        .max();
+    let block_valid_until = program_outputs
+        .iter()
+        .filter_map(|output| output.block_validity_window.end())
+        .min();
+    let ts_valid_from = program_outputs
+        .iter()
+        .filter_map(|output| output.timestamp_validity_window.start())
+        .max();
+    let ts_valid_until = program_outputs
+        .iter()
+        .filter_map(|output| output.timestamp_validity_window.end())
+        .min();
+
+    let block_validity_window: BlockValidityWindow =
+        (block_valid_from, block_valid_until).try_into().expect(
+            "There should be non empty intersection in the program output block validity windows",
+        );
+    let timestamp_validity_window: TimestampValidityWindow = (ts_valid_from, ts_valid_until)
+        .try_into()
+        .expect(
+            "There should be non empty intersection in the program output timestamp validity windows",
+        );
+
+    (block_validity_window, timestamp_validity_window)
+}
+
+fn resolve_external_seed(
+    account_identities: &[InputAccountIdentity],
+    pre_state_position: usize,
+    pre_account_id: AccountId,
+    is_authorized: bool,
+    private_pda_bound_positions: &mut HashMap<usize, (ProgramId, PdaSeed)>,
+    pda_family_binding: &mut HashMap<(ProgramId, PdaSeed), AccountId>,
+) {
+    let external_seed = match account_identities.get(pre_state_position) {
+        Some(InputAccountIdentity::PrivatePdaInit {
+            npk,
+            identifier,
+            seed: Some((seed, authority_program_id)),
+            ..
+        }) => {
+            let expected = AccountId::for_private_pda(authority_program_id, seed, npk, *identifier);
+            assert_eq!(
+                pre_account_id, expected,
+                "External seed mismatch for PrivatePdaInit at position {pre_state_position}"
+            );
+            Some((*seed, *authority_program_id))
+        }
+        Some(InputAccountIdentity::PrivatePdaUpdate {
+            nsk,
+            identifier,
+            seed: Some((seed, authority_program_id)),
+            ..
+        }) => {
+            let npk = NullifierPublicKey::from(nsk);
+            let expected =
+                AccountId::for_private_pda(authority_program_id, seed, &npk, *identifier);
+            assert_eq!(
+                pre_account_id, expected,
+                "External seed mismatch for PrivatePdaUpdate at position {pre_state_position}"
+            );
+            Some((*seed, *authority_program_id))
+        }
+        _ => None,
+    };
+
+    if let Some((seed, authority_program_id)) = external_seed {
+        assert!(
+            !is_authorized,
+            "Private PDA with externally-provided seed must not be authorized at position {pre_state_position}"
+        );
+        bind_private_pda_position(
+            private_pda_bound_positions,
+            pre_state_position,
+            authority_program_id,
+            seed,
+        );
+        assert_family_binding(
+            pda_family_binding,
+            authority_program_id,
+            seed,
+            pre_account_id,
+        );
+    }
+}
+
 /// Record or re-verify the `(program_id, seed) → account_id` family binding for the
 /// transaction. Any claim or caller-seed authorization that resolves a `pre_state` under
 /// `(program_id, seed)` must agree with every prior resolution of the same pair; otherwise a

lee/privacy_preserving_circuit/src/output.rs

@@ -1,13 +1,244 @@
 use lee_core::{
-    Commitment, CommitmentSetDigest, DUMMY_COMMITMENT_HASH, EncryptedAccountData, EncryptionScheme,
-    EphemeralPublicKey, InputAccountIdentity, MembershipProof, Nullifier, NullifierPublicKey,
-    NullifierSecretKey, PrivacyPreservingCircuitOutput, PrivateAccountKind, SharedSecretKey,
+    Commitment, CommitmentSetDigest, EncryptedAccountData, EncryptionScheme, EphemeralPublicKey,
+    InputAccountIdentity, MembershipProof, Nullifier, NullifierPublicKey, NullifierSecretKey,
+    PrivacyPreservingCircuitOutput, PrivateAccountKind, SharedSecretKey,
     account::{Account, AccountId, Nonce},
     compute_digest_for_path,
 };
 
 use crate::execution_state::ExecutionState;
 
+struct PrivateOutputHandler<'ctx> {
+    output: &'ctx mut PrivacyPreservingCircuitOutput,
+    output_index: &'ctx mut u32,
+    pre_state: &'ctx lee_core::account::AccountWithMetadata,
+    post_state: Account,
+    epk: &'ctx EphemeralPublicKey,
+    view_tag: u8,
+    ssk: &'ctx SharedSecretKey,
+    identifier: u128,
+}
+
+impl PrivateOutputHandler<'_> {
+    fn authorized_init(self, nsk: &NullifierSecretKey, commitment_root: &CommitmentSetDigest) {
+        let npk = NullifierPublicKey::from(nsk);
+        let account_id =
+            derive_and_verify_account_id(&npk, self.identifier, self.pre_state.account_id);
+
+        assert!(
+            self.pre_state.is_authorized,
+            "Pre-state not authorized for authenticated private account"
+        );
+        assert_eq!(
+            self.pre_state.account,
+            Account::default(),
+            "Found new private account with non default values"
+        );
+
+        let (new_nullifier, new_nonce) = init_nullifier_and_nonce(&account_id, commitment_root);
+        let kind = PrivateAccountKind::Regular(self.identifier);
+
+        self.emit_private_output(&account_id, &kind, new_nullifier, new_nonce);
+    }
+
+    fn authorized_update(self, nsk: &NullifierSecretKey, membership_proof: &MembershipProof) {
+        let npk = NullifierPublicKey::from(nsk);
+        let account_id =
+            derive_and_verify_account_id(&npk, self.identifier, self.pre_state.account_id);
+
+        assert!(
+            self.pre_state.is_authorized,
+            "Pre-state not authorized for authenticated private account"
+        );
+
+        let new_nullifier = compute_update_nullifier_and_set_digest(
+            membership_proof,
+            &self.pre_state.account,
+            &account_id,
+            nsk,
+        );
+        let new_nonce = self
+            .pre_state
+            .account
+            .nonce
+            .private_account_nonce_increment(nsk);
+        let kind = PrivateAccountKind::Regular(self.identifier);
+
+        self.emit_private_output(&account_id, &kind, new_nullifier, new_nonce);
+    }
+
+    fn unauthorized(self, npk: &NullifierPublicKey, commitment_root: &CommitmentSetDigest) {
+        let account_id =
+            derive_and_verify_account_id(npk, self.identifier, self.pre_state.account_id);
+
+        assert_eq!(
+            self.pre_state.account,
+            Account::default(),
+            "Found new private account with non default values",
+        );
+        assert!(
+            !self.pre_state.is_authorized,
+            "Found new private account marked as authorized."
+        );
+
+        let (new_nullifier, new_nonce) = init_nullifier_and_nonce(&account_id, commitment_root);
+        let kind = PrivateAccountKind::Regular(self.identifier);
+
+        self.emit_private_output(&account_id, &kind, new_nullifier, new_nonce);
+    }
+
+    fn pda_init(
+        self,
+        commitment_root: &CommitmentSetDigest,
+        pos: usize,
+        pda_seed_by_position: &std::collections::HashMap<
+            usize,
+            (lee_core::program::ProgramId, lee_core::program::PdaSeed),
+        >,
+    ) {
+        // The npk-to-account_id binding is established upstream in
+        // `validate_and_sync_states` via `Claim::Pda(seed)` or a caller `pda_seeds`
+        // match. Here we only enforce the init pre-conditions.
+        assert!(
+            !self.pre_state.is_authorized,
+            "PrivatePdaInit requires unauthorized pre_state"
+        );
+        assert_eq!(
+            self.pre_state.account,
+            Account::default(),
+            "New private PDA must be default"
+        );
+
+        let (new_nullifier, new_nonce) =
+            init_nullifier_and_nonce(&self.pre_state.account_id, commitment_root);
+
+        let account_id = self.pre_state.account_id;
+        let (authority_program_id, seed) = pda_seed_by_position
+            .get(&pos)
+            .expect("PrivatePdaInit position must be in pda_seed_by_position");
+        let kind = PrivateAccountKind::Pda {
+            program_id: *authority_program_id,
+            seed: *seed,
+            identifier: self.identifier,
+        };
+
+        self.emit_private_output(&account_id, &kind, new_nullifier, new_nonce);
+    }
+
+    fn pda_update(
+        self,
+        nsk: &NullifierSecretKey,
+        membership_proof: &MembershipProof,
+        external_seed: Option<&(lee_core::program::PdaSeed, lee_core::program::ProgramId)>,
+        pos: usize,
+        pda_seed_by_position: &std::collections::HashMap<
+            usize,
+            (lee_core::program::ProgramId, lee_core::program::PdaSeed),
+        >,
+    ) {
+        // With an external seed the binding comes from the circuit input and the
+        // pre_state is intentionally unauthorized; without one the binding comes from
+        // a Claim or caller pda_seeds, so the pre_state must already be authorized.
+        assert!(
+            self.pre_state.is_authorized ^ external_seed.is_some(),
+            "PrivatePdaUpdate requires authorized pre_state or external seed"
+        );
+
+        let new_nullifier = compute_update_nullifier_and_set_digest(
+            membership_proof,
+            &self.pre_state.account,
+            &self.pre_state.account_id,
+            nsk,
+        );
+        let new_nonce = self
+            .pre_state
+            .account
+            .nonce
+            .private_account_nonce_increment(nsk);
+
+        let account_id = self.pre_state.account_id;
+        let (authority_program_id, seed) = pda_seed_by_position
+            .get(&pos)
+            .expect("PrivatePdaUpdate position must be in pda_seed_by_position");
+        let kind = PrivateAccountKind::Pda {
+            program_id: *authority_program_id,
+            seed: *seed,
+            identifier: self.identifier,
+        };
+
+        self.emit_private_output(&account_id, &kind, new_nullifier, new_nonce);
+    }
+
+    fn emit_private_output(
+        self,
+        account_id: &AccountId,
+        kind: &PrivateAccountKind,
+        new_nullifier: (Nullifier, CommitmentSetDigest),
+        new_nonce: Nonce,
+    ) {
+        self.output.new_nullifiers.push(new_nullifier);
+
+        let mut post_with_updated_nonce = self.post_state;
+        post_with_updated_nonce.nonce = new_nonce;
+
+        let commitment_post = Commitment::new(account_id, &post_with_updated_nonce);
+        let encrypted_account = EncryptionScheme::encrypt(
+            &post_with_updated_nonce,
+            kind,
+            self.ssk,
+            &commitment_post,
+            *self.output_index,
+        );
+
+        self.output.new_commitments.push(commitment_post);
+        self.output
+            .encrypted_private_post_states
+            .push(EncryptedAccountData {
+                ciphertext: encrypted_account,
+                epk: self.epk.clone(),
+                view_tag: self.view_tag,
+            });
+        *self.output_index = self
+            .output_index
+            .checked_add(1)
+            .unwrap_or_else(|| panic!("Too many private accounts, output index overflow"));
+    }
+}
+
+fn init_nullifier_and_nonce(
+    account_id: &AccountId,
+    commitment_root: &CommitmentSetDigest,
+) -> ((Nullifier, CommitmentSetDigest), Nonce) {
+    let nullifier = (
+        Nullifier::for_account_initialization(account_id),
+        *commitment_root,
+    );
+    let nonce = Nonce::private_account_nonce_init(account_id);
+    (nullifier, nonce)
+}
+
+fn derive_and_verify_account_id(
+    npk: &NullifierPublicKey,
+    identifier: u128,
+    pre_state_account_id: AccountId,
+) -> AccountId {
+    let account_id = AccountId::for_regular_private_account(npk, identifier);
+    assert_eq!(account_id, pre_state_account_id, "AccountId mismatch");
+    account_id
+}
+
+fn compute_update_nullifier_and_set_digest(
+    membership_proof: &MembershipProof,
+    pre_account: &Account,
+    account_id: &AccountId,
+    nsk: &NullifierSecretKey,
+) -> (Nullifier, CommitmentSetDigest) {
+    let commitment_pre = Commitment::new(account_id, pre_account);
+    let set_digest = compute_digest_for_path(&commitment_pre, membership_proof);
+    let nullifier = Nullifier::for_account_update(&commitment_pre, nsk);
+    (nullifier, set_digest)
+}
+
 pub fn compute_circuit_output(
     execution_state: ExecutionState,
     account_identities: &[InputAccountIdentity],
@@ -45,40 +276,18 @@ pub fn compute_circuit_output(
                 ssk,
                 nsk,
                 identifier,
-            } => {
-                let npk = NullifierPublicKey::from(nsk);
-                let account_id = AccountId::for_regular_private_account(&npk, *identifier);
-
-                assert_eq!(account_id, pre_state.account_id, "AccountId mismatch");
-                assert!(
-                    pre_state.is_authorized,
-                    "Pre-state not authorized for authenticated private account"
-                );
-                assert_eq!(
-                    pre_state.account,
-                    Account::default(),
-                    "Found new private account with non default values"
-                );
-
-                let new_nullifier = (
-                    Nullifier::for_account_initialization(&account_id),
-                    DUMMY_COMMITMENT_HASH,
-                );
-                let new_nonce = Nonce::private_account_nonce_init(&account_id);
-
-                emit_private_output(
-                    &mut output,
-                    &mut output_index,
-                    post_state,
-                    &account_id,
-                    &PrivateAccountKind::Regular(*identifier),
-                    ssk,
-                    epk,
-                    *view_tag,
-                    new_nullifier,
-                    new_nonce,
-                );
+                commitment_root,
+            } => PrivateOutputHandler {
+                output: &mut output,
+                output_index: &mut output_index,
+                pre_state: &pre_state,
+                post_state,
+                epk,
+                view_tag: *view_tag,
+                ssk,
+                identifier: *identifier,
             }
+            .authorized_init(nsk, commitment_root),
             InputAccountIdentity::PrivateAuthorizedUpdate {
                 epk,
                 view_tag,
@@ -86,127 +295,54 @@ pub fn compute_circuit_output(
                 nsk,
                 membership_proof,
                 identifier,
-            } => {
-                let npk = NullifierPublicKey::from(nsk);
-                let account_id = AccountId::for_regular_private_account(&npk, *identifier);
-
-                assert_eq!(account_id, pre_state.account_id, "AccountId mismatch");
-                assert!(
-                    pre_state.is_authorized,
-                    "Pre-state not authorized for authenticated private account"
-                );
-
-                let new_nullifier = compute_update_nullifier_and_set_digest(
-                    membership_proof,
-                    &pre_state.account,
-                    &account_id,
-                    nsk,
-                );
-                let new_nonce = pre_state.account.nonce.private_account_nonce_increment(nsk);
-
-                emit_private_output(
-                    &mut output,
-                    &mut output_index,
-                    post_state,
-                    &account_id,
-                    &PrivateAccountKind::Regular(*identifier),
-                    ssk,
-                    epk,
-                    *view_tag,
-                    new_nullifier,
-                    new_nonce,
-                );
+            } => PrivateOutputHandler {
+                output: &mut output,
+                output_index: &mut output_index,
+                pre_state: &pre_state,
+                post_state,
+                epk,
+                view_tag: *view_tag,
+                ssk,
+                identifier: *identifier,
             }
+            .authorized_update(nsk, membership_proof),
             InputAccountIdentity::PrivateUnauthorized {
                 epk,
                 view_tag,
                 npk,
                 ssk,
                 identifier,
-            } => {
-                let account_id = AccountId::for_regular_private_account(npk, *identifier);
-
-                assert_eq!(account_id, pre_state.account_id, "AccountId mismatch");
-                assert_eq!(
-                    pre_state.account,
-                    Account::default(),
-                    "Found new private account with non default values",
-                );
-                assert!(
-                    !pre_state.is_authorized,
-                    "Found new private account marked as authorized."
-                );
-
-                let new_nullifier = (
-                    Nullifier::for_account_initialization(&account_id),
-                    DUMMY_COMMITMENT_HASH,
-                );
-                let new_nonce = Nonce::private_account_nonce_init(&account_id);
-
-                emit_private_output(
-                    &mut output,
-                    &mut output_index,
-                    post_state,
-                    &account_id,
-                    &PrivateAccountKind::Regular(*identifier),
-                    ssk,
-                    epk,
-                    *view_tag,
-                    new_nullifier,
-                    new_nonce,
-                );
+                commitment_root,
+            } => PrivateOutputHandler {
+                output: &mut output,
+                output_index: &mut output_index,
+                pre_state: &pre_state,
+                post_state,
+                epk,
+                view_tag: *view_tag,
+                ssk,
+                identifier: *identifier,
             }
+            .unauthorized(npk, commitment_root),
             InputAccountIdentity::PrivatePdaInit {
                 epk,
                 view_tag,
                 npk: _,
                 ssk,
                 identifier,
+                commitment_root,
                 seed: _,
-            } => {
-                // The npk-to-account_id binding is established upstream in
-                // `validate_and_sync_states` via `Claim::Pda(seed)` or a caller `pda_seeds`
-                // match. Here we only enforce the init pre-conditions. The supplied npk on
-                // the variant has been recorded into `private_pda_npk_by_position` and used
-                // for the binding check; we use `pre_state.account_id` directly for nullifier
-                // and commitment derivation.
-                assert!(
-                    !pre_state.is_authorized,
-                    "PrivatePdaInit requires unauthorized pre_state"
-                );
-                assert_eq!(
-                    pre_state.account,
-                    Account::default(),
-                    "New private PDA must be default"
-                );
-
-                let new_nullifier = (
-                    Nullifier::for_account_initialization(&pre_state.account_id),
-                    DUMMY_COMMITMENT_HASH,
-                );
-                let new_nonce = Nonce::private_account_nonce_init(&pre_state.account_id);
-
-                let account_id = pre_state.account_id;
-                let (authority_program_id, seed) = pda_seed_by_position
-                    .get(&pos)
-                    .expect("PrivatePdaInit position must be in pda_seed_by_position");
-                emit_private_output(
-                    &mut output,
-                    &mut output_index,
-                    post_state,
-                    &account_id,
-                    &PrivateAccountKind::Pda {
-                        program_id: *authority_program_id,
-                        seed: *seed,
-                        identifier: *identifier,
-                    },
-                    ssk,
-                    epk,
-                    *view_tag,
-                    new_nullifier,
-                    new_nonce,
-                );
+            } => PrivateOutputHandler {
+                output: &mut output,
+                output_index: &mut output_index,
+                pre_state: &pre_state,
+                post_state,
+                epk,
+                view_tag: *view_tag,
+                ssk,
+                identifier: *identifier,
             }
+            .pda_init(commitment_root, pos, &pda_seed_by_position),
             InputAccountIdentity::PrivatePdaUpdate {
                 epk,
                 view_tag,
@@ -215,103 +351,25 @@ pub fn compute_circuit_output(
                 membership_proof,
                 identifier,
                 seed: external_seed,
-            } => {
-                // With an external seed the binding comes from the circuit input and the
-                // pre_state is intentionally unauthorized; without one the binding comes from
-                // a Claim or caller pda_seeds, so the pre_state must already be authorized.
-                // When `external_seed` is `Some`, execution_state already asserted
-                // `!pre_state.is_authorized`.
-                assert!(
-                    pre_state.is_authorized ^ external_seed.is_some(),
-                    "PrivatePdaUpdate requires authorized pre_state or external seed"
-                );
-
-                let new_nullifier = compute_update_nullifier_and_set_digest(
-                    membership_proof,
-                    &pre_state.account,
-                    &pre_state.account_id,
-                    nsk,
-                );
-                let new_nonce = pre_state.account.nonce.private_account_nonce_increment(nsk);
-
-                let account_id = pre_state.account_id;
-                let (authority_program_id, seed) = pda_seed_by_position
-                    .get(&pos)
-                    .expect("PrivatePdaUpdate position must be in pda_seed_by_position");
-                emit_private_output(
-                    &mut output,
-                    &mut output_index,
-                    post_state,
-                    &account_id,
-                    &PrivateAccountKind::Pda {
-                        program_id: *authority_program_id,
-                        seed: *seed,
-                        identifier: *identifier,
-                    },
-                    ssk,
-                    epk,
-                    *view_tag,
-                    new_nullifier,
-                    new_nonce,
-                );
+            } => PrivateOutputHandler {
+                output: &mut output,
+                output_index: &mut output_index,
+                pre_state: &pre_state,
+                post_state,
+                epk,
+                view_tag: *view_tag,
+                ssk,
+                identifier: *identifier,
             }
+            .pda_update(
+                nsk,
+                membership_proof,
+                external_seed.as_ref(),
+                pos,
+                &pda_seed_by_position,
+            ),
         }
     }
 
     output
 }
-
-#[expect(
-    clippy::too_many_arguments,
-    reason = "Inputs are distinct concerns from the variant arms; bundling would be artificial"
-)]
-fn emit_private_output(
-    output: &mut PrivacyPreservingCircuitOutput,
-    output_index: &mut u32,
-    post_state: Account,
-    account_id: &AccountId,
-    kind: &PrivateAccountKind,
-    shared_secret: &SharedSecretKey,
-    epk: &EphemeralPublicKey,
-    view_tag: u8,
-    new_nullifier: (Nullifier, CommitmentSetDigest),
-    new_nonce: Nonce,
-) {
-    output.new_nullifiers.push(new_nullifier);
-
-    let mut post_with_updated_nonce = post_state;
-    post_with_updated_nonce.nonce = new_nonce;
-
-    let commitment_post = Commitment::new(account_id, &post_with_updated_nonce);
-    let encrypted_account = EncryptionScheme::encrypt(
-        &post_with_updated_nonce,
-        kind,
-        shared_secret,
-        &commitment_post,
-        *output_index,
-    );
-
-    output.new_commitments.push(commitment_post);
-    output
-        .encrypted_private_post_states
-        .push(EncryptedAccountData {
-            ciphertext: encrypted_account,
-            epk: epk.clone(),
-            view_tag,
-        });
-    *output_index = output_index
-        .checked_add(1)
-        .unwrap_or_else(|| panic!("Too many private accounts, output index overflow"));
-}
-
-fn compute_update_nullifier_and_set_digest(
-    membership_proof: &MembershipProof,
-    pre_account: &Account,
-    account_id: &AccountId,
-    nsk: &NullifierSecretKey,
-) -> (Nullifier, CommitmentSetDigest) {
-    let commitment_pre = Commitment::new(account_id, pre_account);
-    let set_digest = compute_digest_for_path(&commitment_pre, membership_proof);
-    let nullifier = Nullifier::for_account_update(&commitment_pre, nsk);
-    (nullifier, set_digest)
-}

lee/state_machine/core/src/circuit_io.rs

@@ -38,6 +38,7 @@ pub enum InputAccountIdentity {
         ssk: SharedSecretKey,
         nsk: NullifierSecretKey,
         identifier: Identifier,
+        commitment_root: CommitmentSetDigest,
     },
     /// Update of an authorized standalone private account: existing on-chain commitment, with
     /// membership proof.
@@ -57,6 +58,7 @@ pub enum InputAccountIdentity {
         npk: NullifierPublicKey,
         ssk: SharedSecretKey,
         identifier: Identifier,
+        commitment_root: CommitmentSetDigest,
     },
     /// Init of a private PDA, unauthorized. The npk-to-account_id binding is proven upstream
     /// via `Claim::Pda(seed)` or a caller's `pda_seeds` match. The identifier diversifies the
@@ -68,6 +70,7 @@ pub enum InputAccountIdentity {
         npk: NullifierPublicKey,
         ssk: SharedSecretKey,
         identifier: Identifier,
+        commitment_root: CommitmentSetDigest,
         /// When `Some((seed, authority_program_id))`, the circuit binds this position via the
         /// external derivation check
         /// `AccountId::for_private_pda(authority_program_id, seed, npk, identifier) ==

lee/state_machine/core/src/encoding.rs

@@ -7,7 +7,7 @@ use std::io::Read as _;
 #[cfg(feature = "host")]
 use crate::Nullifier;
 #[cfg(feature = "host")]
-use crate::encryption::EphemeralPublicKey;
+use crate::encryption::{EphemeralPublicKey, ML_KEM_768_CIPHERTEXT_LEN};
 #[cfg(feature = "host")]
 use crate::error::LeeCoreError;
 use crate::{
@@ -168,7 +168,7 @@ impl EphemeralPublicKey {
     /// Deserializes an ML-KEM-768 ciphertext from a cursor.
     /// Reads exactly 1088 bytes — the fixed ciphertext size for ML-KEM-768.
     pub fn from_cursor(cursor: &mut Cursor<&[u8]>) -> Result<Self, LeeCoreError> {
-        let mut value = vec![0_u8; 1088];
+        let mut value = vec![0_u8; ML_KEM_768_CIPHERTEXT_LEN];
         cursor.read_exact(&mut value)?;
         Ok(Self(value))
     }

lee/state_machine/core/src/encryption/mod.rs

@@ -12,13 +12,16 @@ use crate::{Commitment, account::Account, program::PrivateAccountKind};
 #[cfg(feature = "host")]
 pub mod shared_key_derivation;
 
+/// Length in bytes of an ML-KEM-768 ciphertext (the `EphemeralPublicKey` payload).
+pub const ML_KEM_768_CIPHERTEXT_LEN: usize = 1088;
+
 pub type Scalar = [u8; 32];
 
 #[derive(Serialize, Deserialize, Clone, Copy)]
 pub struct SharedSecretKey(pub [u8; 32]);
 
 /// The ML-KEM-768 ciphertext produced during encapsulation; transmitted on-wire in place of the
-/// former ECDH ephemeral public key. Always 1088 bytes for ML-KEM-768.
+/// former ECDH ephemeral public key. Always `ML_KEM_768_CIPHERTEXT_LEN` (1088) bytes.
 #[derive(Serialize, Deserialize, Clone, Debug, PartialEq, Eq, BorshSerialize, BorshDeserialize)]
 pub struct EphemeralPublicKey(pub Vec<u8>);
 

lee/state_machine/core/src/encryption/shared_key_derivation.rs

@@ -64,7 +64,7 @@ impl SharedSecretKey {
         let ek_bytes: ml_kem::kem::Key<ml_kem::EncapsulationKey768> =
             ek.0.as_slice()
                 .try_into()
-                .expect("MlKem768EncapsulationKey must be 1184 bytes");
+                .expect("MlKem768EncapsulationKey must be MlKem768EncapsulationKey::LEN bytes");
         let ek_obj = ml_kem::EncapsulationKey768::new(&ek_bytes).expect(
             "MlKem768EncapsulationKey bytes must encode a valid ML-KEM-768 encapsulation key",
         );
@@ -104,7 +104,7 @@ impl SharedSecretKey {
         let ek_bytes: ml_kem::kem::Key<ml_kem::EncapsulationKey768> =
             ek.0.as_slice()
                 .try_into()
-                .expect("MlKem768EncapsulationKey must be 1184 bytes");
+                .expect("MlKem768EncapsulationKey must be MlKem768EncapsulationKey::LEN bytes");
         let ek_obj = ml_kem::EncapsulationKey768::new(&ek_bytes).expect(
             "MlKem768EncapsulationKey bytes must encode a valid ML-KEM-768 encapsulation key",
         );
@@ -118,8 +118,9 @@ impl SharedSecretKey {
 
     /// Receiver: decapsulate the shared secret from a KEM ciphertext.
     ///
-    /// Returns `None` if the `EphemeralPublicKey` is not exactly 1088 bytes — callers on
-    /// the wallet scan path should skip the output rather than panic on malformed chain data.
+    /// Returns `None` if the `EphemeralPublicKey` is not exactly [`ML_KEM_768_CIPHERTEXT_LEN`]
+    /// bytes — callers on the wallet scan path should skip the output rather than panic on
+    /// malformed chain data.
     ///
     /// `d` and `z` are the two 32-byte halves of the FIPS 203 `ViewingSecretKey` seed.
     #[must_use]
@@ -146,6 +147,7 @@ mod tests {
     use ml_kem::KeyExport as _;
 
     use super::*;
+    use crate::ML_KEM_768_CIPHERTEXT_LEN;
 
     #[test]
     fn encapsulate_decapsulate_round_trip() {
@@ -164,11 +166,15 @@ mod tests {
         let receiver_ss = SharedSecretKey::decapsulate(&epk, &d, &z).unwrap();
 
         assert_eq!(sender_ss.0, receiver_ss.0, "shared secrets must match");
-        assert_eq!(epk.0.len(), 1088, "ML-KEM-768 ciphertext is 1088 bytes");
+        assert_eq!(
+            epk.0.len(),
+            ML_KEM_768_CIPHERTEXT_LEN,
+            "ML-KEM-768 ciphertext length"
+        );
         assert_eq!(
             ek.0.len(),
-            1184,
-            "ML-KEM-768 encapsulation key is 1184 bytes"
+            MlKem768EncapsulationKey::LEN,
+            "ML-KEM-768 encapsulation key length"
         );
     }
 
@@ -177,15 +183,15 @@ mod tests {
         let d = [1_u8; 32];
         let z = [2_u8; 32];
 
-        // Too short — 100 bytes instead of 1088.
+        // Too short — 100 bytes instead of ML_KEM_768_CIPHERTEXT_LEN.
         let short_epk = EphemeralPublicKey(vec![42_u8; 100]);
         assert!(
             SharedSecretKey::decapsulate(&short_epk, &d, &z).is_none(),
             "short EphemeralPublicKey must return None"
         );
 
-        // Too long — 1089 bytes instead of 1088.
-        let long_epk = EphemeralPublicKey(vec![42_u8; 1089]);
+        // Too long — ML_KEM_768_CIPHERTEXT_LEN + 1.
+        let long_epk = EphemeralPublicKey(vec![42_u8; ML_KEM_768_CIPHERTEXT_LEN + 1]);
         assert!(
             SharedSecretKey::decapsulate(&long_epk, &d, &z).is_none(),
             "long EphemeralPublicKey must return None"

lee/state_machine/core/src/lib.rs

@@ -11,7 +11,8 @@ pub use commitment::{
     compute_digest_for_path,
 };
 pub use encryption::{
-    EncryptedAccountData, EncryptionScheme, EphemeralPublicKey, SharedSecretKey, ViewTag,
+    EncryptedAccountData, EncryptionScheme, EphemeralPublicKey, ML_KEM_768_CIPHERTEXT_LEN,
+    SharedSecretKey, ViewTag,
 };
 pub use nullifier::{Identifier, Nullifier, NullifierPublicKey, NullifierSecretKey};
 pub use program::PrivateAccountKind;

lee/state_machine/core/src/program/mod.rs

@@ -230,7 +230,7 @@ impl ChainedCall {
 /// Represents the final state of an `Account` after a program execution.
 ///
 /// A post state may optionally request that the executing program
-/// becomes the owner of the account (a “claim”). This is used to signal
+/// becomes the owner of the account (a "claim"). This is used to signal
 /// that the program intends to take ownership of the account.
 #[derive(Debug, Clone, Serialize, Deserialize)]
 #[cfg_attr(any(feature = "host", test), derive(PartialEq, Eq))]
@@ -766,338 +766,4 @@ fn validate_uniqueness_of_account_ids(pre_states: &[AccountWithMetadata]) -> boo
 }
 
 #[cfg(test)]
-mod tests {
-    use super::*;
-
-    #[test]
-    fn validity_window_unbounded_accepts_any_value() {
-        let w: ValidityWindow<u64> = ValidityWindow::new_unbounded();
-        assert!(w.is_valid_for(0));
-        assert!(w.is_valid_for(u64::MAX));
-    }
-
-    #[test]
-    fn validity_window_bounded_range_includes_from_excludes_to() {
-        let w: ValidityWindow<u64> = (Some(5), Some(10)).try_into().unwrap();
-        assert!(!w.is_valid_for(4));
-        assert!(w.is_valid_for(5));
-        assert!(w.is_valid_for(9));
-        assert!(!w.is_valid_for(10));
-    }
-
-    #[test]
-    fn validity_window_only_from_bound() {
-        let w: ValidityWindow<u64> = (Some(5), None).try_into().unwrap();
-        assert!(!w.is_valid_for(4));
-        assert!(w.is_valid_for(5));
-        assert!(w.is_valid_for(u64::MAX));
-    }
-
-    #[test]
-    fn validity_window_only_to_bound() {
-        let w: ValidityWindow<u64> = (None, Some(5)).try_into().unwrap();
-        assert!(w.is_valid_for(0));
-        assert!(w.is_valid_for(4));
-        assert!(!w.is_valid_for(5));
-    }
-
-    #[test]
-    fn validity_window_adjacent_bounds_are_invalid() {
-        // [5, 5) is an empty range — from == to
-        assert!(ValidityWindow::<u64>::try_from((Some(5), Some(5))).is_err());
-    }
-
-    #[test]
-    fn validity_window_inverted_bounds_are_invalid() {
-        assert!(ValidityWindow::<u64>::try_from((Some(10), Some(5))).is_err());
-    }
-
-    #[test]
-    fn validity_window_getters_match_construction() {
-        let w: ValidityWindow<u64> = (Some(3), Some(7)).try_into().unwrap();
-        assert_eq!(w.start(), Some(3));
-        assert_eq!(w.end(), Some(7));
-    }
-
-    #[test]
-    fn validity_window_getters_for_unbounded() {
-        let w: ValidityWindow<u64> = ValidityWindow::new_unbounded();
-        assert_eq!(w.start(), None);
-        assert_eq!(w.end(), None);
-    }
-
-    #[test]
-    fn validity_window_from_range() {
-        let w: ValidityWindow<u64> = ValidityWindow::try_from(5_u64..10).unwrap();
-        assert_eq!(w.start(), Some(5));
-        assert_eq!(w.end(), Some(10));
-    }
-
-    #[test]
-    fn validity_window_from_range_empty_is_invalid() {
-        assert!(ValidityWindow::<u64>::try_from(5_u64..5).is_err());
-    }
-
-    #[test]
-    fn validity_window_from_range_inverted_is_invalid() {
-        let from = 10_u64;
-        let to = 5_u64;
-        assert!(ValidityWindow::<u64>::try_from(from..to).is_err());
-    }
-
-    #[test]
-    fn validity_window_from_range_from() {
-        let w: ValidityWindow<u64> = (5_u64..).into();
-        assert_eq!(w.start(), Some(5));
-        assert_eq!(w.end(), None);
-    }
-
-    #[test]
-    fn validity_window_from_range_to() {
-        let w: ValidityWindow<u64> = (..10_u64).into();
-        assert_eq!(w.start(), None);
-        assert_eq!(w.end(), Some(10));
-    }
-
-    #[test]
-    fn validity_window_from_range_full() {
-        let w: ValidityWindow<u64> = (..).into();
-        assert_eq!(w.start(), None);
-        assert_eq!(w.end(), None);
-    }
-
-    #[test]
-    fn program_output_try_with_block_validity_window_range() {
-        let output = ProgramOutput::new(DEFAULT_PROGRAM_ID, None, vec![], vec![], vec![])
-            .try_with_block_validity_window(10_u64..100)
-            .unwrap();
-        assert_eq!(output.block_validity_window.start(), Some(10));
-        assert_eq!(output.block_validity_window.end(), Some(100));
-    }
-
-    #[test]
-    fn program_output_with_block_validity_window_range_from() {
-        let output = ProgramOutput::new(DEFAULT_PROGRAM_ID, None, vec![], vec![], vec![])
-            .with_block_validity_window(10_u64..);
-        assert_eq!(output.block_validity_window.start(), Some(10));
-        assert_eq!(output.block_validity_window.end(), None);
-    }
-
-    #[test]
-    fn program_output_with_block_validity_window_range_to() {
-        let output = ProgramOutput::new(DEFAULT_PROGRAM_ID, None, vec![], vec![], vec![])
-            .with_block_validity_window(..100_u64);
-        assert_eq!(output.block_validity_window.start(), None);
-        assert_eq!(output.block_validity_window.end(), Some(100));
-    }
-
-    #[test]
-    fn program_output_try_with_block_validity_window_empty_range_fails() {
-        let result = ProgramOutput::new(DEFAULT_PROGRAM_ID, None, vec![], vec![], vec![])
-            .try_with_block_validity_window(5_u64..5);
-        assert!(result.is_err());
-    }
-
-    #[test]
-    fn post_state_new_with_claim_constructor() {
-        let account = Account {
-            program_owner: [1, 2, 3, 4, 5, 6, 7, 8],
-            balance: 1337,
-            data: vec![0xde, 0xad, 0xbe, 0xef].try_into().unwrap(),
-            nonce: 10_u128.into(),
-        };
-
-        let account_post_state = AccountPostState::new_claimed(account.clone(), Claim::Authorized);
-
-        assert_eq!(account, account_post_state.account);
-        assert_eq!(account_post_state.required_claim(), Some(Claim::Authorized));
-    }
-
-    #[test]
-    fn post_state_new_without_claim_constructor() {
-        let account = Account {
-            program_owner: [1, 2, 3, 4, 5, 6, 7, 8],
-            balance: 1337,
-            data: vec![0xde, 0xad, 0xbe, 0xef].try_into().unwrap(),
-            nonce: 10_u128.into(),
-        };
-
-        let account_post_state = AccountPostState::new(account.clone());
-
-        assert_eq!(account, account_post_state.account);
-        assert!(account_post_state.required_claim().is_none());
-    }
-
-    #[test]
-    fn post_state_account_getter() {
-        let mut account = Account {
-            program_owner: [1, 2, 3, 4, 5, 6, 7, 8],
-            balance: 1337,
-            data: vec![0xde, 0xad, 0xbe, 0xef].try_into().unwrap(),
-            nonce: 10_u128.into(),
-        };
-
-        let mut account_post_state = AccountPostState::new(account.clone());
-
-        assert_eq!(account_post_state.account(), &account);
-        assert_eq!(account_post_state.account_mut(), &mut account);
-    }
-
-    // ---- AccountId::for_private_pda tests ----
-
-    /// Pins `AccountId::for_private_pda` against a hardcoded expected output for a specific
-    /// `(program_id, seed, npk, identifier)` tuple. Any change to `PRIVATE_PDA_PREFIX`, byte
-    /// ordering, or the underlying hash breaks this test.
-    #[test]
-    fn for_private_pda_matches_pinned_value() {
-        let program_id: ProgramId = [1; 8];
-        let seed = PdaSeed::new([2; 32]);
-        let npk = NullifierPublicKey([3; 32]);
-        let identifier: Identifier = u128::MAX;
-        let expected = AccountId::new([
-            59, 239, 182, 97, 14, 220, 96, 115, 238, 133, 143, 33, 234, 82, 237, 255, 148, 110, 54,
-            124, 98, 159, 245, 101, 146, 182, 150, 54, 37, 62, 25, 17,
-        ]);
-        assert_eq!(
-            AccountId::for_private_pda(&program_id, &seed, &npk, identifier),
-            expected
-        );
-    }
-
-    /// Two groups with different viewing keys at the same (program, seed) get different addresses.
-    #[test]
-    fn for_private_pda_differs_for_different_npk() {
-        let program_id: ProgramId = [1; 8];
-        let seed = PdaSeed::new([2; 32]);
-        let npk_a = NullifierPublicKey([3; 32]);
-        let npk_b = NullifierPublicKey([4; 32]);
-        assert_ne!(
-            AccountId::for_private_pda(&program_id, &seed, &npk_a, u128::MAX),
-            AccountId::for_private_pda(&program_id, &seed, &npk_b, u128::MAX),
-        );
-    }
-
-    /// Different seeds produce different addresses, even with the same program and npk.
-    #[test]
-    fn for_private_pda_differs_for_different_seed() {
-        let program_id: ProgramId = [1; 8];
-        let seed_a = PdaSeed::new([2; 32]);
-        let seed_b = PdaSeed::new([5; 32]);
-        let npk = NullifierPublicKey([3; 32]);
-        assert_ne!(
-            AccountId::for_private_pda(&program_id, &seed_a, &npk, u128::MAX),
-            AccountId::for_private_pda(&program_id, &seed_b, &npk, u128::MAX),
-        );
-    }
-
-    /// Different programs produce different addresses, even with the same seed and npk.
-    #[test]
-    fn for_private_pda_differs_for_different_program_id() {
-        let program_id_a: ProgramId = [1; 8];
-        let program_id_b: ProgramId = [9; 8];
-        let seed = PdaSeed::new([2; 32]);
-        let npk = NullifierPublicKey([3; 32]);
-        assert_ne!(
-            AccountId::for_private_pda(&program_id_a, &seed, &npk, u128::MAX),
-            AccountId::for_private_pda(&program_id_b, &seed, &npk, u128::MAX),
-        );
-    }
-
-    /// Different identifiers produce different addresses for the same `(program_id, seed, npk)`,
-    /// confirming that each `(program_id, seed, npk)` tuple controls a family of 2^128 addresses.
-    #[test]
-    fn for_private_pda_differs_for_different_identifier() {
-        let program_id: ProgramId = [1; 8];
-        let seed = PdaSeed::new([2; 32]);
-        let npk = NullifierPublicKey([3; 32]);
-        assert_ne!(
-            AccountId::for_private_pda(&program_id, &seed, &npk, 0),
-            AccountId::for_private_pda(&program_id, &seed, &npk, 1),
-        );
-        assert_ne!(
-            AccountId::for_private_pda(&program_id, &seed, &npk, 0),
-            AccountId::for_private_pda(&program_id, &seed, &npk, u128::MAX),
-        );
-    }
-
-    /// A private PDA at the same (program, seed) has a different address than a public PDA,
-    /// because the private formula uses a different prefix and includes npk.
-    #[test]
-    fn for_private_pda_differs_from_public_pda() {
-        let program_id: ProgramId = [1; 8];
-        let seed = PdaSeed::new([2; 32]);
-        let npk = NullifierPublicKey([3; 32]);
-        let private_id = AccountId::for_private_pda(&program_id, &seed, &npk, u128::MAX);
-        let public_id = AccountId::for_public_pda(&program_id, &seed);
-        assert_ne!(private_id, public_id);
-    }
-
-    #[cfg(feature = "host")]
-    #[test]
-    fn private_account_kind_header_round_trips() {
-        let regular = PrivateAccountKind::Regular(42);
-        let pda = PrivateAccountKind::Pda {
-            program_id: [1_u32; 8],
-            seed: PdaSeed::new([2_u8; 32]),
-            identifier: u128::MAX,
-        };
-        assert_eq!(
-            PrivateAccountKind::from_header_bytes(&regular.to_header_bytes()),
-            Some(regular)
-        );
-        assert_eq!(
-            PrivateAccountKind::from_header_bytes(&pda.to_header_bytes()),
-            Some(pda)
-        );
-    }
-
-    #[cfg(feature = "host")]
-    #[test]
-    fn private_account_kind_unknown_discriminant_returns_none() {
-        let mut bytes = [0_u8; PrivateAccountKind::HEADER_LEN];
-        bytes[0] = 0xFF;
-        assert_eq!(PrivateAccountKind::from_header_bytes(&bytes), None);
-    }
-
-    #[test]
-    fn for_private_account_dispatches_correctly() {
-        let program_id: ProgramId = [1; 8];
-        let seed = PdaSeed::new([2; 32]);
-        let npk = NullifierPublicKey([3; 32]);
-        let identifier: Identifier = 77;
-
-        assert_eq!(
-            AccountId::for_private_account(&npk, &PrivateAccountKind::Regular(identifier)),
-            AccountId::for_regular_private_account(&npk, identifier),
-        );
-        assert_eq!(
-            AccountId::for_private_account(
-                &npk,
-                &PrivateAccountKind::Pda {
-                    program_id,
-                    seed,
-                    identifier
-                }
-            ),
-            AccountId::for_private_pda(&program_id, &seed, &npk, identifier),
-        );
-    }
-
-    #[test]
-    fn compute_public_authorized_pdas_with_seeds() {
-        let caller: ProgramId = [1; 8];
-        let seed = PdaSeed::new([2; 32]);
-        let result = compute_public_authorized_pdas(Some(caller), &[seed]);
-        let expected = AccountId::for_public_pda(&caller, &seed);
-        assert!(result.contains(&expected));
-        assert_eq!(result.len(), 1);
-    }
-
-    /// With no caller (top-level call), the result is always empty.
-    #[test]
-    fn compute_public_authorized_pdas_no_caller_returns_empty() {
-        let seed = PdaSeed::new([2; 32]);
-        let result = compute_public_authorized_pdas(None, &[seed]);
-        assert!(result.is_empty());
-    }
-}
+mod tests;

lee/state_machine/core/src/program/tests.rs

@@ -0,0 +1,333 @@
+use super::*;
+
+#[test]
+fn validity_window_unbounded_accepts_any_value() {
+    let w: ValidityWindow<u64> = ValidityWindow::new_unbounded();
+    assert!(w.is_valid_for(0));
+    assert!(w.is_valid_for(u64::MAX));
+}
+
+#[test]
+fn validity_window_bounded_range_includes_from_excludes_to() {
+    let w: ValidityWindow<u64> = (Some(5), Some(10)).try_into().unwrap();
+    assert!(!w.is_valid_for(4));
+    assert!(w.is_valid_for(5));
+    assert!(w.is_valid_for(9));
+    assert!(!w.is_valid_for(10));
+}
+
+#[test]
+fn validity_window_only_from_bound() {
+    let w: ValidityWindow<u64> = (Some(5), None).try_into().unwrap();
+    assert!(!w.is_valid_for(4));
+    assert!(w.is_valid_for(5));
+    assert!(w.is_valid_for(u64::MAX));
+}
+
+#[test]
+fn validity_window_only_to_bound() {
+    let w: ValidityWindow<u64> = (None, Some(5)).try_into().unwrap();
+    assert!(w.is_valid_for(0));
+    assert!(w.is_valid_for(4));
+    assert!(!w.is_valid_for(5));
+}
+
+#[test]
+fn validity_window_adjacent_bounds_are_invalid() {
+    // [5, 5) is an empty range — from == to
+    assert!(ValidityWindow::<u64>::try_from((Some(5), Some(5))).is_err());
+}
+
+#[test]
+fn validity_window_inverted_bounds_are_invalid() {
+    assert!(ValidityWindow::<u64>::try_from((Some(10), Some(5))).is_err());
+}
+
+#[test]
+fn validity_window_getters_match_construction() {
+    let w: ValidityWindow<u64> = (Some(3), Some(7)).try_into().unwrap();
+    assert_eq!(w.start(), Some(3));
+    assert_eq!(w.end(), Some(7));
+}
+
+#[test]
+fn validity_window_getters_for_unbounded() {
+    let w: ValidityWindow<u64> = ValidityWindow::new_unbounded();
+    assert_eq!(w.start(), None);
+    assert_eq!(w.end(), None);
+}
+
+#[test]
+fn validity_window_from_range() {
+    let w: ValidityWindow<u64> = ValidityWindow::try_from(5_u64..10).unwrap();
+    assert_eq!(w.start(), Some(5));
+    assert_eq!(w.end(), Some(10));
+}
+
+#[test]
+fn validity_window_from_range_empty_is_invalid() {
+    assert!(ValidityWindow::<u64>::try_from(5_u64..5).is_err());
+}
+
+#[test]
+fn validity_window_from_range_inverted_is_invalid() {
+    let from = 10_u64;
+    let to = 5_u64;
+    assert!(ValidityWindow::<u64>::try_from(from..to).is_err());
+}
+
+#[test]
+fn validity_window_from_range_from() {
+    let w: ValidityWindow<u64> = (5_u64..).into();
+    assert_eq!(w.start(), Some(5));
+    assert_eq!(w.end(), None);
+}
+
+#[test]
+fn validity_window_from_range_to() {
+    let w: ValidityWindow<u64> = (..10_u64).into();
+    assert_eq!(w.start(), None);
+    assert_eq!(w.end(), Some(10));
+}
+
+#[test]
+fn validity_window_from_range_full() {
+    let w: ValidityWindow<u64> = (..).into();
+    assert_eq!(w.start(), None);
+    assert_eq!(w.end(), None);
+}
+
+#[test]
+fn program_output_try_with_block_validity_window_range() {
+    let output = ProgramOutput::new(DEFAULT_PROGRAM_ID, None, vec![], vec![], vec![])
+        .try_with_block_validity_window(10_u64..100)
+        .unwrap();
+    assert_eq!(output.block_validity_window.start(), Some(10));
+    assert_eq!(output.block_validity_window.end(), Some(100));
+}
+
+#[test]
+fn program_output_with_block_validity_window_range_from() {
+    let output = ProgramOutput::new(DEFAULT_PROGRAM_ID, None, vec![], vec![], vec![])
+        .with_block_validity_window(10_u64..);
+    assert_eq!(output.block_validity_window.start(), Some(10));
+    assert_eq!(output.block_validity_window.end(), None);
+}
+
+#[test]
+fn program_output_with_block_validity_window_range_to() {
+    let output = ProgramOutput::new(DEFAULT_PROGRAM_ID, None, vec![], vec![], vec![])
+        .with_block_validity_window(..100_u64);
+    assert_eq!(output.block_validity_window.start(), None);
+    assert_eq!(output.block_validity_window.end(), Some(100));
+}
+
+#[test]
+fn program_output_try_with_block_validity_window_empty_range_fails() {
+    let result = ProgramOutput::new(DEFAULT_PROGRAM_ID, None, vec![], vec![], vec![])
+        .try_with_block_validity_window(5_u64..5);
+    assert!(result.is_err());
+}
+
+#[test]
+fn post_state_new_with_claim_constructor() {
+    let account = Account {
+        program_owner: [1, 2, 3, 4, 5, 6, 7, 8],
+        balance: 1337,
+        data: vec![0xde, 0xad, 0xbe, 0xef].try_into().unwrap(),
+        nonce: 10_u128.into(),
+    };
+
+    let account_post_state = AccountPostState::new_claimed(account.clone(), Claim::Authorized);
+
+    assert_eq!(account, account_post_state.account);
+    assert_eq!(account_post_state.required_claim(), Some(Claim::Authorized));
+}
+
+#[test]
+fn post_state_new_without_claim_constructor() {
+    let account = Account {
+        program_owner: [1, 2, 3, 4, 5, 6, 7, 8],
+        balance: 1337,
+        data: vec![0xde, 0xad, 0xbe, 0xef].try_into().unwrap(),
+        nonce: 10_u128.into(),
+    };
+
+    let account_post_state = AccountPostState::new(account.clone());
+
+    assert_eq!(account, account_post_state.account);
+    assert!(account_post_state.required_claim().is_none());
+}
+
+#[test]
+fn post_state_account_getter() {
+    let mut account = Account {
+        program_owner: [1, 2, 3, 4, 5, 6, 7, 8],
+        balance: 1337,
+        data: vec![0xde, 0xad, 0xbe, 0xef].try_into().unwrap(),
+        nonce: 10_u128.into(),
+    };
+
+    let mut account_post_state = AccountPostState::new(account.clone());
+
+    assert_eq!(account_post_state.account(), &account);
+    assert_eq!(account_post_state.account_mut(), &mut account);
+}
+
+// ---- AccountId::for_private_pda tests ----
+
+/// Pins `AccountId::for_private_pda` against a hardcoded expected output for a specific
+/// `(program_id, seed, npk, identifier)` tuple. Any change to `PRIVATE_PDA_PREFIX`, byte
+/// ordering, or the underlying hash breaks this test.
+#[test]
+fn for_private_pda_matches_pinned_value() {
+    let program_id: ProgramId = [1; 8];
+    let seed = PdaSeed::new([2; 32]);
+    let npk = NullifierPublicKey([3; 32]);
+    let identifier: Identifier = u128::MAX;
+    let expected = AccountId::new([
+        59, 239, 182, 97, 14, 220, 96, 115, 238, 133, 143, 33, 234, 82, 237, 255, 148, 110, 54,
+        124, 98, 159, 245, 101, 146, 182, 150, 54, 37, 62, 25, 17,
+    ]);
+    assert_eq!(
+        AccountId::for_private_pda(&program_id, &seed, &npk, identifier),
+        expected
+    );
+}
+
+/// Two groups with different viewing keys at the same (program, seed) get different addresses.
+#[test]
+fn for_private_pda_differs_for_different_npk() {
+    let program_id: ProgramId = [1; 8];
+    let seed = PdaSeed::new([2; 32]);
+    let npk_a = NullifierPublicKey([3; 32]);
+    let npk_b = NullifierPublicKey([4; 32]);
+    assert_ne!(
+        AccountId::for_private_pda(&program_id, &seed, &npk_a, u128::MAX),
+        AccountId::for_private_pda(&program_id, &seed, &npk_b, u128::MAX),
+    );
+}
+
+/// Different seeds produce different addresses, even with the same program and npk.
+#[test]
+fn for_private_pda_differs_for_different_seed() {
+    let program_id: ProgramId = [1; 8];
+    let seed_a = PdaSeed::new([2; 32]);
+    let seed_b = PdaSeed::new([5; 32]);
+    let npk = NullifierPublicKey([3; 32]);
+    assert_ne!(
+        AccountId::for_private_pda(&program_id, &seed_a, &npk, u128::MAX),
+        AccountId::for_private_pda(&program_id, &seed_b, &npk, u128::MAX),
+    );
+}
+
+/// Different programs produce different addresses, even with the same seed and npk.
+#[test]
+fn for_private_pda_differs_for_different_program_id() {
+    let program_id_a: ProgramId = [1; 8];
+    let program_id_b: ProgramId = [9; 8];
+    let seed = PdaSeed::new([2; 32]);
+    let npk = NullifierPublicKey([3; 32]);
+    assert_ne!(
+        AccountId::for_private_pda(&program_id_a, &seed, &npk, u128::MAX),
+        AccountId::for_private_pda(&program_id_b, &seed, &npk, u128::MAX),
+    );
+}
+
+/// Different identifiers produce different addresses for the same `(program_id, seed, npk)`,
+/// confirming that each `(program_id, seed, npk)` tuple controls a family of 2^128 addresses.
+#[test]
+fn for_private_pda_differs_for_different_identifier() {
+    let program_id: ProgramId = [1; 8];
+    let seed = PdaSeed::new([2; 32]);
+    let npk = NullifierPublicKey([3; 32]);
+    assert_ne!(
+        AccountId::for_private_pda(&program_id, &seed, &npk, 0),
+        AccountId::for_private_pda(&program_id, &seed, &npk, 1),
+    );
+    assert_ne!(
+        AccountId::for_private_pda(&program_id, &seed, &npk, 0),
+        AccountId::for_private_pda(&program_id, &seed, &npk, u128::MAX),
+    );
+}
+
+/// A private PDA at the same (program, seed) has a different address than a public PDA,
+/// because the private formula uses a different prefix and includes npk.
+#[test]
+fn for_private_pda_differs_from_public_pda() {
+    let program_id: ProgramId = [1; 8];
+    let seed = PdaSeed::new([2; 32]);
+    let npk = NullifierPublicKey([3; 32]);
+    let private_id = AccountId::for_private_pda(&program_id, &seed, &npk, u128::MAX);
+    let public_id = AccountId::for_public_pda(&program_id, &seed);
+    assert_ne!(private_id, public_id);
+}
+
+#[cfg(feature = "host")]
+#[test]
+fn private_account_kind_header_round_trips() {
+    let regular = PrivateAccountKind::Regular(42);
+    let pda = PrivateAccountKind::Pda {
+        program_id: [1_u32; 8],
+        seed: PdaSeed::new([2_u8; 32]),
+        identifier: u128::MAX,
+    };
+    assert_eq!(
+        PrivateAccountKind::from_header_bytes(&regular.to_header_bytes()),
+        Some(regular)
+    );
+    assert_eq!(
+        PrivateAccountKind::from_header_bytes(&pda.to_header_bytes()),
+        Some(pda)
+    );
+}
+
+#[cfg(feature = "host")]
+#[test]
+fn private_account_kind_unknown_discriminant_returns_none() {
+    let mut bytes = [0_u8; PrivateAccountKind::HEADER_LEN];
+    bytes[0] = 0xFF;
+    assert_eq!(PrivateAccountKind::from_header_bytes(&bytes), None);
+}
+
+#[test]
+fn for_private_account_dispatches_correctly() {
+    let program_id: ProgramId = [1; 8];
+    let seed = PdaSeed::new([2; 32]);
+    let npk = NullifierPublicKey([3; 32]);
+    let identifier: Identifier = 77;
+
+    assert_eq!(
+        AccountId::for_private_account(&npk, &PrivateAccountKind::Regular(identifier)),
+        AccountId::for_regular_private_account(&npk, identifier),
+    );
+    assert_eq!(
+        AccountId::for_private_account(
+            &npk,
+            &PrivateAccountKind::Pda {
+                program_id,
+                seed,
+                identifier
+            }
+        ),
+        AccountId::for_private_pda(&program_id, &seed, &npk, identifier),
+    );
+}
+
+#[test]
+fn compute_public_authorized_pdas_with_seeds() {
+    let caller: ProgramId = [1; 8];
+    let seed = PdaSeed::new([2; 32]);
+    let result = compute_public_authorized_pdas(Some(caller), &[seed]);
+    let expected = AccountId::for_public_pda(&caller, &seed);
+    assert!(result.contains(&expected));
+    assert_eq!(result.len(), 1);
+}
+
+/// With no caller (top-level call), the result is always empty.
+#[test]
+fn compute_public_authorized_pdas_no_caller_returns_empty() {
+    let seed = PdaSeed::new([2; 32]);
+    let result = compute_public_authorized_pdas(None, &[seed]);
+    assert!(result.is_empty());
+}

lee/state_machine/src/merkle_tree/mod.rs

@@ -164,398 +164,4 @@ const fn prev_power_of_two(x: usize) -> usize {
 }
 
 #[cfg(test)]
-mod tests {
-    use hex_literal::hex;
-
-    use super::*;
-
-    impl MerkleTree {
-        pub fn new(values: &[Value]) -> Self {
-            let mut this = Self::with_capacity(values.len());
-            for value in values.iter().copied() {
-                this.insert(value);
-            }
-            this
-        }
-    }
-
-    #[test]
-    fn empty_merkle_tree() {
-        let tree = MerkleTree::with_capacity(4);
-        let expected_root =
-            hex!("0000000000000000000000000000000000000000000000000000000000000000");
-        assert_eq!(tree.root(), expected_root);
-        assert_eq!(tree.capacity, 4);
-        assert_eq!(tree.length, 0);
-    }
-
-    #[test]
-    fn merkle_tree_0() {
-        let values = [[0; 32]];
-        let tree = MerkleTree::new(&values);
-        assert_eq!(tree.root(), hash_value(&[0; 32]));
-        assert_eq!(tree.capacity, 1);
-        assert_eq!(tree.length, 1);
-    }
-
-    #[test]
-    fn merkle_tree_1() {
-        let values = [[1; 32], [2; 32], [3; 32], [4; 32]];
-        let tree = MerkleTree::new(&values);
-        let expected_root =
-            hex!("48c73f7821a58a8d2a703e5b39c571c0aa20cf14abcd0af8f2b955bc202998de");
-        assert_eq!(tree.root(), expected_root);
-        assert_eq!(tree.capacity, 4);
-        assert_eq!(tree.length, 4);
-    }
-
-    #[test]
-    fn merkle_tree_2() {
-        let values = [[1; 32], [2; 32], [3; 32], [0; 32]];
-        let tree = MerkleTree::new(&values);
-        let expected_root =
-            hex!("c9bbb83096df85157a146e7d770455a98412dee0633187ee86fee6c8a45b831a");
-        assert_eq!(tree.root(), expected_root);
-        assert_eq!(tree.capacity, 4);
-        assert_eq!(tree.length, 4);
-    }
-
-    #[test]
-    fn merkle_tree_3() {
-        let values = [[1; 32], [2; 32], [3; 32]];
-        let tree = MerkleTree::new(&values);
-        let expected_root =
-            hex!("c8d3d8d2b13f27ceeccdc699119871f9f32ea7ed86ff45d0ad11f77b28cd7568");
-        assert_eq!(tree.root(), expected_root);
-        assert_eq!(tree.capacity, 4);
-        assert_eq!(tree.length, 3);
-    }
-
-    #[test]
-    fn merkle_tree_4() {
-        let values = [[11; 32], [12; 32], [13; 32], [14; 32], [15; 32]];
-        let tree = MerkleTree::new(&values);
-        let expected_root =
-            hex!("ef418aed5aa20702d4d94c92da79a4012f2e36f1008bfdb3cd1e38749dca2499");
-
-        assert_eq!(tree.root(), expected_root);
-        assert_eq!(tree.capacity, 8);
-        assert_eq!(tree.length, 5);
-    }
-
-    #[test]
-    fn merkle_tree_5() {
-        let values = [
-            [11; 32], [12; 32], [12; 32], [13; 32], [14; 32], [15; 32], [15; 32], [13; 32],
-            [13; 32], [15; 32], [11; 32],
-        ];
-        let tree = MerkleTree::new(&values);
-        let expected_root =
-            hex!("3f72d2ff55921a86c48e5988ec3e19ee9d0d5aa3e23197842970a903508ed767");
-        assert_eq!(tree.root(), expected_root);
-        assert_eq!(tree.capacity, 16);
-        assert_eq!(tree.length, 11);
-    }
-
-    #[test]
-    fn merkle_tree_6() {
-        let values = [[1; 32], [2; 32], [3; 32], [4; 32], [5; 32]];
-        let tree = MerkleTree::new(&values);
-        let expected_root =
-            hex!("069cb8259a06fe6edb3fa7ff7933a6dd7dca6fca299314379794a688926c3792");
-        assert_eq!(tree.root(), expected_root);
-    }
-
-    #[test]
-    fn with_capacity_4() {
-        let tree = MerkleTree::with_capacity(4);
-
-        assert_eq!(tree.length, 0);
-        assert_eq!(tree.nodes.len(), 7);
-        for i in 3..7 {
-            assert_eq!(*tree.get_node(i), default_values::DEFAULT_VALUES[0], "{i}");
-        }
-        for i in 1..3 {
-            assert_eq!(*tree.get_node(i), default_values::DEFAULT_VALUES[1], "{i}");
-        }
-        assert_eq!(*tree.get_node(0), default_values::DEFAULT_VALUES[2]);
-    }
-
-    #[test]
-    fn with_capacity_5() {
-        let tree = MerkleTree::with_capacity(5);
-
-        assert_eq!(tree.length, 0);
-        assert_eq!(tree.nodes.len(), 15);
-        for i in 7..15 {
-            assert_eq!(*tree.get_node(i), default_values::DEFAULT_VALUES[0]);
-        }
-        for i in 3..7 {
-            assert_eq!(*tree.get_node(i), default_values::DEFAULT_VALUES[1]);
-        }
-        for i in 1..3 {
-            assert_eq!(*tree.get_node(i), default_values::DEFAULT_VALUES[2]);
-        }
-        assert_eq!(*tree.get_node(0), default_values::DEFAULT_VALUES[3]);
-    }
-
-    #[test]
-    fn with_capacity_6() {
-        let mut tree = MerkleTree::with_capacity(100);
-
-        let values = [[1; 32], [2; 32], [3; 32], [4; 32]];
-
-        let expected_root =
-            hex!("48c73f7821a58a8d2a703e5b39c571c0aa20cf14abcd0af8f2b955bc202998de");
-
-        assert_eq!(0, tree.insert(values[0]));
-        assert_eq!(1, tree.insert(values[1]));
-        assert_eq!(2, tree.insert(values[2]));
-        assert_eq!(3, tree.insert(values[3]));
-
-        assert_eq!(tree.root(), expected_root);
-    }
-
-    #[test]
-    fn with_capacity_7() {
-        let mut tree = MerkleTree::with_capacity(599);
-
-        let values = [[1; 32], [2; 32], [3; 32]];
-
-        let expected_root =
-            hex!("c8d3d8d2b13f27ceeccdc699119871f9f32ea7ed86ff45d0ad11f77b28cd7568");
-
-        assert_eq!(0, tree.insert(values[0]));
-        assert_eq!(1, tree.insert(values[1]));
-        assert_eq!(2, tree.insert(values[2]));
-
-        assert_eq!(tree.root(), expected_root);
-    }
-
-    #[test]
-    fn with_capacity_8() {
-        let mut tree = MerkleTree::with_capacity(1);
-
-        let values = [[1; 32], [2; 32], [3; 32]];
-
-        let expected_root =
-            hex!("c8d3d8d2b13f27ceeccdc699119871f9f32ea7ed86ff45d0ad11f77b28cd7568");
-
-        assert_eq!(0, tree.insert(values[0]));
-        assert_eq!(1, tree.insert(values[1]));
-        assert_eq!(2, tree.insert(values[2]));
-
-        assert_eq!(tree.root(), expected_root);
-    }
-
-    #[test]
-    fn insert_value_1() {
-        let mut tree = MerkleTree::with_capacity(1);
-
-        let values = [[1; 32], [2; 32], [3; 32]];
-        let expected_tree = MerkleTree::new(&values);
-
-        assert_eq!(0, tree.insert(values[0]));
-        assert_eq!(1, tree.insert(values[1]));
-        assert_eq!(2, tree.insert(values[2]));
-
-        assert_eq!(expected_tree, tree);
-    }
-
-    #[test]
-    fn insert_value_2() {
-        let mut tree = MerkleTree::with_capacity(1);
-
-        let values = [[1; 32], [2; 32], [3; 32], [4; 32]];
-        let expected_tree = MerkleTree::new(&values);
-
-        assert_eq!(0, tree.insert(values[0]));
-        assert_eq!(1, tree.insert(values[1]));
-        assert_eq!(2, tree.insert(values[2]));
-        assert_eq!(3, tree.insert(values[3]));
-
-        assert_eq!(expected_tree, tree);
-    }
-
-    #[test]
-    fn insert_value_3() {
-        let mut tree = MerkleTree::with_capacity(1);
-
-        let values = [[11; 32], [12; 32], [13; 32], [14; 32], [15; 32]];
-        let expected_tree = MerkleTree::new(&values);
-
-        tree.insert(values[0]);
-        tree.insert(values[1]);
-        tree.insert(values[2]);
-        tree.insert(values[3]);
-        tree.insert(values[4]);
-
-        assert_eq!(expected_tree, tree);
-    }
-
-    // Reference implementation
-    fn verify_authentication_path(value: &Value, index: usize, path: &[Node], root: &Node) -> bool {
-        let mut result = hash_value(value);
-        let mut level_index = index;
-        for node in path {
-            let is_left_child = level_index & 1 == 0;
-            if is_left_child {
-                result = hash_two(&result, node);
-            } else {
-                result = hash_two(node, &result);
-            }
-            level_index >>= 1;
-        }
-        &result == root
-    }
-
-    #[test]
-    fn authentication_path_1() {
-        let values = [[1; 32], [2; 32], [3; 32], [4; 32]];
-        let tree = MerkleTree::new(&values);
-        let expected_authentication_path = vec![
-            hex!("9f4fb68f3e1dac82202f9aa581ce0bbf1f765df0e9ac3c8c57e20f685abab8ed"),
-            hex!("50a27d4746f357cb700cbe9d4883b77fb64f0128828a3489dc6a6f21ddbf2414"),
-        ];
-
-        let authentication_path = tree.get_authentication_path_for(2).unwrap();
-        assert_eq!(authentication_path, expected_authentication_path);
-    }
-
-    #[test]
-    fn authentication_path_2() {
-        let values = [[1; 32], [2; 32], [3; 32]];
-        let tree = MerkleTree::new(&values);
-        let expected_authentication_path = vec![
-            hex!("75877bb41d393b5fb8455ce60ecd8dda001d06316496b14dfa7f895656eeca4a"),
-            hex!("a41b855d2db4de9052cd7be5ec67d6586629cb9f6e3246a4afa5ba313f07a9c5"),
-        ];
-
-        let authentication_path = tree.get_authentication_path_for(0).unwrap();
-        assert_eq!(authentication_path, expected_authentication_path);
-    }
-
-    #[test]
-    fn authentication_path_3() {
-        let values = [[1; 32], [2; 32], [3; 32], [4; 32], [5; 32]];
-        let tree = MerkleTree::new(&values);
-        let expected_authentication_path = vec![
-            hex!("0000000000000000000000000000000000000000000000000000000000000000"),
-            hex!("f5a5fd42d16a20302798ef6ed309979b43003d2320d9f0e8ea9831a92759fb4b"),
-            hex!("48c73f7821a58a8d2a703e5b39c571c0aa20cf14abcd0af8f2b955bc202998de"),
-        ];
-
-        let authentication_path = tree.get_authentication_path_for(4).unwrap();
-        assert_eq!(authentication_path, expected_authentication_path);
-    }
-
-    #[test]
-    fn authentication_path_4() {
-        let values = [[1; 32], [2; 32], [3; 32], [4; 32], [5; 32]];
-        let tree = MerkleTree::new(&values);
-        assert!(tree.get_authentication_path_for(5).is_none());
-    }
-
-    #[test]
-    fn authentication_path_5() {
-        let values = [[1; 32], [2; 32], [3; 32], [4; 32], [5; 32]];
-        let tree = MerkleTree::new(&values);
-        let index = 4;
-        let value = values[index];
-        let path = tree.get_authentication_path_for(index).unwrap();
-        assert!(verify_authentication_path(
-            &value,
-            index,
-            &path,
-            &tree.root()
-        ));
-    }
-
-    #[test]
-    fn tree_with_63_insertions() {
-        let values = [
-            hex!("cd00acab0f45736e6c6311f1953becc0b69a062e7c2a7310875d28bdf9ef9c5b"),
-            hex!("0df5a6afbcc7bf126caf7084acfc593593ab512e6ca433c61c1a922be40a04ea"),
-            hex!("23c1258620266c7bedb6d1ee32f6da9413e4010ace975239dccb34e727e07c40"),
-            hex!("f33ccc3a11476b0ef62326ca5ec292056759b05e6a28023d2d1ce66165611353"),
-            hex!("77f914ab016b8049f6bea7704000e413a393865918a3824f9285c3db0aacff23"),
-            hex!("910a1c23188e54d57fd167ddb0f8bf68c6b70ed9ec76ef56c4b7f2632f82ca7f"),
-            hex!("047ee85526197d1e7403a559cf6d2f22c1926c8ad59481a2e2f1b697af45e40b"),
-            hex!("9d355cf89fb382ae34bf80566b28489278d10f2cebb5b0ea42fab1bac5adae0c"),
-            hex!("604018b95232596b2685a9bc737b6cccb53b10e483d2d9a2f4a755410b02a188"),
-            hex!("a16708ef7b6bf1796063addaf57d6a566b6f87b0bbe42af43a4590d05f1684cb"),
-            hex!("820f2dfa271cd2fd41e1452406d5dad552c85c1223c45d45dbd7446759fdc6b8"),
-            hex!("680b6912d7e219f8805d4d28adb4428dd78fea0dc1b8cdb2412645c4b1962c88"),
-            hex!("14d5471ce6c45506753982b17cac5790ac7bc29e6f388f31052d7dfd62b294e5"),
-            hex!("8b364200172b777d4aa16d2098b5eb98ac3dd4a1b9597e5c2bf6f6930031f230"),
-            hex!("9bb45b910711874339dda8a21a9aad73822286f5e52d7d3de0ed78dfbba329a5"),
-            hex!("d6806d5df5cb25ce5d531042f09b3cb34fb9e47c61182b63cccd9d44392f6027"),
-            hex!("b8cfa90ebc8fd09c04682d93a08fddd3e8e57715174dcc92451edd191264a58b"),
-            hex!("3463c7f81d00f809b3dfa83195447c927fb4045b3913dac6f45bee6c4010d7ed"),
-            hex!("1d6ad7f7d677905feb506c58f4b404a79370ebc567296abea3a368b61d5a8239"),
-            hex!("a58085ecf00963cb22da23c901b9b3ddc56462bb96ff03c923d67708e10dd29c"),
-            hex!("c3319f4a65fb5bbb8447137b0972c03cbd84ebf7d9da194e0fcbd68c2d4d5bdb"),
-            hex!("4aa31e90e0090faf3648d05e5d5499df2c78ebed4d6e6c23d8147de5d67dae73"),
-            hex!("9f33b1d2c8bc7bd265336de1033ede6344bc41260313bdcb43f1108b83b9be92"),
-            hex!("6500d4ad93d41c16ec81eaa5e70f173194aabe5c1072ac263b5727296f5b7cac"),
-            hex!("3584f5d260003669fad98786e13171376b0f19410cb232ce65606cbff79e6768"),
-            hex!("c8410946ebf56f13141c894a34ced85a5230088af70dcea581e44f52847830ac"),
-            hex!("71dd90281cdebb70422f2d04ae446d5d2d5ea64b803c16128d37e3fcd5d1a4cc"),
-            hex!("c05acf8d77ab4d659a538bd35af590864a7ad9c055ff5d6cda9d5aecfccecba3"),
-            hex!("f1df98822ea084cce9021aa9cc81b1746cd1e84a75690da63e10fd877633ed77"),
-            hex!("2ca822bc8f67bceb0a71a0d06fea7349036ef3e5ec21795a851e4182bd35ce01"),
-            hex!("7fd2179abc3bcf89b4d8092988ba8c23952b3bbd3d7caea6b5ea0c13cf19f68b"),
-            hex!("91b6ad516e017f6aa5a2e95776538bd3a3e933c1b1d32bb5e0f00a9db63c9c24"),
-            hex!("cd31a8b5eef5ca0be5ef1cb261d0bf0a74d774a3152bb99739cfd296a1d0b85e"),
-            hex!("3fb16f48b2bf93f3815979e6638f975d7f935088ec37db0be0f07965fbc78339"),
-            hex!("c60c61b99bf486af5f4bf780a69860dafcd35c1474306a8575666fb5449bcec0"),
-            hex!("8048d0d7e14091251f3f6c6b10bf6b5880a014b513f9f8c2395501dbffa6192a"),
-            hex!("778b5af10b9dbe80b60a8e4f0bb91caf4476bcb812801099760754ae623fbd84"),
-            hex!("d3ac25467920a4e08998b7a3226b8b54bfe66ac58cfedc71f15b2402fee0054a"),
-            hex!("029aa94598fae2961a0d43937b8a9a3138bcfeae99a7cb15f77fac7c506f8432"),
-            hex!("2eee5ef52fe669cb6882a68c893abdc1262dcf4424e4ba7a479da7cf1c10171d"),
-            hex!("de3fb3d070e3a90f0eed8b5e65088a8dc0e4e3c342b9c0bf33bab714eae5dfec"),
-            hex!("14d40177e833ab45bbfdc5f2b11fba7efaebb3f69facc554f24b549a2efe8538"),
-            hex!("5734355069702448774fb2df95f1d562e1b9fe1514aeb6b922554ee9d2d01068"),
-            hex!("8a273d49ac110343cec2cf3359d16eb2906b446bd9ec9833e2a640cebc8d5155"),
-            hex!("e3fa984dd3cbeb9a7e827ed32d3d4e6a6ba643a55d82be97d9ddb06ee809fa3e"),
-            hex!("90b1d5a364e17c8b7965396b06ec6e13749b5fc16500731518ad8fc30ae33e77"),
-            hex!("7517376541b2e8ec83cbab04522b54a26610908a9872feb663451385aea58eb1"),
-            hex!("5cba2e4cf7448e526d161133c4b2ea7c919ac4813a7308612595f46f11dea6cd"),
-            hex!("c721911b300bec0691c8a2dfaabfef1d66b7b6258918914d3c3ad690729f05b7"),
-            hex!("d0d0a70d8ae0d27806fa0b711c507290c260a89cbca0436d339d1dccdd087d62"),
-            hex!("2a625c28ea763c5e82dd0a93ecfca7ec371ccbb363cd42be359c2c875f58009d"),
-            hex!("174ef0119932ed890397d9f3837dd85f9100558b6fc9085d4af947ae8cf74bbc"),
-            hex!("b497bc267151e8efa3c6daa461e6804b01a3f05f44f1f4d5b41d5f0d3f5219b1"),
-            hex!("e987e91f5734630ddd7e6b58733b4fcdbc316ee9e8cac0e94c36c91cf58e59cc"),
-            hex!("55019ad8bbe656c51eb042190c1c8da53f42baf43fd2350ebea38fc7cca2fae3"),
-            hex!("c45a638edd18a6d9f5ad20b870c81b8626459bcb22dae7d58add7a6b6c6a84a8"),
-            hex!("d42d3a5fb2ad50b2027fe5a36d59dd71e49a63e4b1b299073c96bbf7ba5d68a1"),
-            hex!("9599e561054bcd3f647eb018ab0b069d3176497d42be9c4466551cbb959be47c"),
-            hex!("42f33b23775327ff71aea6569548255f3cc9929da73373cc9bb1743d417f7cda"),
-            hex!("ab24294f44fc6fdbeb96e0f6e93c4f6d97d035b73b9a337c353e18c6d0603bdd"),
-            hex!("33954ec63520334f99b640a2982ac966b68c363fed383d621a1ab573934f1d33"),
-            hex!("5e2a1f7df963d1fd8f50a285387cfbb5df581426619b325563e20bf7886c62b7"),
-            hex!("13ffde471d4e27c473254e766fd1328ad80c42cab4d4955cffeae43d866f86e5"),
-        ];
-
-        let expected_root =
-            hex!("1cf9b214217d7823f9de51b8f6cb34d0a99436a3a1bb762f90b815672a6afcc0");
-
-        let mut tree_less_capacity = MerkleTree::with_capacity(1);
-        let mut tree_exact_capacity = MerkleTree::with_capacity(64);
-        let mut tree_more_capacity = MerkleTree::with_capacity(128);
-
-        for value in &values {
-            tree_less_capacity.insert(*value);
-            tree_exact_capacity.insert(*value);
-            tree_more_capacity.insert(*value);
-        }
-
-        assert_eq!(tree_more_capacity.root(), expected_root);
-        assert_eq!(tree_less_capacity.root(), expected_root);
-        assert_eq!(tree_exact_capacity.root(), expected_root);
-    }
-}
-
-//
+mod tests;

lee/state_machine/src/merkle_tree/tests.rs

@@ -0,0 +1,380 @@
+use hex_literal::hex;
+
+use super::*;
+
+impl MerkleTree {
+    pub fn new(values: &[Value]) -> Self {
+        let mut this = Self::with_capacity(values.len());
+        for value in values.iter().copied() {
+            this.insert(value);
+        }
+        this
+    }
+}
+
+#[test]
+fn empty_merkle_tree() {
+    let tree = MerkleTree::with_capacity(4);
+    let expected_root = hex!("0000000000000000000000000000000000000000000000000000000000000000");
+    assert_eq!(tree.root(), expected_root);
+    assert_eq!(tree.capacity, 4);
+    assert_eq!(tree.length, 0);
+}
+
+#[test]
+fn merkle_tree_0() {
+    let values = [[0; 32]];
+    let tree = MerkleTree::new(&values);
+    assert_eq!(tree.root(), hash_value(&[0; 32]));
+    assert_eq!(tree.capacity, 1);
+    assert_eq!(tree.length, 1);
+}
+
+#[test]
+fn merkle_tree_1() {
+    let values = [[1; 32], [2; 32], [3; 32], [4; 32]];
+    let tree = MerkleTree::new(&values);
+    let expected_root = hex!("48c73f7821a58a8d2a703e5b39c571c0aa20cf14abcd0af8f2b955bc202998de");
+    assert_eq!(tree.root(), expected_root);
+    assert_eq!(tree.capacity, 4);
+    assert_eq!(tree.length, 4);
+}
+
+#[test]
+fn merkle_tree_2() {
+    let values = [[1; 32], [2; 32], [3; 32], [0; 32]];
+    let tree = MerkleTree::new(&values);
+    let expected_root = hex!("c9bbb83096df85157a146e7d770455a98412dee0633187ee86fee6c8a45b831a");
+    assert_eq!(tree.root(), expected_root);
+    assert_eq!(tree.capacity, 4);
+    assert_eq!(tree.length, 4);
+}
+
+#[test]
+fn merkle_tree_3() {
+    let values = [[1; 32], [2; 32], [3; 32]];
+    let tree = MerkleTree::new(&values);
+    let expected_root = hex!("c8d3d8d2b13f27ceeccdc699119871f9f32ea7ed86ff45d0ad11f77b28cd7568");
+    assert_eq!(tree.root(), expected_root);
+    assert_eq!(tree.capacity, 4);
+    assert_eq!(tree.length, 3);
+}
+
+#[test]
+fn merkle_tree_4() {
+    let values = [[11; 32], [12; 32], [13; 32], [14; 32], [15; 32]];
+    let tree = MerkleTree::new(&values);
+    let expected_root = hex!("ef418aed5aa20702d4d94c92da79a4012f2e36f1008bfdb3cd1e38749dca2499");
+
+    assert_eq!(tree.root(), expected_root);
+    assert_eq!(tree.capacity, 8);
+    assert_eq!(tree.length, 5);
+}
+
+#[test]
+fn merkle_tree_5() {
+    let values = [
+        [11; 32], [12; 32], [12; 32], [13; 32], [14; 32], [15; 32], [15; 32], [13; 32], [13; 32],
+        [15; 32], [11; 32],
+    ];
+    let tree = MerkleTree::new(&values);
+    let expected_root = hex!("3f72d2ff55921a86c48e5988ec3e19ee9d0d5aa3e23197842970a903508ed767");
+    assert_eq!(tree.root(), expected_root);
+    assert_eq!(tree.capacity, 16);
+    assert_eq!(tree.length, 11);
+}
+
+#[test]
+fn merkle_tree_6() {
+    let values = [[1; 32], [2; 32], [3; 32], [4; 32], [5; 32]];
+    let tree = MerkleTree::new(&values);
+    let expected_root = hex!("069cb8259a06fe6edb3fa7ff7933a6dd7dca6fca299314379794a688926c3792");
+    assert_eq!(tree.root(), expected_root);
+}
+
+#[test]
+fn with_capacity_4() {
+    let tree = MerkleTree::with_capacity(4);
+
+    assert_eq!(tree.length, 0);
+    assert_eq!(tree.nodes.len(), 7);
+    for i in 3..7 {
+        assert_eq!(*tree.get_node(i), default_values::DEFAULT_VALUES[0], "{i}");
+    }
+    for i in 1..3 {
+        assert_eq!(*tree.get_node(i), default_values::DEFAULT_VALUES[1], "{i}");
+    }
+    assert_eq!(*tree.get_node(0), default_values::DEFAULT_VALUES[2]);
+}
+
+#[test]
+fn with_capacity_5() {
+    let tree = MerkleTree::with_capacity(5);
+
+    assert_eq!(tree.length, 0);
+    assert_eq!(tree.nodes.len(), 15);
+    for i in 7..15 {
+        assert_eq!(*tree.get_node(i), default_values::DEFAULT_VALUES[0]);
+    }
+    for i in 3..7 {
+        assert_eq!(*tree.get_node(i), default_values::DEFAULT_VALUES[1]);
+    }
+    for i in 1..3 {
+        assert_eq!(*tree.get_node(i), default_values::DEFAULT_VALUES[2]);
+    }
+    assert_eq!(*tree.get_node(0), default_values::DEFAULT_VALUES[3]);
+}
+
+#[test]
+fn with_capacity_6() {
+    let mut tree = MerkleTree::with_capacity(100);
+
+    let values = [[1; 32], [2; 32], [3; 32], [4; 32]];
+
+    let expected_root = hex!("48c73f7821a58a8d2a703e5b39c571c0aa20cf14abcd0af8f2b955bc202998de");
+
+    assert_eq!(0, tree.insert(values[0]));
+    assert_eq!(1, tree.insert(values[1]));
+    assert_eq!(2, tree.insert(values[2]));
+    assert_eq!(3, tree.insert(values[3]));
+
+    assert_eq!(tree.root(), expected_root);
+}
+
+#[test]
+fn with_capacity_7() {
+    let mut tree = MerkleTree::with_capacity(599);
+
+    let values = [[1; 32], [2; 32], [3; 32]];
+
+    let expected_root = hex!("c8d3d8d2b13f27ceeccdc699119871f9f32ea7ed86ff45d0ad11f77b28cd7568");
+
+    assert_eq!(0, tree.insert(values[0]));
+    assert_eq!(1, tree.insert(values[1]));
+    assert_eq!(2, tree.insert(values[2]));
+
+    assert_eq!(tree.root(), expected_root);
+}
+
+#[test]
+fn with_capacity_8() {
+    let mut tree = MerkleTree::with_capacity(1);
+
+    let values = [[1; 32], [2; 32], [3; 32]];
+
+    let expected_root = hex!("c8d3d8d2b13f27ceeccdc699119871f9f32ea7ed86ff45d0ad11f77b28cd7568");
+
+    assert_eq!(0, tree.insert(values[0]));
+    assert_eq!(1, tree.insert(values[1]));
+    assert_eq!(2, tree.insert(values[2]));
+
+    assert_eq!(tree.root(), expected_root);
+}
+
+#[test]
+fn insert_value_1() {
+    let mut tree = MerkleTree::with_capacity(1);
+
+    let values = [[1; 32], [2; 32], [3; 32]];
+    let expected_tree = MerkleTree::new(&values);
+
+    assert_eq!(0, tree.insert(values[0]));
+    assert_eq!(1, tree.insert(values[1]));
+    assert_eq!(2, tree.insert(values[2]));
+
+    assert_eq!(expected_tree, tree);
+}
+
+#[test]
+fn insert_value_2() {
+    let mut tree = MerkleTree::with_capacity(1);
+
+    let values = [[1; 32], [2; 32], [3; 32], [4; 32]];
+    let expected_tree = MerkleTree::new(&values);
+
+    assert_eq!(0, tree.insert(values[0]));
+    assert_eq!(1, tree.insert(values[1]));
+    assert_eq!(2, tree.insert(values[2]));
+    assert_eq!(3, tree.insert(values[3]));
+
+    assert_eq!(expected_tree, tree);
+}
+
+#[test]
+fn insert_value_3() {
+    let mut tree = MerkleTree::with_capacity(1);
+
+    let values = [[11; 32], [12; 32], [13; 32], [14; 32], [15; 32]];
+    let expected_tree = MerkleTree::new(&values);
+
+    tree.insert(values[0]);
+    tree.insert(values[1]);
+    tree.insert(values[2]);
+    tree.insert(values[3]);
+    tree.insert(values[4]);
+
+    assert_eq!(expected_tree, tree);
+}
+
+// Reference implementation
+fn verify_authentication_path(value: &Value, index: usize, path: &[Node], root: &Node) -> bool {
+    let mut result = hash_value(value);
+    let mut level_index = index;
+    for node in path {
+        let is_left_child = level_index & 1 == 0;
+        if is_left_child {
+            result = hash_two(&result, node);
+        } else {
+            result = hash_two(node, &result);
+        }
+        level_index >>= 1;
+    }
+    &result == root
+}
+
+#[test]
+fn authentication_path_1() {
+    let values = [[1; 32], [2; 32], [3; 32], [4; 32]];
+    let tree = MerkleTree::new(&values);
+    let expected_authentication_path = vec![
+        hex!("9f4fb68f3e1dac82202f9aa581ce0bbf1f765df0e9ac3c8c57e20f685abab8ed"),
+        hex!("50a27d4746f357cb700cbe9d4883b77fb64f0128828a3489dc6a6f21ddbf2414"),
+    ];
+
+    let authentication_path = tree.get_authentication_path_for(2).unwrap();
+    assert_eq!(authentication_path, expected_authentication_path);
+}
+
+#[test]
+fn authentication_path_2() {
+    let values = [[1; 32], [2; 32], [3; 32]];
+    let tree = MerkleTree::new(&values);
+    let expected_authentication_path = vec![
+        hex!("75877bb41d393b5fb8455ce60ecd8dda001d06316496b14dfa7f895656eeca4a"),
+        hex!("a41b855d2db4de9052cd7be5ec67d6586629cb9f6e3246a4afa5ba313f07a9c5"),
+    ];
+
+    let authentication_path = tree.get_authentication_path_for(0).unwrap();
+    assert_eq!(authentication_path, expected_authentication_path);
+}
+
+#[test]
+fn authentication_path_3() {
+    let values = [[1; 32], [2; 32], [3; 32], [4; 32], [5; 32]];
+    let tree = MerkleTree::new(&values);
+    let expected_authentication_path = vec![
+        hex!("0000000000000000000000000000000000000000000000000000000000000000"),
+        hex!("f5a5fd42d16a20302798ef6ed309979b43003d2320d9f0e8ea9831a92759fb4b"),
+        hex!("48c73f7821a58a8d2a703e5b39c571c0aa20cf14abcd0af8f2b955bc202998de"),
+    ];
+
+    let authentication_path = tree.get_authentication_path_for(4).unwrap();
+    assert_eq!(authentication_path, expected_authentication_path);
+}
+
+#[test]
+fn authentication_path_4() {
+    let values = [[1; 32], [2; 32], [3; 32], [4; 32], [5; 32]];
+    let tree = MerkleTree::new(&values);
+    assert!(tree.get_authentication_path_for(5).is_none());
+}
+
+#[test]
+fn authentication_path_5() {
+    let values = [[1; 32], [2; 32], [3; 32], [4; 32], [5; 32]];
+    let tree = MerkleTree::new(&values);
+    let index = 4;
+    let value = values[index];
+    let path = tree.get_authentication_path_for(index).unwrap();
+    assert!(verify_authentication_path(
+        &value,
+        index,
+        &path,
+        &tree.root()
+    ));
+}
+
+#[test]
+fn tree_with_63_insertions() {
+    let values = [
+        hex!("cd00acab0f45736e6c6311f1953becc0b69a062e7c2a7310875d28bdf9ef9c5b"),
+        hex!("0df5a6afbcc7bf126caf7084acfc593593ab512e6ca433c61c1a922be40a04ea"),
+        hex!("23c1258620266c7bedb6d1ee32f6da9413e4010ace975239dccb34e727e07c40"),
+        hex!("f33ccc3a11476b0ef62326ca5ec292056759b05e6a28023d2d1ce66165611353"),
+        hex!("77f914ab016b8049f6bea7704000e413a393865918a3824f9285c3db0aacff23"),
+        hex!("910a1c23188e54d57fd167ddb0f8bf68c6b70ed9ec76ef56c4b7f2632f82ca7f"),
+        hex!("047ee85526197d1e7403a559cf6d2f22c1926c8ad59481a2e2f1b697af45e40b"),
+        hex!("9d355cf89fb382ae34bf80566b28489278d10f2cebb5b0ea42fab1bac5adae0c"),
+        hex!("604018b95232596b2685a9bc737b6cccb53b10e483d2d9a2f4a755410b02a188"),
+        hex!("a16708ef7b6bf1796063addaf57d6a566b6f87b0bbe42af43a4590d05f1684cb"),
+        hex!("820f2dfa271cd2fd41e1452406d5dad552c85c1223c45d45dbd7446759fdc6b8"),
+        hex!("680b6912d7e219f8805d4d28adb4428dd78fea0dc1b8cdb2412645c4b1962c88"),
+        hex!("14d5471ce6c45506753982b17cac5790ac7bc29e6f388f31052d7dfd62b294e5"),
+        hex!("8b364200172b777d4aa16d2098b5eb98ac3dd4a1b9597e5c2bf6f6930031f230"),
+        hex!("9bb45b910711874339dda8a21a9aad73822286f5e52d7d3de0ed78dfbba329a5"),
+        hex!("d6806d5df5cb25ce5d531042f09b3cb34fb9e47c61182b63cccd9d44392f6027"),
+        hex!("b8cfa90ebc8fd09c04682d93a08fddd3e8e57715174dcc92451edd191264a58b"),
+        hex!("3463c7f81d00f809b3dfa83195447c927fb4045b3913dac6f45bee6c4010d7ed"),
+        hex!("1d6ad7f7d677905feb506c58f4b404a79370ebc567296abea3a368b61d5a8239"),
+        hex!("a58085ecf00963cb22da23c901b9b3ddc56462bb96ff03c923d67708e10dd29c"),
+        hex!("c3319f4a65fb5bbb8447137b0972c03cbd84ebf7d9da194e0fcbd68c2d4d5bdb"),
+        hex!("4aa31e90e0090faf3648d05e5d5499df2c78ebed4d6e6c23d8147de5d67dae73"),
+        hex!("9f33b1d2c8bc7bd265336de1033ede6344bc41260313bdcb43f1108b83b9be92"),
+        hex!("6500d4ad93d41c16ec81eaa5e70f173194aabe5c1072ac263b5727296f5b7cac"),
+        hex!("3584f5d260003669fad98786e13171376b0f19410cb232ce65606cbff79e6768"),
+        hex!("c8410946ebf56f13141c894a34ced85a5230088af70dcea581e44f52847830ac"),
+        hex!("71dd90281cdebb70422f2d04ae446d5d2d5ea64b803c16128d37e3fcd5d1a4cc"),
+        hex!("c05acf8d77ab4d659a538bd35af590864a7ad9c055ff5d6cda9d5aecfccecba3"),
+        hex!("f1df98822ea084cce9021aa9cc81b1746cd1e84a75690da63e10fd877633ed77"),
+        hex!("2ca822bc8f67bceb0a71a0d06fea7349036ef3e5ec21795a851e4182bd35ce01"),
+        hex!("7fd2179abc3bcf89b4d8092988ba8c23952b3bbd3d7caea6b5ea0c13cf19f68b"),
+        hex!("91b6ad516e017f6aa5a2e95776538bd3a3e933c1b1d32bb5e0f00a9db63c9c24"),
+        hex!("cd31a8b5eef5ca0be5ef1cb261d0bf0a74d774a3152bb99739cfd296a1d0b85e"),
+        hex!("3fb16f48b2bf93f3815979e6638f975d7f935088ec37db0be0f07965fbc78339"),
+        hex!("c60c61b99bf486af5f4bf780a69860dafcd35c1474306a8575666fb5449bcec0"),
+        hex!("8048d0d7e14091251f3f6c6b10bf6b5880a014b513f9f8c2395501dbffa6192a"),
+        hex!("778b5af10b9dbe80b60a8e4f0bb91caf4476bcb812801099760754ae623fbd84"),
+        hex!("d3ac25467920a4e08998b7a3226b8b54bfe66ac58cfedc71f15b2402fee0054a"),
+        hex!("029aa94598fae2961a0d43937b8a9a3138bcfeae99a7cb15f77fac7c506f8432"),
+        hex!("2eee5ef52fe669cb6882a68c893abdc1262dcf4424e4ba7a479da7cf1c10171d"),
+        hex!("de3fb3d070e3a90f0eed8b5e65088a8dc0e4e3c342b9c0bf33bab714eae5dfec"),
+        hex!("14d40177e833ab45bbfdc5f2b11fba7efaebb3f69facc554f24b549a2efe8538"),
+        hex!("5734355069702448774fb2df95f1d562e1b9fe1514aeb6b922554ee9d2d01068"),
+        hex!("8a273d49ac110343cec2cf3359d16eb2906b446bd9ec9833e2a640cebc8d5155"),
+        hex!("e3fa984dd3cbeb9a7e827ed32d3d4e6a6ba643a55d82be97d9ddb06ee809fa3e"),
+        hex!("90b1d5a364e17c8b7965396b06ec6e13749b5fc16500731518ad8fc30ae33e77"),
+        hex!("7517376541b2e8ec83cbab04522b54a26610908a9872feb663451385aea58eb1"),
+        hex!("5cba2e4cf7448e526d161133c4b2ea7c919ac4813a7308612595f46f11dea6cd"),
+        hex!("c721911b300bec0691c8a2dfaabfef1d66b7b6258918914d3c3ad690729f05b7"),
+        hex!("d0d0a70d8ae0d27806fa0b711c507290c260a89cbca0436d339d1dccdd087d62"),
+        hex!("2a625c28ea763c5e82dd0a93ecfca7ec371ccbb363cd42be359c2c875f58009d"),
+        hex!("174ef0119932ed890397d9f3837dd85f9100558b6fc9085d4af947ae8cf74bbc"),
+        hex!("b497bc267151e8efa3c6daa461e6804b01a3f05f44f1f4d5b41d5f0d3f5219b1"),
+        hex!("e987e91f5734630ddd7e6b58733b4fcdbc316ee9e8cac0e94c36c91cf58e59cc"),
+        hex!("55019ad8bbe656c51eb042190c1c8da53f42baf43fd2350ebea38fc7cca2fae3"),
+        hex!("c45a638edd18a6d9f5ad20b870c81b8626459bcb22dae7d58add7a6b6c6a84a8"),
+        hex!("d42d3a5fb2ad50b2027fe5a36d59dd71e49a63e4b1b299073c96bbf7ba5d68a1"),
+        hex!("9599e561054bcd3f647eb018ab0b069d3176497d42be9c4466551cbb959be47c"),
+        hex!("42f33b23775327ff71aea6569548255f3cc9929da73373cc9bb1743d417f7cda"),
+        hex!("ab24294f44fc6fdbeb96e0f6e93c4f6d97d035b73b9a337c353e18c6d0603bdd"),
+        hex!("33954ec63520334f99b640a2982ac966b68c363fed383d621a1ab573934f1d33"),
+        hex!("5e2a1f7df963d1fd8f50a285387cfbb5df581426619b325563e20bf7886c62b7"),
+        hex!("13ffde471d4e27c473254e766fd1328ad80c42cab4d4955cffeae43d866f86e5"),
+    ];
+
+    let expected_root = hex!("1cf9b214217d7823f9de51b8f6cb34d0a99436a3a1bb762f90b815672a6afcc0");
+
+    let mut tree_less_capacity = MerkleTree::with_capacity(1);
+    let mut tree_exact_capacity = MerkleTree::with_capacity(64);
+    let mut tree_more_capacity = MerkleTree::with_capacity(128);
+
+    for value in &values {
+        tree_less_capacity.insert(*value);
+        tree_exact_capacity.insert(*value);
+        tree_more_capacity.insert(*value);
+    }
+
+    assert_eq!(tree_more_capacity.root(), expected_root);
+    assert_eq!(tree_less_capacity.root(), expected_root);
+    assert_eq!(tree_exact_capacity.root(), expected_root);
+}

lee/state_machine/src/privacy_preserving_transaction/circuit.rs

@@ -1,906 +0,0 @@
-use std::collections::{HashMap, VecDeque};
-
-use borsh::{BorshDeserialize, BorshSerialize};
-use lee_core::{
-    InputAccountIdentity, PrivacyPreservingCircuitInput, PrivacyPreservingCircuitOutput,
-    account::AccountWithMetadata,
-    program::{ChainedCall, InstructionData, ProgramId, ProgramOutput},
-};
-use risc0_zkvm::{ExecutorEnv, InnerReceipt, ProverOpts, Receipt, default_prover};
-
-use crate::{
-    PRIVACY_PRESERVING_CIRCUIT_ELF, PRIVACY_PRESERVING_CIRCUIT_ID,
-    error::{InvalidProgramBehaviorError, LeeError},
-    program::Program,
-    state::MAX_NUMBER_CHAINED_CALLS,
-};
-
-/// Proof of the privacy preserving execution circuit.
-#[derive(Debug, Clone, PartialEq, Eq, BorshSerialize, BorshDeserialize)]
-pub struct Proof(pub(crate) Vec<u8>);
-
-impl Proof {
-    #[must_use]
-    pub fn into_inner(self) -> Vec<u8> {
-        self.0
-    }
-
-    #[must_use]
-    pub const fn from_inner(inner: Vec<u8>) -> Self {
-        Self(inner)
-    }
-
-    pub(crate) fn is_valid_for(&self, circuit_output: &PrivacyPreservingCircuitOutput) -> bool {
-        let Ok(inner) = borsh::from_slice::<InnerReceipt>(&self.0) else {
-            return false;
-        };
-        let receipt = Receipt::new(inner, circuit_output.to_bytes());
-        receipt.verify(PRIVACY_PRESERVING_CIRCUIT_ID).is_ok()
-    }
-}
-
-#[derive(Clone)]
-pub struct ProgramWithDependencies {
-    pub program: Program,
-    // TODO: avoid having a copy of the bytecode of each dependency.
-    pub dependencies: HashMap<ProgramId, Program>,
-}
-
-impl ProgramWithDependencies {
-    #[must_use]
-    pub const fn new(program: Program, dependencies: HashMap<ProgramId, Program>) -> Self {
-        Self {
-            program,
-            dependencies,
-        }
-    }
-}
-
-impl From<Program> for ProgramWithDependencies {
-    fn from(program: Program) -> Self {
-        Self::new(program, HashMap::new())
-    }
-}
-
-/// Generates a proof of the execution of a LEE program inside the privacy preserving execution
-/// circuit.
-pub fn execute_and_prove(
-    pre_states: Vec<AccountWithMetadata>,
-    instruction_data: InstructionData,
-    account_identities: Vec<InputAccountIdentity>,
-    program_with_dependencies: &ProgramWithDependencies,
-) -> Result<(PrivacyPreservingCircuitOutput, Proof), LeeError> {
-    let ProgramWithDependencies {
-        program: initial_program,
-        dependencies,
-    } = program_with_dependencies;
-    let mut env_builder = ExecutorEnv::builder();
-    let mut program_outputs = Vec::new();
-
-    let initial_call = ChainedCall {
-        program_id: initial_program.id(),
-        instruction_data,
-        pre_states,
-        pda_seeds: vec![],
-    };
-
-    let mut chained_calls = VecDeque::from_iter([(initial_call, initial_program, None)]);
-    let mut chain_calls_counter = 0;
-    while let Some((chained_call, program, caller_program_id)) = chained_calls.pop_front() {
-        if chain_calls_counter >= MAX_NUMBER_CHAINED_CALLS {
-            return Err(LeeError::MaxChainedCallsDepthExceeded);
-        }
-
-        let inner_receipt = execute_and_prove_program(
-            program,
-            caller_program_id,
-            &chained_call.pre_states,
-            &chained_call.instruction_data,
-        )?;
-
-        let program_output: ProgramOutput = inner_receipt
-            .journal
-            .decode()
-            .map_err(|e| LeeError::ProgramOutputDeserializationError(e.to_string()))?;
-
-        // TODO: remove clone
-        program_outputs.push(program_output.clone());
-
-        // Prove circuit.
-        env_builder.add_assumption(inner_receipt);
-
-        for new_call in program_output.chained_calls.into_iter().rev() {
-            let next_program = dependencies.get(&new_call.program_id).ok_or(
-                InvalidProgramBehaviorError::UndeclaredProgramDependency {
-                    program_id: new_call.program_id,
-                },
-            )?;
-            chained_calls.push_front((new_call, next_program, Some(chained_call.program_id)));
-        }
-
-        chain_calls_counter = chain_calls_counter
-            .checked_add(1)
-            .expect("we check the max depth at the beginning of the loop");
-    }
-
-    let circuit_input = PrivacyPreservingCircuitInput {
-        program_outputs,
-        account_identities,
-        program_id: program_with_dependencies.program.id(),
-    };
-
-    env_builder.write(&circuit_input).unwrap();
-    let env = env_builder.build().unwrap();
-    let prover = default_prover();
-    let opts = ProverOpts::succinct();
-    let prove_info = prover
-        .prove_with_opts(env, PRIVACY_PRESERVING_CIRCUIT_ELF, &opts)
-        .map_err(|e| LeeError::CircuitProvingError(e.to_string()))?;
-
-    let proof = Proof(borsh::to_vec(&prove_info.receipt.inner)?);
-
-    let circuit_output: PrivacyPreservingCircuitOutput = prove_info
-        .receipt
-        .journal
-        .decode()
-        .map_err(|e| LeeError::CircuitOutputDeserializationError(e.to_string()))?;
-
-    Ok((circuit_output, proof))
-}
-
-fn execute_and_prove_program(
-    program: &Program,
-    caller_program_id: Option<ProgramId>,
-    pre_states: &[AccountWithMetadata],
-    instruction_data: &InstructionData,
-) -> Result<Receipt, LeeError> {
-    // Write inputs to the program
-    let mut env_builder = ExecutorEnv::builder();
-    Program::write_inputs(
-        program.id(),
-        caller_program_id,
-        pre_states,
-        instruction_data,
-        &mut env_builder,
-    )?;
-    let env = env_builder.build().unwrap();
-
-    // Prove the program
-    let prover = default_prover();
-    Ok(prover
-        .prove(env, program.elf())
-        .map_err(|e| LeeError::ProgramProveFailed(e.to_string()))?
-        .receipt)
-}
-
-#[cfg(test)]
-mod tests {
-    #![expect(clippy::shadow_unrelated, reason = "We don't care about it in tests")]
-
-    use lee_core::{
-        Commitment, DUMMY_COMMITMENT_HASH, EncryptedAccountData, EncryptionScheme,
-        EphemeralPublicKey, Nullifier, PrivacyPreservingCircuitOutput, SharedSecretKey,
-        account::{Account, AccountId, AccountWithMetadata, Nonce, data::Data},
-        program::{PdaSeed, PrivateAccountKind},
-    };
-
-    use super::*;
-    use crate::{
-        error::LeeError,
-        privacy_preserving_transaction::circuit::execute_and_prove,
-        program::Program,
-        state::{
-            CommitmentSet,
-            tests::{test_private_account_keys_1, test_private_account_keys_2},
-        },
-    };
-
-    fn decrypt_kind(
-        output: &PrivacyPreservingCircuitOutput,
-        ssk: &SharedSecretKey,
-        idx: usize,
-    ) -> PrivateAccountKind {
-        let (kind, _) = EncryptionScheme::decrypt(
-            &output.encrypted_private_post_states[idx].ciphertext,
-            ssk,
-            &output.new_commitments[idx],
-            u32::try_from(idx).expect("idx fits in u32"),
-        )
-        .unwrap();
-        kind
-    }
-
-    #[test]
-    fn proof_inner_roundtrip() {
-        // `Proof::from_inner(b).into_inner()` must return exactly `b`. Catches
-        // mutations of `into_inner` returning `vec![]`, `vec![0]`, or `vec![1]`,
-        // and of `from_inner` discarding its argument.
-        let bytes = vec![0xDE_u8, 0xAD, 0xBE, 0xEF];
-        assert_eq!(Proof::from_inner(bytes.clone()).into_inner(), bytes);
-        assert!(Proof::from_inner(vec![]).into_inner().is_empty());
-        assert_eq!(Proof::from_inner(vec![0xFF]).into_inner(), vec![0xFF_u8]);
-    }
-
-    #[test]
-    fn prove_privacy_preserving_execution_circuit_public_and_private_pre_accounts() {
-        let recipient_keys = test_private_account_keys_1();
-        let program = crate::test_methods::simple_balance_transfer();
-        let sender = AccountWithMetadata::new(
-            Account {
-                program_owner: program.id(),
-                balance: 100,
-                ..Account::default()
-            },
-            true,
-            AccountId::new([0; 32]),
-        );
-
-        let recipient_account_id = AccountId::for_regular_private_account(&recipient_keys.npk(), 0);
-        let recipient = AccountWithMetadata::new(Account::default(), false, recipient_account_id);
-
-        let balance_to_move: u128 = 37;
-
-        let expected_sender_post = Account {
-            program_owner: program.id(),
-            balance: 100 - balance_to_move,
-            nonce: Nonce::default(),
-            data: Data::default(),
-        };
-
-        let expected_recipient_post = Account {
-            program_owner: program.id(),
-            balance: balance_to_move,
-            nonce: Nonce::private_account_nonce_init(&recipient_account_id),
-            data: Data::default(),
-        };
-
-        let expected_sender_pre = sender.clone();
-
-        let shared_secret =
-            SharedSecretKey::encapsulate_deterministic(&recipient_keys.vpk(), &[0_u8; 32], 0).0;
-
-        let (output, proof) = execute_and_prove(
-            vec![sender, recipient],
-            Program::serialize_instruction(balance_to_move).unwrap(),
-            vec![
-                InputAccountIdentity::Public,
-                InputAccountIdentity::PrivateUnauthorized {
-                    epk: EphemeralPublicKey(Vec::new()),
-                    view_tag: EncryptedAccountData::compute_view_tag(
-                        &recipient_keys.npk(),
-                        &recipient_keys.vpk(),
-                    ),
-                    npk: recipient_keys.npk(),
-                    ssk: shared_secret,
-                    identifier: 0,
-                },
-            ],
-            &crate::test_methods::simple_balance_transfer().into(),
-        )
-        .unwrap();
-
-        assert!(proof.is_valid_for(&output));
-
-        let [sender_pre] = output.public_pre_states.try_into().unwrap();
-        let [sender_post] = output.public_post_states.try_into().unwrap();
-        assert_eq!(sender_pre, expected_sender_pre);
-        assert_eq!(sender_post, expected_sender_post);
-        assert_eq!(output.new_commitments.len(), 1);
-        assert_eq!(output.new_nullifiers.len(), 1);
-        assert_eq!(output.encrypted_private_post_states.len(), 1);
-
-        let (_identifier, recipient_post) = EncryptionScheme::decrypt(
-            &output.encrypted_private_post_states[0].ciphertext,
-            &shared_secret,
-            &output.new_commitments[0],
-            0,
-        )
-        .unwrap();
-        assert_eq!(recipient_post, expected_recipient_post);
-    }
-
-    #[test]
-    fn prove_privacy_preserving_execution_circuit_fully_private() {
-        let program = crate::test_methods::simple_balance_transfer();
-        let sender_keys = test_private_account_keys_1();
-        let recipient_keys = test_private_account_keys_2();
-
-        let sender_nonce = Nonce(0xdead_beef);
-        let sender_pre = AccountWithMetadata::new(
-            Account {
-                balance: 100,
-                nonce: sender_nonce,
-                program_owner: program.id(),
-                data: Data::default(),
-            },
-            true,
-            AccountId::for_regular_private_account(&sender_keys.npk(), 0),
-        );
-        let sender_account_id = AccountId::for_regular_private_account(&sender_keys.npk(), 0);
-        let commitment_sender = Commitment::new(&sender_account_id, &sender_pre.account);
-
-        let recipient_account_id = AccountId::for_regular_private_account(&recipient_keys.npk(), 0);
-        let recipient = AccountWithMetadata::new(Account::default(), false, recipient_account_id);
-        let balance_to_move: u128 = 37;
-
-        let mut commitment_set = CommitmentSet::with_capacity(2);
-        commitment_set.extend(std::slice::from_ref(&commitment_sender));
-        let expected_new_nullifiers = vec![
-            (
-                Nullifier::for_account_update(&commitment_sender, &sender_keys.nsk),
-                commitment_set.digest(),
-            ),
-            (
-                Nullifier::for_account_initialization(&recipient_account_id),
-                DUMMY_COMMITMENT_HASH,
-            ),
-        ];
-
-        let program = crate::test_methods::simple_balance_transfer();
-
-        let expected_private_account_1 = Account {
-            program_owner: program.id(),
-            balance: 100 - balance_to_move,
-            nonce: sender_nonce.private_account_nonce_increment(&sender_keys.nsk),
-            ..Default::default()
-        };
-        let expected_private_account_2 = Account {
-            program_owner: program.id(),
-            balance: balance_to_move,
-            nonce: Nonce::private_account_nonce_init(&recipient_account_id),
-            ..Default::default()
-        };
-        let expected_new_commitments = vec![
-            Commitment::new(&sender_account_id, &expected_private_account_1),
-            Commitment::new(&recipient_account_id, &expected_private_account_2),
-        ];
-
-        let shared_secret_1 =
-            SharedSecretKey::encapsulate_deterministic(&sender_keys.vpk(), &[0_u8; 32], 0).0;
-
-        let shared_secret_2 =
-            SharedSecretKey::encapsulate_deterministic(&recipient_keys.vpk(), &[0_u8; 32], 1).0;
-
-        let (output, proof) = execute_and_prove(
-            vec![sender_pre, recipient],
-            Program::serialize_instruction(balance_to_move).unwrap(),
-            vec![
-                InputAccountIdentity::PrivateAuthorizedUpdate {
-                    epk: EphemeralPublicKey(Vec::new()),
-                    view_tag: EncryptedAccountData::compute_view_tag(
-                        &sender_keys.npk(),
-                        &sender_keys.vpk(),
-                    ),
-                    ssk: shared_secret_1,
-                    nsk: sender_keys.nsk,
-                    membership_proof: commitment_set
-                        .get_proof_for(&commitment_sender)
-                        .expect("sender's commitment must be in the set"),
-                    identifier: 0,
-                },
-                InputAccountIdentity::PrivateUnauthorized {
-                    epk: EphemeralPublicKey(Vec::new()),
-                    view_tag: EncryptedAccountData::compute_view_tag(
-                        &recipient_keys.npk(),
-                        &recipient_keys.vpk(),
-                    ),
-                    npk: recipient_keys.npk(),
-                    ssk: shared_secret_2,
-                    identifier: 0,
-                },
-            ],
-            &program.into(),
-        )
-        .unwrap();
-
-        assert!(proof.is_valid_for(&output));
-        assert!(output.public_pre_states.is_empty());
-        assert!(output.public_post_states.is_empty());
-        assert_eq!(output.new_commitments, expected_new_commitments);
-        assert_eq!(output.new_nullifiers, expected_new_nullifiers);
-        assert_eq!(output.encrypted_private_post_states.len(), 2);
-
-        let (_identifier, sender_post) = EncryptionScheme::decrypt(
-            &output.encrypted_private_post_states[0].ciphertext,
-            &shared_secret_1,
-            &expected_new_commitments[0],
-            0,
-        )
-        .unwrap();
-        assert_eq!(sender_post, expected_private_account_1);
-
-        let (_identifier, recipient_post) = EncryptionScheme::decrypt(
-            &output.encrypted_private_post_states[1].ciphertext,
-            &shared_secret_2,
-            &expected_new_commitments[1],
-            1,
-        )
-        .unwrap();
-        assert_eq!(recipient_post, expected_private_account_2);
-    }
-
-    #[test]
-    fn circuit_fails_when_chained_validity_windows_have_empty_intersection() {
-        let account_keys = test_private_account_keys_1();
-        let pre = AccountWithMetadata::new(
-            Account::default(),
-            false,
-            AccountId::for_regular_private_account(&account_keys.npk(), 0),
-        );
-
-        let validity_window_chain_caller = crate::test_methods::validity_window_chain_caller();
-        let validity_window = crate::test_methods::validity_window();
-
-        let instruction = Program::serialize_instruction((
-            Some(1_u64),
-            Some(4_u64),
-            validity_window.id(),
-            Some(4_u64),
-            Some(7_u64),
-        ))
-        .unwrap();
-
-        let shared_secret =
-            SharedSecretKey::encapsulate_deterministic(&account_keys.vpk(), &[0_u8; 32], 0).0;
-
-        let program_with_deps = ProgramWithDependencies::new(
-            validity_window_chain_caller,
-            [(validity_window.id(), validity_window)].into(),
-        );
-
-        let result = execute_and_prove(
-            vec![pre],
-            instruction,
-            vec![InputAccountIdentity::PrivateUnauthorized {
-                epk: EphemeralPublicKey(Vec::new()),
-                view_tag: EncryptedAccountData::compute_view_tag(
-                    &account_keys.npk(),
-                    &account_keys.vpk(),
-                ),
-                npk: account_keys.npk(),
-                ssk: shared_secret,
-                identifier: 0,
-            }],
-            &program_with_deps,
-        );
-
-        assert!(matches!(result, Err(LeeError::CircuitProvingError(_))));
-    }
-
-    /// A private PDA claimed with a non-default identifier produces a ciphertext that decrypts
-    /// to `PrivateAccountKind::Pda` carrying the correct `(program_id, seed, identifier)`.
-    #[test]
-    fn private_pda_claim_with_custom_identifier_encrypts_correct_kind() {
-        let program = crate::test_methods::pda_claimer();
-        let keys = test_private_account_keys_1();
-        let npk = keys.npk();
-        let seed = PdaSeed::new([42; 32]);
-        let identifier: u128 = 99;
-        let shared_secret =
-            SharedSecretKey::encapsulate_deterministic(&keys.vpk(), &[0_u8; 32], 0).0;
-
-        let account_id = AccountId::for_private_pda(&program.id(), &seed, &npk, identifier);
-        let pre_state = AccountWithMetadata::new(Account::default(), false, account_id);
-
-        let (output, _proof) = execute_and_prove(
-            vec![pre_state],
-            Program::serialize_instruction(seed).unwrap(),
-            vec![InputAccountIdentity::PrivatePdaInit {
-                epk: EphemeralPublicKey(Vec::new()),
-                view_tag: EncryptedAccountData::compute_view_tag(&npk, &keys.vpk()),
-                npk,
-                ssk: shared_secret,
-                identifier,
-                seed: None,
-            }],
-            &program.clone().into(),
-        )
-        .unwrap();
-
-        assert_eq!(
-            decrypt_kind(&output, &shared_secret, 0),
-            PrivateAccountKind::Pda {
-                program_id: program.id(),
-                seed,
-                identifier
-            },
-        );
-    }
-
-    /// PDA init: initializes a new PDA under `simple_balance_transfer`'s ownership.
-    /// The `simple_transfer_proxy` program chains to `simple_balance_transfer` with `pda_seeds`
-    /// to establish authorization and the private PDA binding.
-    #[test]
-    fn private_pda_init() {
-        let program = crate::test_methods::simple_transfer_proxy();
-        let simple_transfer = crate::test_methods::simple_balance_transfer();
-        let keys = test_private_account_keys_1();
-        let npk = keys.npk();
-        let seed = PdaSeed::new([42; 32]);
-        let shared_secret_pda =
-            SharedSecretKey::encapsulate_deterministic(&keys.vpk(), &[0_u8; 32], 0).0;
-
-        // PDA (new, private PDA)
-        let pda_id = AccountId::for_private_pda(&program.id(), &seed, &npk, 0);
-        let pda_pre = AccountWithMetadata::new(Account::default(), false, pda_id);
-
-        let auth_id = simple_transfer.id();
-        let program_with_deps =
-            ProgramWithDependencies::new(program, [(auth_id, simple_transfer)].into());
-
-        // is_withdraw=false triggers init path (1 pre-state)
-        let instruction = Program::serialize_instruction((seed, auth_id, 0_u128, false)).unwrap();
-
-        let result = execute_and_prove(
-            vec![pda_pre],
-            instruction,
-            vec![InputAccountIdentity::PrivatePdaInit {
-                epk: EphemeralPublicKey(Vec::new()),
-                view_tag: EncryptedAccountData::compute_view_tag(&npk, &keys.vpk()),
-                npk,
-                ssk: shared_secret_pda,
-                identifier: 0,
-                seed: None,
-            }],
-            &program_with_deps,
-        );
-
-        let (output, _proof) = result.expect("PDA init should succeed");
-        assert_eq!(output.new_commitments.len(), 1);
-    }
-
-    /// PDA withdraw: chains to `simple_balance_transfer` to move balance from PDA to recipient.
-    /// Uses a default PDA (amount=0) because testing with a pre-funded PDA requires a
-    /// two-tx sequence with membership proofs.
-    #[test]
-    fn private_pda_withdraw() {
-        let program = crate::test_methods::simple_transfer_proxy();
-        let simple_transfer = crate::test_methods::simple_balance_transfer();
-        let keys = test_private_account_keys_1();
-        let npk = keys.npk();
-        let seed = PdaSeed::new([42; 32]);
-        let shared_secret_pda =
-            SharedSecretKey::encapsulate_deterministic(&keys.vpk(), &[0_u8; 32], 0).0;
-
-        // PDA (new, private PDA)
-        let pda_id = AccountId::for_private_pda(&program.id(), &seed, &npk, 0);
-        let pda_pre = AccountWithMetadata::new(Account::default(), false, pda_id);
-
-        // Recipient (public)
-        let recipient_id = AccountId::new([88; 32]);
-        let recipient_pre = AccountWithMetadata::new(
-            Account {
-                program_owner: simple_transfer.id(),
-                balance: 10000,
-                ..Account::default()
-            },
-            true,
-            recipient_id,
-        );
-
-        let auth_id = simple_transfer.id();
-        let program_with_deps =
-            ProgramWithDependencies::new(program, [(auth_id, simple_transfer)].into());
-
-        // is_withdraw=true, amount=0 (PDA has no balance yet)
-        let instruction = Program::serialize_instruction((seed, auth_id, 0_u128, true)).unwrap();
-
-        let result = execute_and_prove(
-            vec![pda_pre, recipient_pre],
-            instruction,
-            vec![
-                InputAccountIdentity::PrivatePdaInit {
-                    epk: EphemeralPublicKey(Vec::new()),
-                    view_tag: EncryptedAccountData::compute_view_tag(&npk, &keys.vpk()),
-                    npk,
-                    ssk: shared_secret_pda,
-                    identifier: 0,
-                    seed: None,
-                },
-                InputAccountIdentity::Public,
-            ],
-            &program_with_deps,
-        );
-
-        let (output, _proof) = result.expect("PDA withdraw should succeed");
-        assert_eq!(output.new_commitments.len(), 1);
-    }
-
-    /// Shared regular private account: receives funds via `authenticated_transfer` directly,
-    /// no custom program needed. This demonstrates the non-PDA shared account flow where
-    /// keys are derived from GMS via `derive_keys_for_shared_account`. The shared account
-    /// uses the standard unauthorized private account path and works with auth-transfer's
-    /// transfer path like any other private account.
-    #[test]
-    fn shared_account_receives_via_simple_transfer() {
-        let program = crate::test_methods::simple_balance_transfer();
-        let shared_keys = test_private_account_keys_1();
-        let shared_npk = shared_keys.npk();
-        let shared_identifier: u128 = 42;
-        let shared_secret =
-            SharedSecretKey::encapsulate_deterministic(&shared_keys.vpk(), &[0_u8; 32], 0).0;
-
-        // Sender: public account with balance, owned by auth-transfer
-        let sender_id = AccountId::new([99; 32]);
-        let sender = AccountWithMetadata::new(
-            Account {
-                program_owner: program.id(),
-                balance: 1000,
-                ..Account::default()
-            },
-            true,
-            sender_id,
-        );
-
-        // Recipient: shared private account (new, unauthorized)
-        let shared_account_id = AccountId::from((&shared_npk, shared_identifier));
-        let recipient = AccountWithMetadata::new(Account::default(), false, shared_account_id);
-
-        let balance_to_move: u128 = 100;
-        let instruction = Program::serialize_instruction(balance_to_move).unwrap();
-
-        let result = execute_and_prove(
-            vec![sender, recipient],
-            instruction,
-            vec![
-                InputAccountIdentity::Public,
-                InputAccountIdentity::PrivateUnauthorized {
-                    epk: EphemeralPublicKey(Vec::new()),
-                    view_tag: EncryptedAccountData::compute_view_tag(
-                        &shared_npk,
-                        &shared_keys.vpk(),
-                    ),
-                    npk: shared_npk,
-                    ssk: shared_secret,
-                    identifier: shared_identifier,
-                },
-            ],
-            &program.into(),
-        );
-
-        let (output, _proof) = result.expect("shared account receive should succeed");
-        // Sender is public (no commitment), recipient is private (1 commitment)
-        assert_eq!(output.new_commitments.len(), 1);
-    }
-
-    /// `PrivateAuthorizedInit` with a non-default identifier produces a ciphertext that decrypts
-    /// to `PrivateAccountKind::Regular` carrying the correct identifier.
-    #[test]
-    fn private_authorized_init_encrypts_regular_kind_with_identifier() {
-        let program = crate::test_methods::claimer();
-        let keys = test_private_account_keys_1();
-        let identifier: u128 = 99;
-        let ssk = SharedSecretKey::encapsulate_deterministic(&keys.vpk(), &[0_u8; 32], 0).0;
-        let account_id = AccountId::for_regular_private_account(&keys.npk(), identifier);
-        let pre = AccountWithMetadata::new(Account::default(), true, account_id);
-
-        let (output, _) = execute_and_prove(
-            vec![pre],
-            Program::serialize_instruction(()).unwrap(),
-            vec![InputAccountIdentity::PrivateAuthorizedInit {
-                epk: EphemeralPublicKey(Vec::new()),
-                view_tag: EncryptedAccountData::compute_view_tag(&keys.npk(), &keys.vpk()),
-                ssk,
-                nsk: keys.nsk,
-                identifier,
-            }],
-            &program.into(),
-        )
-        .unwrap();
-
-        assert_eq!(
-            decrypt_kind(&output, &ssk, 0),
-            PrivateAccountKind::Regular(identifier)
-        );
-    }
-
-    /// `PrivateUnauthorized` with a non-default identifier produces a ciphertext that decrypts
-    /// to `PrivateAccountKind::Regular` carrying the correct identifier.
-    #[test]
-    fn private_unauthorized_init_encrypts_regular_kind_with_identifier() {
-        let program = crate::test_methods::claimer();
-        let keys = test_private_account_keys_1();
-        let identifier: u128 = 99;
-        let ssk = SharedSecretKey::encapsulate_deterministic(&keys.vpk(), &[0_u8; 32], 0).0;
-        let recipient_id = AccountId::for_regular_private_account(&keys.npk(), identifier);
-        let recipient = AccountWithMetadata::new(Account::default(), false, recipient_id);
-
-        let (output, _) = execute_and_prove(
-            vec![recipient],
-            Program::serialize_instruction(()).unwrap(),
-            vec![InputAccountIdentity::PrivateUnauthorized {
-                epk: EphemeralPublicKey(Vec::new()),
-                view_tag: EncryptedAccountData::compute_view_tag(&keys.npk(), &keys.vpk()),
-                npk: keys.npk(),
-                ssk,
-                identifier,
-            }],
-            &program.into(),
-        )
-        .unwrap();
-
-        assert_eq!(
-            decrypt_kind(&output, &ssk, 0),
-            PrivateAccountKind::Regular(identifier)
-        );
-    }
-
-    /// `PrivateAuthorizedUpdate` with a non-default identifier produces a ciphertext that decrypts
-    /// to `PrivateAccountKind::Regular` carrying the correct identifier.
-    #[test]
-    fn private_authorized_update_encrypts_regular_kind_with_identifier() {
-        let program = crate::test_methods::noop();
-        let keys = test_private_account_keys_1();
-        let identifier: u128 = 99;
-        let ssk = SharedSecretKey::encapsulate_deterministic(&keys.vpk(), &[0_u8; 32], 0).0;
-        let account_id = AccountId::for_regular_private_account(&keys.npk(), identifier);
-        let account = Account {
-            program_owner: program.id(),
-            balance: 1,
-            ..Account::default()
-        };
-        let commitment = Commitment::new(&account_id, &account);
-        let mut commitment_set = CommitmentSet::with_capacity(1);
-        commitment_set.extend(std::slice::from_ref(&commitment));
-
-        let sender = AccountWithMetadata::new(account, true, account_id);
-
-        let (output, _) = execute_and_prove(
-            vec![sender],
-            Program::serialize_instruction(()).unwrap(),
-            vec![InputAccountIdentity::PrivateAuthorizedUpdate {
-                epk: EphemeralPublicKey(Vec::new()),
-                view_tag: EncryptedAccountData::compute_view_tag(&keys.npk(), &keys.vpk()),
-                ssk,
-                nsk: keys.nsk,
-                membership_proof: commitment_set.get_proof_for(&commitment).unwrap(),
-                identifier,
-            }],
-            &program.into(),
-        )
-        .unwrap();
-
-        assert_eq!(
-            decrypt_kind(&output, &ssk, 0),
-            PrivateAccountKind::Regular(identifier)
-        );
-    }
-
-    /// `PrivatePdaUpdate` with a non-default identifier produces a ciphertext that decrypts
-    /// to `PrivateAccountKind::Pda` carrying the correct `(program_id, seed, identifier)`.
-    #[test]
-    fn private_pda_update_encrypts_pda_kind_with_identifier() {
-        let program = crate::test_methods::pda_spend_proxy();
-        let simple_transfer = crate::test_methods::simple_balance_transfer();
-        let keys = test_private_account_keys_1();
-        let npk = keys.npk();
-        let seed = PdaSeed::new([42; 32]);
-        let identifier: u128 = 99;
-        let ssk = SharedSecretKey::encapsulate_deterministic(&keys.vpk(), &[0_u8; 32], 0).0;
-
-        let simple_transfer_id = simple_transfer.id();
-        let pda_id = AccountId::for_private_pda(&program.id(), &seed, &npk, identifier);
-        let pda_account = Account {
-            program_owner: simple_transfer_id,
-            balance: 1,
-            ..Account::default()
-        };
-        let pda_commitment = Commitment::new(&pda_id, &pda_account);
-        let mut commitment_set = CommitmentSet::with_capacity(1);
-        commitment_set.extend(std::slice::from_ref(&pda_commitment));
-
-        let pda_pre = AccountWithMetadata::new(pda_account, true, pda_id);
-        let recipient_pre =
-            AccountWithMetadata::new(Account::default(), true, AccountId::new([0; 32]));
-
-        let program_with_deps = ProgramWithDependencies::new(
-            program.clone(),
-            [(simple_transfer_id, simple_transfer)].into(),
-        );
-
-        let (output, _) = execute_and_prove(
-            vec![pda_pre, recipient_pre],
-            Program::serialize_instruction((seed, 1_u128, simple_transfer_id, false)).unwrap(),
-            vec![
-                InputAccountIdentity::PrivatePdaUpdate {
-                    epk: EphemeralPublicKey(Vec::new()),
-                    view_tag: EncryptedAccountData::compute_view_tag(&npk, &keys.vpk()),
-                    ssk,
-                    nsk: keys.nsk,
-                    membership_proof: commitment_set.get_proof_for(&pda_commitment).unwrap(),
-                    identifier,
-                    seed: None,
-                },
-                InputAccountIdentity::Public,
-            ],
-            &program_with_deps,
-        )
-        .unwrap();
-
-        assert_eq!(
-            decrypt_kind(&output, &ssk, 0),
-            PrivateAccountKind::Pda {
-                program_id: program.id(),
-                seed,
-                identifier
-            },
-        );
-    }
-
-    #[test]
-    fn private_pda_init_identifier_mismatch_fails() {
-        let program = crate::test_methods::pda_claimer();
-        let keys = test_private_account_keys_1();
-        let npk = keys.npk();
-        let seed = PdaSeed::new([42; 32]);
-        let shared_secret =
-            SharedSecretKey::encapsulate_deterministic(&keys.vpk(), &[0_u8; 32], 0).0;
-
-        let account_id = AccountId::for_private_pda(&program.id(), &seed, &npk, 5);
-        let pre_state = AccountWithMetadata::new(Account::default(), false, account_id);
-
-        let result = execute_and_prove(
-            vec![pre_state],
-            Program::serialize_instruction(seed).unwrap(),
-            vec![InputAccountIdentity::PrivatePdaInit {
-                epk: EphemeralPublicKey(Vec::new()),
-                view_tag: EncryptedAccountData::compute_view_tag(&npk, &keys.vpk()),
-                npk,
-                ssk: shared_secret,
-                identifier: 99,
-                seed: None,
-            }],
-            &program.into(),
-        );
-
-        assert!(matches!(result, Err(LeeError::CircuitProvingError(_))));
-    }
-
-    #[test]
-    fn private_pda_update_identifier_mismatch_fails() {
-        let program = crate::test_methods::pda_spend_proxy();
-        let simple_transfer = crate::test_methods::simple_balance_transfer();
-        let keys = test_private_account_keys_1();
-        let npk = keys.npk();
-        let seed = PdaSeed::new([42; 32]);
-        let ssk = SharedSecretKey::encapsulate_deterministic(&keys.vpk(), &[0_u8; 32], 0).0;
-
-        let simple_transfer_id = simple_transfer.id();
-        let pda_id = AccountId::for_private_pda(&program.id(), &seed, &npk, 5);
-        let pda_account = Account {
-            program_owner: simple_transfer_id,
-            balance: 1,
-            ..Account::default()
-        };
-        let pda_commitment = Commitment::new(&pda_id, &pda_account);
-        let mut commitment_set = CommitmentSet::with_capacity(1);
-        commitment_set.extend(std::slice::from_ref(&pda_commitment));
-
-        let pda_pre = AccountWithMetadata::new(pda_account, true, pda_id);
-        let recipient_pre =
-            AccountWithMetadata::new(Account::default(), true, AccountId::new([0; 32]));
-
-        let program_with_deps =
-            ProgramWithDependencies::new(program, [(simple_transfer_id, simple_transfer)].into());
-
-        let result = execute_and_prove(
-            vec![pda_pre, recipient_pre],
-            Program::serialize_instruction((seed, 1_u128, simple_transfer_id, false)).unwrap(),
-            vec![
-                InputAccountIdentity::PrivatePdaUpdate {
-                    epk: EphemeralPublicKey(Vec::new()),
-                    view_tag: EncryptedAccountData::compute_view_tag(&npk, &keys.vpk()),
-                    ssk,
-                    nsk: keys.nsk,
-                    membership_proof: commitment_set.get_proof_for(&pda_commitment).unwrap(),
-                    identifier: 99,
-                    seed: None,
-                },
-                InputAccountIdentity::Public,
-            ],
-            &program_with_deps,
-        );
-
-        assert!(matches!(result, Err(LeeError::CircuitProvingError(_))));
-    }
-}

lee/state_machine/src/privacy_preserving_transaction/circuit/mod.rs

@@ -0,0 +1,177 @@
+use std::collections::{HashMap, VecDeque};
+
+use borsh::{BorshDeserialize, BorshSerialize};
+use lee_core::{
+    InputAccountIdentity, PrivacyPreservingCircuitInput, PrivacyPreservingCircuitOutput,
+    account::AccountWithMetadata,
+    program::{ChainedCall, InstructionData, ProgramId, ProgramOutput},
+};
+use risc0_zkvm::{ExecutorEnv, InnerReceipt, ProverOpts, Receipt, default_prover};
+
+use crate::{
+    PRIVACY_PRESERVING_CIRCUIT_ELF, PRIVACY_PRESERVING_CIRCUIT_ID,
+    error::{InvalidProgramBehaviorError, LeeError},
+    program::Program,
+    state::MAX_NUMBER_CHAINED_CALLS,
+};
+
+/// Proof of the privacy preserving execution circuit.
+#[derive(Debug, Clone, PartialEq, Eq, BorshSerialize, BorshDeserialize)]
+pub struct Proof(pub(crate) Vec<u8>);
+
+impl Proof {
+    #[must_use]
+    pub fn into_inner(self) -> Vec<u8> {
+        self.0
+    }
+
+    #[must_use]
+    pub const fn from_inner(inner: Vec<u8>) -> Self {
+        Self(inner)
+    }
+
+    pub(crate) fn is_valid_for(&self, circuit_output: &PrivacyPreservingCircuitOutput) -> bool {
+        let Ok(inner) = borsh::from_slice::<InnerReceipt>(&self.0) else {
+            return false;
+        };
+        let receipt = Receipt::new(inner, circuit_output.to_bytes());
+        receipt.verify(PRIVACY_PRESERVING_CIRCUIT_ID).is_ok()
+    }
+}
+
+#[derive(Clone)]
+pub struct ProgramWithDependencies {
+    pub program: Program,
+    // TODO: avoid having a copy of the bytecode of each dependency.
+    pub dependencies: HashMap<ProgramId, Program>,
+}
+
+impl ProgramWithDependencies {
+    #[must_use]
+    pub const fn new(program: Program, dependencies: HashMap<ProgramId, Program>) -> Self {
+        Self {
+            program,
+            dependencies,
+        }
+    }
+}
+
+impl From<Program> for ProgramWithDependencies {
+    fn from(program: Program) -> Self {
+        Self::new(program, HashMap::new())
+    }
+}
+
+/// Generates a proof of the execution of a LEE program inside the privacy preserving execution
+/// circuit.
+pub fn execute_and_prove(
+    pre_states: Vec<AccountWithMetadata>,
+    instruction_data: InstructionData,
+    account_identities: Vec<InputAccountIdentity>,
+    program_with_dependencies: &ProgramWithDependencies,
+) -> Result<(PrivacyPreservingCircuitOutput, Proof), LeeError> {
+    let ProgramWithDependencies {
+        program: initial_program,
+        dependencies,
+    } = program_with_dependencies;
+    let mut env_builder = ExecutorEnv::builder();
+    let mut program_outputs = Vec::new();
+
+    let initial_call = ChainedCall {
+        program_id: initial_program.id(),
+        instruction_data,
+        pre_states,
+        pda_seeds: vec![],
+    };
+
+    let mut chained_calls = VecDeque::from_iter([(initial_call, initial_program, None)]);
+    let mut chain_calls_counter = 0;
+    while let Some((chained_call, program, caller_program_id)) = chained_calls.pop_front() {
+        if chain_calls_counter >= MAX_NUMBER_CHAINED_CALLS {
+            return Err(LeeError::MaxChainedCallsDepthExceeded);
+        }
+
+        let inner_receipt = execute_and_prove_program(
+            program,
+            caller_program_id,
+            &chained_call.pre_states,
+            &chained_call.instruction_data,
+        )?;
+
+        let program_output: ProgramOutput = inner_receipt
+            .journal
+            .decode()
+            .map_err(|e| LeeError::ProgramOutputDeserializationError(e.to_string()))?;
+
+        // TODO: remove clone
+        program_outputs.push(program_output.clone());
+
+        // Prove circuit.
+        env_builder.add_assumption(inner_receipt);
+
+        for new_call in program_output.chained_calls.into_iter().rev() {
+            let next_program = dependencies.get(&new_call.program_id).ok_or(
+                InvalidProgramBehaviorError::UndeclaredProgramDependency {
+                    program_id: new_call.program_id,
+                },
+            )?;
+            chained_calls.push_front((new_call, next_program, Some(chained_call.program_id)));
+        }
+
+        chain_calls_counter = chain_calls_counter
+            .checked_add(1)
+            .expect("we check the max depth at the beginning of the loop");
+    }
+
+    let circuit_input = PrivacyPreservingCircuitInput {
+        program_outputs,
+        account_identities,
+        program_id: program_with_dependencies.program.id(),
+    };
+
+    env_builder.write(&circuit_input).unwrap();
+    let env = env_builder.build().unwrap();
+    let prover = default_prover();
+    let opts = ProverOpts::succinct();
+    let prove_info = prover
+        .prove_with_opts(env, PRIVACY_PRESERVING_CIRCUIT_ELF, &opts)
+        .map_err(|e| LeeError::CircuitProvingError(e.to_string()))?;
+
+    let proof = Proof(borsh::to_vec(&prove_info.receipt.inner)?);
+
+    let circuit_output: PrivacyPreservingCircuitOutput = prove_info
+        .receipt
+        .journal
+        .decode()
+        .map_err(|e| LeeError::CircuitOutputDeserializationError(e.to_string()))?;
+
+    Ok((circuit_output, proof))
+}
+
+fn execute_and_prove_program(
+    program: &Program,
+    caller_program_id: Option<ProgramId>,
+    pre_states: &[AccountWithMetadata],
+    instruction_data: &InstructionData,
+) -> Result<Receipt, LeeError> {
+    // Write inputs to the program
+    let mut env_builder = ExecutorEnv::builder();
+    Program::write_inputs(
+        program.id(),
+        caller_program_id,
+        pre_states,
+        instruction_data,
+        &mut env_builder,
+    )?;
+    let env = env_builder.build().unwrap();
+
+    // Prove the program
+    let prover = default_prover();
+    Ok(prover
+        .prove(env, program.elf())
+        .map_err(|e| LeeError::ProgramProveFailed(e.to_string()))?
+        .receipt)
+}
+
+#[cfg(test)]
+mod tests;

lee/state_machine/src/privacy_preserving_transaction/circuit/tests.rs

@@ -0,0 +1,731 @@
+#![expect(clippy::shadow_unrelated, reason = "We don't care about it in tests")]
+
+use lee_core::{
+    Commitment, DUMMY_COMMITMENT_HASH, EncryptedAccountData, EncryptionScheme, EphemeralPublicKey,
+    Nullifier, PrivacyPreservingCircuitOutput, SharedSecretKey,
+    account::{Account, AccountId, AccountWithMetadata, Nonce, data::Data},
+    program::{PdaSeed, PrivateAccountKind},
+};
+
+use super::*;
+use crate::{
+    error::LeeError,
+    privacy_preserving_transaction::circuit::execute_and_prove,
+    program::Program,
+    state::{
+        CommitmentSet,
+        tests::{test_private_account_keys_1, test_private_account_keys_2},
+    },
+};
+
+fn decrypt_kind(
+    output: &PrivacyPreservingCircuitOutput,
+    ssk: &SharedSecretKey,
+    idx: usize,
+) -> PrivateAccountKind {
+    let (kind, _) = EncryptionScheme::decrypt(
+        &output.encrypted_private_post_states[idx].ciphertext,
+        ssk,
+        &output.new_commitments[idx],
+        u32::try_from(idx).expect("idx fits in u32"),
+    )
+    .unwrap();
+    kind
+}
+
+#[test]
+fn proof_inner_roundtrip() {
+    // `Proof::from_inner(b).into_inner()` must return exactly `b`. Catches
+    // mutations of `into_inner` returning `vec![]`, `vec![0]`, or `vec![1]`,
+    // and of `from_inner` discarding its argument.
+    let bytes = vec![0xDE_u8, 0xAD, 0xBE, 0xEF];
+    assert_eq!(Proof::from_inner(bytes.clone()).into_inner(), bytes);
+    assert!(Proof::from_inner(vec![]).into_inner().is_empty());
+    assert_eq!(Proof::from_inner(vec![0xFF]).into_inner(), vec![0xFF_u8]);
+}
+
+#[test]
+fn prove_privacy_preserving_execution_circuit_public_and_private_pre_accounts() {
+    let recipient_keys = test_private_account_keys_1();
+    let program = crate::test_methods::simple_balance_transfer();
+    let sender = AccountWithMetadata::new(
+        Account {
+            program_owner: program.id(),
+            balance: 100,
+            ..Account::default()
+        },
+        true,
+        AccountId::new([0; 32]),
+    );
+
+    let recipient_account_id = AccountId::for_regular_private_account(&recipient_keys.npk(), 0);
+    let recipient = AccountWithMetadata::new(Account::default(), false, recipient_account_id);
+
+    let balance_to_move: u128 = 37;
+
+    let expected_sender_post = Account {
+        program_owner: program.id(),
+        balance: 100 - balance_to_move,
+        nonce: Nonce::default(),
+        data: Data::default(),
+    };
+
+    let expected_recipient_post = Account {
+        program_owner: program.id(),
+        balance: balance_to_move,
+        nonce: Nonce::private_account_nonce_init(&recipient_account_id),
+        data: Data::default(),
+    };
+
+    let expected_sender_pre = sender.clone();
+
+    let shared_secret =
+        SharedSecretKey::encapsulate_deterministic(&recipient_keys.vpk(), &[0_u8; 32], 0).0;
+
+    let (output, proof) = execute_and_prove(
+        vec![sender, recipient],
+        Program::serialize_instruction(balance_to_move).unwrap(),
+        vec![
+            InputAccountIdentity::Public,
+            InputAccountIdentity::PrivateUnauthorized {
+                epk: EphemeralPublicKey(Vec::new()),
+                view_tag: EncryptedAccountData::compute_view_tag(
+                    &recipient_keys.npk(),
+                    &recipient_keys.vpk(),
+                ),
+                npk: recipient_keys.npk(),
+                ssk: shared_secret,
+                identifier: 0,
+                commitment_root: DUMMY_COMMITMENT_HASH,
+            },
+        ],
+        &crate::test_methods::simple_balance_transfer().into(),
+    )
+    .unwrap();
+
+    assert!(proof.is_valid_for(&output));
+
+    let [sender_pre] = output.public_pre_states.try_into().unwrap();
+    let [sender_post] = output.public_post_states.try_into().unwrap();
+    assert_eq!(sender_pre, expected_sender_pre);
+    assert_eq!(sender_post, expected_sender_post);
+    assert_eq!(output.new_commitments.len(), 1);
+    assert_eq!(output.new_nullifiers.len(), 1);
+    assert_eq!(output.encrypted_private_post_states.len(), 1);
+
+    let (_identifier, recipient_post) = EncryptionScheme::decrypt(
+        &output.encrypted_private_post_states[0].ciphertext,
+        &shared_secret,
+        &output.new_commitments[0],
+        0,
+    )
+    .unwrap();
+    assert_eq!(recipient_post, expected_recipient_post);
+}
+
+#[test]
+fn prove_privacy_preserving_execution_circuit_fully_private() {
+    let program = crate::test_methods::simple_balance_transfer();
+    let sender_keys = test_private_account_keys_1();
+    let recipient_keys = test_private_account_keys_2();
+
+    let sender_nonce = Nonce(0xdead_beef);
+    let sender_pre = AccountWithMetadata::new(
+        Account {
+            balance: 100,
+            nonce: sender_nonce,
+            program_owner: program.id(),
+            data: Data::default(),
+        },
+        true,
+        AccountId::for_regular_private_account(&sender_keys.npk(), 0),
+    );
+    let sender_account_id = AccountId::for_regular_private_account(&sender_keys.npk(), 0);
+    let commitment_sender = Commitment::new(&sender_account_id, &sender_pre.account);
+
+    let recipient_account_id = AccountId::for_regular_private_account(&recipient_keys.npk(), 0);
+    let recipient = AccountWithMetadata::new(Account::default(), false, recipient_account_id);
+    let balance_to_move: u128 = 37;
+
+    let mut commitment_set = CommitmentSet::with_capacity(2);
+    commitment_set.extend(std::slice::from_ref(&commitment_sender));
+    let expected_new_nullifiers = vec![
+        (
+            Nullifier::for_account_update(&commitment_sender, &sender_keys.nsk),
+            commitment_set.digest(),
+        ),
+        (
+            Nullifier::for_account_initialization(&recipient_account_id),
+            DUMMY_COMMITMENT_HASH,
+        ),
+    ];
+
+    let program = crate::test_methods::simple_balance_transfer();
+
+    let expected_private_account_1 = Account {
+        program_owner: program.id(),
+        balance: 100 - balance_to_move,
+        nonce: sender_nonce.private_account_nonce_increment(&sender_keys.nsk),
+        ..Default::default()
+    };
+    let expected_private_account_2 = Account {
+        program_owner: program.id(),
+        balance: balance_to_move,
+        nonce: Nonce::private_account_nonce_init(&recipient_account_id),
+        ..Default::default()
+    };
+    let expected_new_commitments = vec![
+        Commitment::new(&sender_account_id, &expected_private_account_1),
+        Commitment::new(&recipient_account_id, &expected_private_account_2),
+    ];
+
+    let shared_secret_1 =
+        SharedSecretKey::encapsulate_deterministic(&sender_keys.vpk(), &[0_u8; 32], 0).0;
+
+    let shared_secret_2 =
+        SharedSecretKey::encapsulate_deterministic(&recipient_keys.vpk(), &[0_u8; 32], 1).0;
+
+    let (output, proof) = execute_and_prove(
+        vec![sender_pre, recipient],
+        Program::serialize_instruction(balance_to_move).unwrap(),
+        vec![
+            InputAccountIdentity::PrivateAuthorizedUpdate {
+                epk: EphemeralPublicKey(Vec::new()),
+                view_tag: EncryptedAccountData::compute_view_tag(
+                    &sender_keys.npk(),
+                    &sender_keys.vpk(),
+                ),
+                ssk: shared_secret_1,
+                nsk: sender_keys.nsk,
+                membership_proof: commitment_set
+                    .get_proof_for(&commitment_sender)
+                    .expect("sender's commitment must be in the set"),
+                identifier: 0,
+            },
+            InputAccountIdentity::PrivateUnauthorized {
+                epk: EphemeralPublicKey(Vec::new()),
+                view_tag: EncryptedAccountData::compute_view_tag(
+                    &recipient_keys.npk(),
+                    &recipient_keys.vpk(),
+                ),
+                npk: recipient_keys.npk(),
+                ssk: shared_secret_2,
+                identifier: 0,
+                commitment_root: DUMMY_COMMITMENT_HASH,
+            },
+        ],
+        &program.into(),
+    )
+    .unwrap();
+
+    assert!(proof.is_valid_for(&output));
+    assert!(output.public_pre_states.is_empty());
+    assert!(output.public_post_states.is_empty());
+    assert_eq!(output.new_commitments, expected_new_commitments);
+    assert_eq!(output.new_nullifiers, expected_new_nullifiers);
+    assert_eq!(output.encrypted_private_post_states.len(), 2);
+
+    let (_identifier, sender_post) = EncryptionScheme::decrypt(
+        &output.encrypted_private_post_states[0].ciphertext,
+        &shared_secret_1,
+        &expected_new_commitments[0],
+        0,
+    )
+    .unwrap();
+    assert_eq!(sender_post, expected_private_account_1);
+
+    let (_identifier, recipient_post) = EncryptionScheme::decrypt(
+        &output.encrypted_private_post_states[1].ciphertext,
+        &shared_secret_2,
+        &expected_new_commitments[1],
+        1,
+    )
+    .unwrap();
+    assert_eq!(recipient_post, expected_private_account_2);
+}
+
+#[test]
+fn circuit_fails_when_chained_validity_windows_have_empty_intersection() {
+    let account_keys = test_private_account_keys_1();
+    let pre = AccountWithMetadata::new(
+        Account::default(),
+        false,
+        AccountId::for_regular_private_account(&account_keys.npk(), 0),
+    );
+
+    let validity_window_chain_caller = crate::test_methods::validity_window_chain_caller();
+    let validity_window = crate::test_methods::validity_window();
+
+    let instruction = Program::serialize_instruction((
+        Some(1_u64),
+        Some(4_u64),
+        validity_window.id(),
+        Some(4_u64),
+        Some(7_u64),
+    ))
+    .unwrap();
+
+    let shared_secret =
+        SharedSecretKey::encapsulate_deterministic(&account_keys.vpk(), &[0_u8; 32], 0).0;
+
+    let program_with_deps = ProgramWithDependencies::new(
+        validity_window_chain_caller,
+        [(validity_window.id(), validity_window)].into(),
+    );
+
+    let result = execute_and_prove(
+        vec![pre],
+        instruction,
+        vec![InputAccountIdentity::PrivateUnauthorized {
+            epk: EphemeralPublicKey(Vec::new()),
+            view_tag: EncryptedAccountData::compute_view_tag(
+                &account_keys.npk(),
+                &account_keys.vpk(),
+            ),
+            npk: account_keys.npk(),
+            ssk: shared_secret,
+            identifier: 0,
+            commitment_root: DUMMY_COMMITMENT_HASH,
+        }],
+        &program_with_deps,
+    );
+
+    assert!(matches!(result, Err(LeeError::CircuitProvingError(_))));
+}
+
+/// A private PDA claimed with a non-default identifier produces a ciphertext that decrypts
+/// to `PrivateAccountKind::Pda` carrying the correct `(program_id, seed, identifier)`.
+#[test]
+fn private_pda_claim_with_custom_identifier_encrypts_correct_kind() {
+    let program = crate::test_methods::pda_claimer();
+    let keys = test_private_account_keys_1();
+    let npk = keys.npk();
+    let seed = PdaSeed::new([42; 32]);
+    let identifier: u128 = 99;
+    let shared_secret = SharedSecretKey::encapsulate_deterministic(&keys.vpk(), &[0_u8; 32], 0).0;
+
+    let account_id = AccountId::for_private_pda(&program.id(), &seed, &npk, identifier);
+    let pre_state = AccountWithMetadata::new(Account::default(), false, account_id);
+
+    let (output, _proof) = execute_and_prove(
+        vec![pre_state],
+        Program::serialize_instruction(seed).unwrap(),
+        vec![InputAccountIdentity::PrivatePdaInit {
+            epk: EphemeralPublicKey(Vec::new()),
+            view_tag: EncryptedAccountData::compute_view_tag(&npk, &keys.vpk()),
+            npk,
+            ssk: shared_secret,
+            identifier,
+            commitment_root: DUMMY_COMMITMENT_HASH,
+            seed: None,
+        }],
+        &program.clone().into(),
+    )
+    .unwrap();
+
+    assert_eq!(
+        decrypt_kind(&output, &shared_secret, 0),
+        PrivateAccountKind::Pda {
+            program_id: program.id(),
+            seed,
+            identifier
+        },
+    );
+}
+
+/// PDA init: initializes a new PDA under `simple_balance_transfer`'s ownership.
+/// The `simple_transfer_proxy` program chains to `simple_balance_transfer` with `pda_seeds`
+/// to establish authorization and the private PDA binding.
+#[test]
+fn private_pda_init() {
+    let program = crate::test_methods::simple_transfer_proxy();
+    let simple_transfer = crate::test_methods::simple_balance_transfer();
+    let keys = test_private_account_keys_1();
+    let npk = keys.npk();
+    let seed = PdaSeed::new([42; 32]);
+    let shared_secret_pda =
+        SharedSecretKey::encapsulate_deterministic(&keys.vpk(), &[0_u8; 32], 0).0;
+
+    // PDA (new, private PDA)
+    let pda_id = AccountId::for_private_pda(&program.id(), &seed, &npk, 0);
+    let pda_pre = AccountWithMetadata::new(Account::default(), false, pda_id);
+
+    let auth_id = simple_transfer.id();
+    let program_with_deps =
+        ProgramWithDependencies::new(program, [(auth_id, simple_transfer)].into());
+
+    // is_withdraw=false triggers init path (1 pre-state)
+    let instruction = Program::serialize_instruction((seed, auth_id, 0_u128, false)).unwrap();
+
+    let result = execute_and_prove(
+        vec![pda_pre],
+        instruction,
+        vec![InputAccountIdentity::PrivatePdaInit {
+            epk: EphemeralPublicKey(Vec::new()),
+            view_tag: EncryptedAccountData::compute_view_tag(&npk, &keys.vpk()),
+            npk,
+            ssk: shared_secret_pda,
+            identifier: 0,
+            commitment_root: DUMMY_COMMITMENT_HASH,
+            seed: None,
+        }],
+        &program_with_deps,
+    );
+
+    let (output, _proof) = result.expect("PDA init should succeed");
+    assert_eq!(output.new_commitments.len(), 1);
+}
+
+/// PDA withdraw: chains to `simple_balance_transfer` to move balance from PDA to recipient.
+/// Uses a default PDA (amount=0) because testing with a pre-funded PDA requires a
+/// two-tx sequence with membership proofs.
+#[test]
+fn private_pda_withdraw() {
+    let program = crate::test_methods::simple_transfer_proxy();
+    let simple_transfer = crate::test_methods::simple_balance_transfer();
+    let keys = test_private_account_keys_1();
+    let npk = keys.npk();
+    let seed = PdaSeed::new([42; 32]);
+    let shared_secret_pda =
+        SharedSecretKey::encapsulate_deterministic(&keys.vpk(), &[0_u8; 32], 0).0;
+
+    // PDA (new, private PDA)
+    let pda_id = AccountId::for_private_pda(&program.id(), &seed, &npk, 0);
+    let pda_pre = AccountWithMetadata::new(Account::default(), false, pda_id);
+
+    // Recipient (public)
+    let recipient_id = AccountId::new([88; 32]);
+    let recipient_pre = AccountWithMetadata::new(
+        Account {
+            program_owner: simple_transfer.id(),
+            balance: 10000,
+            ..Account::default()
+        },
+        true,
+        recipient_id,
+    );
+
+    let auth_id = simple_transfer.id();
+    let program_with_deps =
+        ProgramWithDependencies::new(program, [(auth_id, simple_transfer)].into());
+
+    // is_withdraw=true, amount=0 (PDA has no balance yet)
+    let instruction = Program::serialize_instruction((seed, auth_id, 0_u128, true)).unwrap();
+
+    let result = execute_and_prove(
+        vec![pda_pre, recipient_pre],
+        instruction,
+        vec![
+            InputAccountIdentity::PrivatePdaInit {
+                epk: EphemeralPublicKey(Vec::new()),
+                view_tag: EncryptedAccountData::compute_view_tag(&npk, &keys.vpk()),
+                npk,
+                ssk: shared_secret_pda,
+                identifier: 0,
+                commitment_root: DUMMY_COMMITMENT_HASH,
+                seed: None,
+            },
+            InputAccountIdentity::Public,
+        ],
+        &program_with_deps,
+    );
+
+    let (output, _proof) = result.expect("PDA withdraw should succeed");
+    assert_eq!(output.new_commitments.len(), 1);
+}
+
+/// Shared regular private account: receives funds via `authenticated_transfer` directly,
+/// no custom program needed. This demonstrates the non-PDA shared account flow where
+/// keys are derived from GMS via `derive_keys_for_shared_account`. The shared account
+/// uses the standard unauthorized private account path and works with auth-transfer's
+/// transfer path like any other private account.
+#[test]
+fn shared_account_receives_via_simple_transfer() {
+    let program = crate::test_methods::simple_balance_transfer();
+    let shared_keys = test_private_account_keys_1();
+    let shared_npk = shared_keys.npk();
+    let shared_identifier: u128 = 42;
+    let shared_secret =
+        SharedSecretKey::encapsulate_deterministic(&shared_keys.vpk(), &[0_u8; 32], 0).0;
+
+    // Sender: public account with balance, owned by auth-transfer
+    let sender_id = AccountId::new([99; 32]);
+    let sender = AccountWithMetadata::new(
+        Account {
+            program_owner: program.id(),
+            balance: 1000,
+            ..Account::default()
+        },
+        true,
+        sender_id,
+    );
+
+    // Recipient: shared private account (new, unauthorized)
+    let shared_account_id = AccountId::from((&shared_npk, shared_identifier));
+    let recipient = AccountWithMetadata::new(Account::default(), false, shared_account_id);
+
+    let balance_to_move: u128 = 100;
+    let instruction = Program::serialize_instruction(balance_to_move).unwrap();
+
+    let result = execute_and_prove(
+        vec![sender, recipient],
+        instruction,
+        vec![
+            InputAccountIdentity::Public,
+            InputAccountIdentity::PrivateUnauthorized {
+                epk: EphemeralPublicKey(Vec::new()),
+                view_tag: EncryptedAccountData::compute_view_tag(&shared_npk, &shared_keys.vpk()),
+                npk: shared_npk,
+                ssk: shared_secret,
+                identifier: shared_identifier,
+                commitment_root: DUMMY_COMMITMENT_HASH,
+            },
+        ],
+        &program.into(),
+    );
+
+    let (output, _proof) = result.expect("shared account receive should succeed");
+    // Sender is public (no commitment), recipient is private (1 commitment)
+    assert_eq!(output.new_commitments.len(), 1);
+}
+
+/// `PrivateAuthorizedInit` with a non-default identifier produces a ciphertext that decrypts
+/// to `PrivateAccountKind::Regular` carrying the correct identifier.
+#[test]
+fn private_authorized_init_encrypts_regular_kind_with_identifier() {
+    let program = crate::test_methods::claimer();
+    let keys = test_private_account_keys_1();
+    let identifier: u128 = 99;
+    let ssk = SharedSecretKey::encapsulate_deterministic(&keys.vpk(), &[0_u8; 32], 0).0;
+    let account_id = AccountId::for_regular_private_account(&keys.npk(), identifier);
+    let pre = AccountWithMetadata::new(Account::default(), true, account_id);
+
+    let (output, _) = execute_and_prove(
+        vec![pre],
+        Program::serialize_instruction(()).unwrap(),
+        vec![InputAccountIdentity::PrivateAuthorizedInit {
+            epk: EphemeralPublicKey(Vec::new()),
+            view_tag: EncryptedAccountData::compute_view_tag(&keys.npk(), &keys.vpk()),
+            ssk,
+            nsk: keys.nsk,
+            identifier,
+            commitment_root: DUMMY_COMMITMENT_HASH,
+        }],
+        &program.into(),
+    )
+    .unwrap();
+
+    assert_eq!(
+        decrypt_kind(&output, &ssk, 0),
+        PrivateAccountKind::Regular(identifier)
+    );
+}
+
+/// `PrivateUnauthorized` with a non-default identifier produces a ciphertext that decrypts
+/// to `PrivateAccountKind::Regular` carrying the correct identifier.
+#[test]
+fn private_unauthorized_init_encrypts_regular_kind_with_identifier() {
+    let program = crate::test_methods::claimer();
+    let keys = test_private_account_keys_1();
+    let identifier: u128 = 99;
+    let ssk = SharedSecretKey::encapsulate_deterministic(&keys.vpk(), &[0_u8; 32], 0).0;
+    let recipient_id = AccountId::for_regular_private_account(&keys.npk(), identifier);
+    let recipient = AccountWithMetadata::new(Account::default(), false, recipient_id);
+
+    let (output, _) = execute_and_prove(
+        vec![recipient],
+        Program::serialize_instruction(()).unwrap(),
+        vec![InputAccountIdentity::PrivateUnauthorized {
+            epk: EphemeralPublicKey(Vec::new()),
+            view_tag: EncryptedAccountData::compute_view_tag(&keys.npk(), &keys.vpk()),
+            npk: keys.npk(),
+            ssk,
+            identifier,
+            commitment_root: DUMMY_COMMITMENT_HASH,
+        }],
+        &program.into(),
+    )
+    .unwrap();
+
+    assert_eq!(
+        decrypt_kind(&output, &ssk, 0),
+        PrivateAccountKind::Regular(identifier)
+    );
+}
+
+/// `PrivateAuthorizedUpdate` with a non-default identifier produces a ciphertext that decrypts
+/// to `PrivateAccountKind::Regular` carrying the correct identifier.
+#[test]
+fn private_authorized_update_encrypts_regular_kind_with_identifier() {
+    let program = crate::test_methods::noop();
+    let keys = test_private_account_keys_1();
+    let identifier: u128 = 99;
+    let ssk = SharedSecretKey::encapsulate_deterministic(&keys.vpk(), &[0_u8; 32], 0).0;
+    let account_id = AccountId::for_regular_private_account(&keys.npk(), identifier);
+    let account = Account {
+        program_owner: program.id(),
+        balance: 1,
+        ..Account::default()
+    };
+    let commitment = Commitment::new(&account_id, &account);
+    let mut commitment_set = CommitmentSet::with_capacity(1);
+    commitment_set.extend(std::slice::from_ref(&commitment));
+
+    let sender = AccountWithMetadata::new(account, true, account_id);
+
+    let (output, _) = execute_and_prove(
+        vec![sender],
+        Program::serialize_instruction(()).unwrap(),
+        vec![InputAccountIdentity::PrivateAuthorizedUpdate {
+            epk: EphemeralPublicKey(Vec::new()),
+            view_tag: EncryptedAccountData::compute_view_tag(&keys.npk(), &keys.vpk()),
+            ssk,
+            nsk: keys.nsk,
+            membership_proof: commitment_set.get_proof_for(&commitment).unwrap(),
+            identifier,
+        }],
+        &program.into(),
+    )
+    .unwrap();
+
+    assert_eq!(
+        decrypt_kind(&output, &ssk, 0),
+        PrivateAccountKind::Regular(identifier)
+    );
+}
+
+/// `PrivatePdaUpdate` with a non-default identifier produces a ciphertext that decrypts
+/// to `PrivateAccountKind::Pda` carrying the correct `(program_id, seed, identifier)`.
+#[test]
+fn private_pda_update_encrypts_pda_kind_with_identifier() {
+    let program = crate::test_methods::pda_spend_proxy();
+    let simple_transfer = crate::test_methods::simple_balance_transfer();
+    let keys = test_private_account_keys_1();
+    let npk = keys.npk();
+    let seed = PdaSeed::new([42; 32]);
+    let identifier: u128 = 99;
+    let ssk = SharedSecretKey::encapsulate_deterministic(&keys.vpk(), &[0_u8; 32], 0).0;
+
+    let simple_transfer_id = simple_transfer.id();
+    let pda_id = AccountId::for_private_pda(&program.id(), &seed, &npk, identifier);
+    let pda_account = Account {
+        program_owner: simple_transfer_id,
+        balance: 1,
+        ..Account::default()
+    };
+    let pda_commitment = Commitment::new(&pda_id, &pda_account);
+    let mut commitment_set = CommitmentSet::with_capacity(1);
+    commitment_set.extend(std::slice::from_ref(&pda_commitment));
+
+    let pda_pre = AccountWithMetadata::new(pda_account, true, pda_id);
+    let recipient_pre = AccountWithMetadata::new(Account::default(), true, AccountId::new([0; 32]));
+
+    let program_with_deps = ProgramWithDependencies::new(
+        program.clone(),
+        [(simple_transfer_id, simple_transfer)].into(),
+    );
+
+    let (output, _) = execute_and_prove(
+        vec![pda_pre, recipient_pre],
+        Program::serialize_instruction((seed, 1_u128, simple_transfer_id, false)).unwrap(),
+        vec![
+            InputAccountIdentity::PrivatePdaUpdate {
+                epk: EphemeralPublicKey(Vec::new()),
+                view_tag: EncryptedAccountData::compute_view_tag(&npk, &keys.vpk()),
+                ssk,
+                nsk: keys.nsk,
+                membership_proof: commitment_set.get_proof_for(&pda_commitment).unwrap(),
+                identifier,
+                seed: None,
+            },
+            InputAccountIdentity::Public,
+        ],
+        &program_with_deps,
+    )
+    .unwrap();
+
+    assert_eq!(
+        decrypt_kind(&output, &ssk, 0),
+        PrivateAccountKind::Pda {
+            program_id: program.id(),
+            seed,
+            identifier
+        },
+    );
+}
+
+#[test]
+fn private_pda_init_identifier_mismatch_fails() {
+    let program = crate::test_methods::pda_claimer();
+    let keys = test_private_account_keys_1();
+    let npk = keys.npk();
+    let seed = PdaSeed::new([42; 32]);
+    let shared_secret = SharedSecretKey::encapsulate_deterministic(&keys.vpk(), &[0_u8; 32], 0).0;
+
+    let account_id = AccountId::for_private_pda(&program.id(), &seed, &npk, 5);
+    let pre_state = AccountWithMetadata::new(Account::default(), false, account_id);
+
+    let result = execute_and_prove(
+        vec![pre_state],
+        Program::serialize_instruction(seed).unwrap(),
+        vec![InputAccountIdentity::PrivatePdaInit {
+            epk: EphemeralPublicKey(Vec::new()),
+            view_tag: EncryptedAccountData::compute_view_tag(&npk, &keys.vpk()),
+            npk,
+            ssk: shared_secret,
+            identifier: 99,
+            commitment_root: DUMMY_COMMITMENT_HASH,
+            seed: None,
+        }],
+        &program.into(),
+    );
+
+    assert!(matches!(result, Err(LeeError::CircuitProvingError(_))));
+}
+
+#[test]
+fn private_pda_update_identifier_mismatch_fails() {
+    let program = crate::test_methods::pda_spend_proxy();
+    let simple_transfer = crate::test_methods::simple_balance_transfer();
+    let keys = test_private_account_keys_1();
+    let npk = keys.npk();
+    let seed = PdaSeed::new([42; 32]);
+    let ssk = SharedSecretKey::encapsulate_deterministic(&keys.vpk(), &[0_u8; 32], 0).0;
+
+    let simple_transfer_id = simple_transfer.id();
+    let pda_id = AccountId::for_private_pda(&program.id(), &seed, &npk, 5);
+    let pda_account = Account {
+        program_owner: simple_transfer_id,
+        balance: 1,
+        ..Account::default()
+    };
+    let pda_commitment = Commitment::new(&pda_id, &pda_account);
+    let mut commitment_set = CommitmentSet::with_capacity(1);
+    commitment_set.extend(std::slice::from_ref(&pda_commitment));
+
+    let pda_pre = AccountWithMetadata::new(pda_account, true, pda_id);
+    let recipient_pre = AccountWithMetadata::new(Account::default(), true, AccountId::new([0; 32]));
+
+    let program_with_deps =
+        ProgramWithDependencies::new(program, [(simple_transfer_id, simple_transfer)].into());
+
+    let result = execute_and_prove(
+        vec![pda_pre, recipient_pre],
+        Program::serialize_instruction((seed, 1_u128, simple_transfer_id, false)).unwrap(),
+        vec![
+            InputAccountIdentity::PrivatePdaUpdate {
+                epk: EphemeralPublicKey(Vec::new()),
+                view_tag: EncryptedAccountData::compute_view_tag(&npk, &keys.vpk()),
+                ssk,
+                nsk: keys.nsk,
+                membership_proof: commitment_set.get_proof_for(&pda_commitment).unwrap(),
+                identifier: 99,
+                seed: None,
+            },
+            InputAccountIdentity::Public,
+        ],
+        &program_with_deps,
+    );
+
+    assert!(matches!(result, Err(LeeError::CircuitProvingError(_))));
+}

lee/state_machine/src/program/mod.rs

@@ -111,42 +111,4 @@ impl Program {
 }
 
 #[cfg(test)]
-mod tests {
-    use lee_core::account::{Account, AccountId, AccountWithMetadata};
-
-    use crate::program::Program;
-
-    #[test]
-    fn program_execution() {
-        let program = crate::test_methods::simple_balance_transfer();
-        let balance_to_move: u128 = 11_223_344_556_677;
-        let instruction_data = Program::serialize_instruction(balance_to_move).unwrap();
-        let sender = AccountWithMetadata::new(
-            Account {
-                balance: 77_665_544_332_211,
-                ..Account::default()
-            },
-            true,
-            AccountId::new([0; 32]),
-        );
-        let recipient =
-            AccountWithMetadata::new(Account::default(), false, AccountId::new([1; 32]));
-
-        let expected_sender_post = Account {
-            balance: 77_665_544_332_211 - balance_to_move,
-            ..Account::default()
-        };
-        let expected_recipient_post = Account {
-            balance: balance_to_move,
-            ..Account::default()
-        };
-        let program_output = program
-            .execute(None, &[sender, recipient], &instruction_data)
-            .unwrap();
-
-        let [sender_post, recipient_post] = program_output.post_states.try_into().unwrap();
-
-        assert_eq!(sender_post.account(), &expected_sender_post);
-        assert_eq!(recipient_post.account(), &expected_recipient_post);
-    }
-}
+mod tests;

lee/state_machine/src/program/tests.rs

@@ -0,0 +1,36 @@
+use lee_core::account::{Account, AccountId, AccountWithMetadata};
+
+use crate::program::Program;
+
+#[test]
+fn program_execution() {
+    let program = crate::test_methods::simple_balance_transfer();
+    let balance_to_move: u128 = 11_223_344_556_677;
+    let instruction_data = Program::serialize_instruction(balance_to_move).unwrap();
+    let sender = AccountWithMetadata::new(
+        Account {
+            balance: 77_665_544_332_211,
+            ..Account::default()
+        },
+        true,
+        AccountId::new([0; 32]),
+    );
+    let recipient = AccountWithMetadata::new(Account::default(), false, AccountId::new([1; 32]));
+
+    let expected_sender_post = Account {
+        balance: 77_665_544_332_211 - balance_to_move,
+        ..Account::default()
+    };
+    let expected_recipient_post = Account {
+        balance: balance_to_move,
+        ..Account::default()
+    };
+    let program_output = program
+        .execute(None, &[sender, recipient], &instruction_data)
+        .unwrap();
+
+    let [sender_post, recipient_post] = program_output.post_states.try_into().unwrap();
+
+    assert_eq!(sender_post.account(), &expected_sender_post);
+    assert_eq!(recipient_post.account(), &expected_recipient_post);
+}

lee/state_machine/src/state/mod.rs

@@ -0,0 +1,328 @@
+use std::collections::{BTreeSet, HashMap, HashSet};
+
+use borsh::{BorshDeserialize, BorshSerialize};
+use lee_core::{
+    BlockId, Commitment, CommitmentSetDigest, DUMMY_COMMITMENT, MembershipProof, Nullifier,
+    Timestamp,
+    account::{Account, AccountId},
+    program::ProgramId,
+};
+
+use crate::{
+    error::LeeError,
+    merkle_tree::MerkleTree,
+    privacy_preserving_transaction::PrivacyPreservingTransaction,
+    program::Program,
+    program_deployment_transaction::ProgramDeploymentTransaction,
+    public_transaction::PublicTransaction,
+    validated_state_diff::{StateDiff, ValidatedStateDiff},
+};
+
+pub const MAX_NUMBER_CHAINED_CALLS: usize = 10;
+
+#[derive(Clone, BorshSerialize, BorshDeserialize)]
+#[cfg_attr(test, derive(Debug, PartialEq, Eq))]
+pub struct CommitmentSet {
+    merkle_tree: MerkleTree,
+    commitments: HashMap<Commitment, usize>,
+    root_history: HashSet<CommitmentSetDigest>,
+}
+
+impl CommitmentSet {
+    pub(crate) fn digest(&self) -> CommitmentSetDigest {
+        self.merkle_tree.root()
+    }
+
+    /// Queries the `CommitmentSet` for a membership proof of commitment.
+    pub fn get_proof_for(&self, commitment: &Commitment) -> Option<MembershipProof> {
+        let index = *self.commitments.get(commitment)?;
+
+        self.merkle_tree
+            .get_authentication_path_for(index)
+            .map(|path| (index, path))
+    }
+
+    /// Inserts a list of commitments to the `CommitmentSet`.
+    pub(crate) fn extend(&mut self, commitments: &[Commitment]) {
+        for commitment in commitments.iter().cloned() {
+            let index = self.merkle_tree.insert(commitment.to_byte_array());
+            self.commitments.insert(commitment, index);
+        }
+        self.root_history.insert(self.digest());
+    }
+
+    fn contains(&self, commitment: &Commitment) -> bool {
+        self.commitments.contains_key(commitment)
+    }
+
+    /// Initializes an empty `CommitmentSet` with a given capacity.
+    /// If the capacity is not a `power_of_two`, then capacity is taken
+    /// to be the next `power_of_two`.
+    pub(crate) fn with_capacity(capacity: usize) -> Self {
+        Self {
+            merkle_tree: MerkleTree::with_capacity(capacity),
+            commitments: HashMap::new(),
+            root_history: HashSet::new(),
+        }
+    }
+}
+
+#[cfg_attr(test, derive(Debug, PartialEq, Eq))]
+#[derive(Clone)]
+struct NullifierSet(BTreeSet<Nullifier>);
+
+impl NullifierSet {
+    const fn new() -> Self {
+        Self(BTreeSet::new())
+    }
+
+    fn extend(&mut self, new_nullifiers: &[Nullifier]) {
+        self.0.extend(new_nullifiers);
+    }
+
+    fn contains(&self, nullifier: &Nullifier) -> bool {
+        self.0.contains(nullifier)
+    }
+}
+
+impl BorshSerialize for NullifierSet {
+    fn serialize<W: std::io::Write>(&self, writer: &mut W) -> std::io::Result<()> {
+        self.0.iter().collect::<Vec<_>>().serialize(writer)
+    }
+}
+
+impl BorshDeserialize for NullifierSet {
+    fn deserialize_reader<R: std::io::Read>(reader: &mut R) -> std::io::Result<Self> {
+        let vec = Vec::<Nullifier>::deserialize_reader(reader)?;
+
+        let mut set = BTreeSet::new();
+        for n in vec {
+            if !set.insert(n) {
+                return Err(std::io::Error::new(
+                    std::io::ErrorKind::InvalidData,
+                    "duplicate nullifier in NullifierSet",
+                ));
+            }
+        }
+
+        Ok(Self(set))
+    }
+}
+
+#[derive(Clone, BorshSerialize, BorshDeserialize)]
+#[cfg_attr(test, derive(Debug, PartialEq, Eq))]
+pub struct V03State {
+    public_state: HashMap<AccountId, Account>,
+    private_state: (CommitmentSet, NullifierSet),
+    programs: HashMap<ProgramId, Program>,
+}
+
+impl Default for V03State {
+    fn default() -> Self {
+        let mut commitment_set = CommitmentSet::with_capacity(32);
+        commitment_set.extend(&[DUMMY_COMMITMENT]);
+        let nullifier_set = NullifierSet::new();
+        let private_state = (commitment_set, nullifier_set);
+
+        Self {
+            public_state: HashMap::default(),
+            private_state,
+            programs: HashMap::default(),
+        }
+    }
+}
+
+impl V03State {
+    #[must_use]
+    pub fn new() -> Self {
+        Self::default()
+    }
+
+    /// Initializes state with given public account balances leaving other account fields at their
+    /// default values.
+    #[must_use]
+    pub fn with_public_account_balances(
+        mut self,
+        balances: impl IntoIterator<Item = (AccountId, u128)>,
+    ) -> Self {
+        let public_accounts = balances.into_iter().map(|(account_id, balance)| {
+            (
+                account_id,
+                Account {
+                    balance,
+                    ..Account::default()
+                },
+            )
+        });
+        self.public_state.extend(public_accounts);
+        self
+    }
+
+    /// Initializes state with given public accounts.
+    #[must_use]
+    pub fn with_public_accounts(
+        mut self,
+        public_accounts: impl IntoIterator<Item = (AccountId, Account)>,
+    ) -> Self {
+        self.public_state.extend(public_accounts);
+        self
+    }
+
+    /// Initializes state with given private accounts.
+    #[must_use]
+    pub fn with_private_accounts(
+        mut self,
+        private_accounts: impl IntoIterator<Item = (Commitment, Nullifier)>,
+    ) -> Self {
+        let (commitments, nullifiers): (Vec<Commitment>, Vec<Nullifier>) =
+            private_accounts.into_iter().unzip();
+        self.private_state.0.extend(&commitments);
+        self.private_state.1.extend(&nullifiers);
+        self
+    }
+
+    /// Initializes state with given builtin programs.
+    #[must_use]
+    pub fn with_programs(mut self, programs: impl IntoIterator<Item = Program>) -> Self {
+        for program in programs {
+            self.insert_program(program);
+        }
+        self
+    }
+
+    pub(crate) fn insert_program(&mut self, program: Program) {
+        self.programs.insert(program.id(), program);
+    }
+
+    /// Seeds a single genesis account that is not produced by any transaction
+    /// (e.g. the cross-zone inbox config or a bridge-lock holding). Lets the
+    /// sequencer and indexer seed identical zone-specific state after building
+    /// the shared initial state.
+    pub fn insert_genesis_account(&mut self, account_id: AccountId, account: Account) {
+        self.public_state.insert(account_id, account);
+    }
+
+    pub fn apply_state_diff(&mut self, diff: ValidatedStateDiff) {
+        let StateDiff {
+            signer_account_ids,
+            public_diff,
+            new_commitments,
+            new_nullifiers,
+            program,
+        } = diff.into_state_diff();
+        #[expect(
+            clippy::iter_over_hash_type,
+            reason = "Iteration order doesn't matter here"
+        )]
+        for (account_id, account) in public_diff {
+            *self.get_account_by_id_mut(account_id) = account;
+        }
+        for account_id in signer_account_ids {
+            self.get_account_by_id_mut(account_id)
+                .nonce
+                .public_account_nonce_increment();
+        }
+        self.private_state.0.extend(&new_commitments);
+        self.private_state.1.extend(&new_nullifiers);
+        if let Some(program) = program {
+            self.insert_program(program);
+        }
+    }
+
+    pub fn transition_from_public_transaction(
+        &mut self,
+        tx: &PublicTransaction,
+        block_id: BlockId,
+        timestamp: Timestamp,
+    ) -> Result<(), LeeError> {
+        let diff = ValidatedStateDiff::from_public_transaction(tx, self, block_id, timestamp)?;
+        self.apply_state_diff(diff);
+        Ok(())
+    }
+
+    pub fn transition_from_privacy_preserving_transaction(
+        &mut self,
+        tx: &PrivacyPreservingTransaction,
+        block_id: BlockId,
+        timestamp: Timestamp,
+    ) -> Result<(), LeeError> {
+        let diff =
+            ValidatedStateDiff::from_privacy_preserving_transaction(tx, self, block_id, timestamp)?;
+        self.apply_state_diff(diff);
+        Ok(())
+    }
+
+    pub fn transition_from_program_deployment_transaction(
+        &mut self,
+        tx: &ProgramDeploymentTransaction,
+    ) -> Result<(), LeeError> {
+        let diff = ValidatedStateDiff::from_program_deployment_transaction(tx, self)?;
+        self.apply_state_diff(diff);
+        Ok(())
+    }
+
+    fn get_account_by_id_mut(&mut self, account_id: AccountId) -> &mut Account {
+        self.public_state.entry(account_id).or_default()
+    }
+
+    #[must_use]
+    pub fn get_account_by_id(&self, account_id: AccountId) -> Account {
+        self.public_state
+            .get(&account_id)
+            .cloned()
+            .unwrap_or_else(Account::default)
+    }
+
+    #[must_use]
+    pub fn get_proof_for_commitment(&self, commitment: &Commitment) -> Option<MembershipProof> {
+        self.private_state.0.get_proof_for(commitment)
+    }
+
+    pub(crate) const fn programs(&self) -> &HashMap<ProgramId, Program> {
+        &self.programs
+    }
+
+    #[must_use]
+    pub fn commitment_set_digest(&self) -> CommitmentSetDigest {
+        self.private_state.0.digest()
+    }
+
+    pub(crate) fn check_commitments_are_new(
+        &self,
+        new_commitments: &[Commitment],
+    ) -> Result<(), LeeError> {
+        for commitment in new_commitments {
+            if self.private_state.0.contains(commitment) {
+                return Err(LeeError::InvalidInput("Commitment already seen".to_owned()));
+            }
+        }
+        Ok(())
+    }
+
+    pub(crate) fn check_nullifiers_are_valid(
+        &self,
+        new_nullifiers: &[(Nullifier, CommitmentSetDigest)],
+    ) -> Result<(), LeeError> {
+        for (nullifier, digest) in new_nullifiers {
+            if self.private_state.1.contains(nullifier) {
+                return Err(LeeError::InvalidInput("Nullifier already seen".to_owned()));
+            }
+            if !self.private_state.0.root_history.contains(digest) {
+                return Err(LeeError::InvalidInput(
+                    "Unrecognized commitment set digest".to_owned(),
+                ));
+            }
+        }
+        Ok(())
+    }
+}
+
+#[cfg(any(test, feature = "test-utils"))]
+impl V03State {
+    pub fn force_insert_account(&mut self, account_id: AccountId, account: Account) {
+        self.public_state.insert(account_id, account);
+    }
+}
+
+#[cfg(test)]
+pub mod tests;

lee/state_machine/src/state/tests/authenticated_transfer.rs

@@ -0,0 +1,149 @@
+use super::*;
+
+#[test]
+fn transition_from_authenticated_transfer_program_invocation_default_account_destination() {
+    let key = PrivateKey::try_new([1; 32]).unwrap();
+    let account_id = AccountId::from(&PublicKey::new_from_private_key(&key));
+    let initial_data = [(
+        account_id,
+        Account {
+            program_owner: crate::test_methods::simple_balance_transfer().id(),
+            balance: 100,
+            ..Account::default()
+        },
+    )];
+    let mut state = V03State::new()
+        .with_public_accounts(initial_data)
+        .with_test_programs();
+    let from = account_id;
+    let to_key = PrivateKey::try_new([2; 32]).unwrap();
+    let to = AccountId::from(&PublicKey::new_from_private_key(&to_key));
+    assert_eq!(state.get_account_by_id(to), Account::default());
+    let balance_to_move = 5;
+
+    let tx = transfer_transaction(from, &key, 0, to, &to_key, 0, balance_to_move);
+    state.transition_from_public_transaction(&tx, 1, 0).unwrap();
+
+    assert_eq!(state.get_account_by_id(from).balance, 95);
+    assert_eq!(state.get_account_by_id(to).balance, 5);
+    assert_eq!(state.get_account_by_id(from).nonce, Nonce(1));
+    assert_eq!(state.get_account_by_id(to).nonce, Nonce(1));
+}
+
+#[test]
+fn transition_from_authenticated_transfer_program_invocation_insuficient_balance() {
+    let key = PrivateKey::try_new([1; 32]).unwrap();
+    let account_id = AccountId::from(&PublicKey::new_from_private_key(&key));
+    let mut state = V03State::new()
+        .with_public_account_balances([(account_id, 100)])
+        .with_test_programs();
+    let from = account_id;
+    let from_key = key;
+    let to_key = PrivateKey::try_new([2; 32]).unwrap();
+    let to = AccountId::from(&PublicKey::new_from_private_key(&to_key));
+    let balance_to_move = 101;
+    assert!(state.get_account_by_id(from).balance < balance_to_move);
+
+    let tx = transfer_transaction(from, &from_key, 0, to, &to_key, 0, balance_to_move);
+    let result = state.transition_from_public_transaction(&tx, 1, 0);
+
+    assert!(matches!(result, Err(LeeError::ProgramExecutionFailed(_))));
+    assert_eq!(state.get_account_by_id(from).balance, 100);
+    assert_eq!(state.get_account_by_id(to).balance, 0);
+    assert_eq!(state.get_account_by_id(from).nonce, Nonce(0));
+    assert_eq!(state.get_account_by_id(to).nonce, Nonce(0));
+}
+
+#[test]
+fn transition_from_authenticated_transfer_program_invocation_non_default_account_destination() {
+    let key1 = PrivateKey::try_new([1; 32]).unwrap();
+    let key2 = PrivateKey::try_new([2; 32]).unwrap();
+    let account_id1 = AccountId::from(&PublicKey::new_from_private_key(&key1));
+    let account_id2 = AccountId::from(&PublicKey::new_from_private_key(&key2));
+    let initial_data = [
+        (
+            account_id1,
+            Account {
+                program_owner: crate::test_methods::simple_balance_transfer().id(),
+                balance: 100,
+                ..Account::default()
+            },
+        ),
+        (
+            account_id2,
+            Account {
+                program_owner: crate::test_methods::simple_balance_transfer().id(),
+                balance: 200,
+                ..Account::default()
+            },
+        ),
+    ];
+    let mut state = V03State::new()
+        .with_public_accounts(initial_data)
+        .with_test_programs();
+    let from = account_id2;
+    let from_key = key2;
+    let to = account_id1;
+    let to_key = key1;
+    assert_ne!(state.get_account_by_id(to), Account::default());
+    let balance_to_move = 8;
+
+    let tx = transfer_transaction(from, &from_key, 0, to, &to_key, 0, balance_to_move);
+    state.transition_from_public_transaction(&tx, 1, 0).unwrap();
+
+    assert_eq!(state.get_account_by_id(from).balance, 192);
+    assert_eq!(state.get_account_by_id(to).balance, 108);
+    assert_eq!(state.get_account_by_id(from).nonce, Nonce(1));
+    assert_eq!(state.get_account_by_id(to).nonce, Nonce(1));
+}
+
+#[test]
+fn transition_from_sequence_of_authenticated_transfer_program_invocations() {
+    let key1 = PrivateKey::try_new([8; 32]).unwrap();
+    let account_id1 = AccountId::from(&PublicKey::new_from_private_key(&key1));
+    let key2 = PrivateKey::try_new([2; 32]).unwrap();
+    let account_id2 = AccountId::from(&PublicKey::new_from_private_key(&key2));
+    let initial_data = [(
+        account_id1,
+        Account {
+            program_owner: crate::test_methods::simple_balance_transfer().id(),
+            balance: 100,
+            ..Account::default()
+        },
+    )];
+    let mut state = V03State::new()
+        .with_public_accounts(initial_data)
+        .with_test_programs();
+    let key3 = PrivateKey::try_new([3; 32]).unwrap();
+    let account_id3 = AccountId::from(&PublicKey::new_from_private_key(&key3));
+    let balance_to_move = 5;
+
+    let tx = transfer_transaction(
+        account_id1,
+        &key1,
+        0,
+        account_id2,
+        &key2,
+        0,
+        balance_to_move,
+    );
+    state.transition_from_public_transaction(&tx, 1, 0).unwrap();
+    let balance_to_move = 3;
+    let tx = transfer_transaction(
+        account_id2,
+        &key2,
+        1,
+        account_id3,
+        &key3,
+        0,
+        balance_to_move,
+    );
+    state.transition_from_public_transaction(&tx, 1, 0).unwrap();
+
+    assert_eq!(state.get_account_by_id(account_id1).balance, 95);
+    assert_eq!(state.get_account_by_id(account_id2).balance, 2);
+    assert_eq!(state.get_account_by_id(account_id3).balance, 3);
+    assert_eq!(state.get_account_by_id(account_id1).nonce, Nonce(1));
+    assert_eq!(state.get_account_by_id(account_id2).nonce, Nonce(2));
+    assert_eq!(state.get_account_by_id(account_id3).nonce, Nonce(1));
+}

lee/state_machine/src/state/tests/changer_claimer.rs

@@ -0,0 +1,118 @@
+use super::*;
+
+#[test]
+fn public_changer_claimer_no_data_change_no_claim_succeeds() {
+    let initial_data = [];
+    let mut state = V03State::new()
+        .with_public_accounts(public_state_from_balances(&initial_data))
+        .with_test_programs();
+    let account_id = AccountId::new([1; 32]);
+    let program_id = crate::test_methods::changer_claimer().id();
+    // Don't change data (None) and don't claim (false)
+    let instruction: (Option<Vec<u8>>, bool) = (None, false);
+
+    let message =
+        public_transaction::Message::try_new(program_id, vec![account_id], vec![], instruction)
+            .unwrap();
+    let witness_set = public_transaction::WitnessSet::for_message(&message, &[]);
+    let tx = PublicTransaction::new(message, witness_set);
+
+    let result = state.transition_from_public_transaction(&tx, 1, 0);
+
+    // Should succeed - no changes made, no claim needed
+    assert!(result.is_ok());
+    // Account should remain default/unclaimed
+    assert_eq!(state.get_account_by_id(account_id), Account::default());
+}
+
+#[test]
+fn public_changer_claimer_data_change_no_claim_fails() {
+    let initial_data = [];
+    let mut state = V03State::new()
+        .with_public_accounts(public_state_from_balances(&initial_data))
+        .with_test_programs();
+    let account_id = AccountId::new([1; 32]);
+    let program_id = crate::test_methods::changer_claimer().id();
+    // Change data but don't claim (false) - should fail
+    let new_data = vec![1, 2, 3, 4, 5];
+    let instruction: (Option<Vec<u8>>, bool) = (Some(new_data), false);
+
+    let message =
+        public_transaction::Message::try_new(program_id, vec![account_id], vec![], instruction)
+            .unwrap();
+    let witness_set = public_transaction::WitnessSet::for_message(&message, &[]);
+    let tx = PublicTransaction::new(message, witness_set);
+
+    let result = state.transition_from_public_transaction(&tx, 1, 0);
+
+    // Should fail - cannot modify data without claiming the account
+    assert!(matches!(
+        result,
+        Err(LeeError::InvalidProgramBehavior(
+            InvalidProgramBehaviorError::DefaultAccountModifiedWithoutClaim {
+                account_id: err_account_id
+            }
+        )) if err_account_id == account_id
+    ));
+}
+
+#[test]
+fn private_changer_claimer_no_data_change_no_claim_succeeds() {
+    let program = crate::test_methods::changer_claimer();
+    let sender_keys = test_private_account_keys_1();
+    let private_account =
+        AccountWithMetadata::new(Account::default(), true, (&sender_keys.npk(), 0));
+    // Don't change data (None) and don't claim (false)
+    let instruction: (Option<Vec<u8>>, bool) = (None, false);
+
+    let result = execute_and_prove(
+        vec![private_account],
+        Program::serialize_instruction(instruction).unwrap(),
+        vec![InputAccountIdentity::PrivateAuthorizedUpdate {
+            epk: EphemeralPublicKey(Vec::new()),
+            view_tag: EncryptedAccountData::compute_view_tag(
+                &sender_keys.npk(),
+                &sender_keys.vpk(),
+            ),
+            ssk: SharedSecretKey::encapsulate_deterministic(&sender_keys.vpk(), &[0_u8; 32], 0).0,
+            nsk: sender_keys.nsk,
+            membership_proof: (0, vec![]),
+            identifier: 0,
+        }],
+        &program.into(),
+    );
+
+    // Should succeed - no changes made, no claim needed
+    assert!(result.is_ok());
+}
+
+#[test]
+fn private_changer_claimer_data_change_no_claim_fails() {
+    let program = crate::test_methods::changer_claimer();
+    let sender_keys = test_private_account_keys_1();
+    let private_account =
+        AccountWithMetadata::new(Account::default(), true, (&sender_keys.npk(), 0));
+    // Change data but don't claim (false) - should fail
+    let new_data = vec![1, 2, 3, 4, 5];
+    let instruction: (Option<Vec<u8>>, bool) = (Some(new_data), false);
+
+    let result = execute_and_prove(
+        vec![private_account],
+        Program::serialize_instruction(instruction).unwrap(),
+        vec![InputAccountIdentity::PrivateAuthorizedUpdate {
+            epk: EphemeralPublicKey(Vec::new()),
+            view_tag: EncryptedAccountData::compute_view_tag(
+                &sender_keys.npk(),
+                &sender_keys.vpk(),
+            ),
+            ssk: SharedSecretKey::encapsulate_deterministic(&sender_keys.vpk(), &[0_u8; 32], 0).0,
+            nsk: sender_keys.nsk,
+            membership_proof: (0, vec![]),
+            identifier: 0,
+        }],
+        &program.into(),
+    );
+
+    // Should fail - cannot modify data without claiming the account
+    assert!(matches!(result, Err(LeeError::CircuitProvingError(_))));
+}

lee/state_machine/src/state/tests/claiming.rs

@@ -0,0 +1,619 @@
+use super::*;
+
+#[test]
+fn claiming_mechanism() {
+    let program = crate::test_methods::simple_balance_transfer();
+    let from_key = PrivateKey::try_new([1; 32]).unwrap();
+    let from = AccountId::from(&PublicKey::new_from_private_key(&from_key));
+    let initial_balance = 100;
+    let initial_data = [(from, initial_balance)];
+    let mut state = V03State::new()
+        .with_public_accounts(public_state_from_balances(&initial_data))
+        .with_test_programs();
+    let to_key = PrivateKey::try_new([2; 32]).unwrap();
+    let to = AccountId::from(&PublicKey::new_from_private_key(&to_key));
+    let amount: u128 = 37;
+
+    // Check the recipient is an uninitialized account
+    assert_eq!(state.get_account_by_id(to), Account::default());
+
+    let expected_recipient_post = Account {
+        program_owner: program.id(),
+        balance: amount,
+        nonce: Nonce(1),
+        ..Account::default()
+    };
+
+    let message = public_transaction::Message::try_new(
+        program.id(),
+        vec![from, to],
+        vec![Nonce(0), Nonce(0)],
+        amount,
+    )
+    .unwrap();
+    let witness_set = public_transaction::WitnessSet::for_message(&message, &[&from_key, &to_key]);
+    let tx = PublicTransaction::new(message, witness_set);
+
+    state.transition_from_public_transaction(&tx, 1, 0).unwrap();
+
+    let recipient_post = state.get_account_by_id(to);
+
+    assert_eq!(recipient_post, expected_recipient_post);
+}
+
+#[test]
+fn unauthorized_public_account_claiming_fails() {
+    let program = crate::test_methods::simple_balance_transfer();
+    let account_key = PrivateKey::try_new([9; 32]).unwrap();
+    let account_id = AccountId::from(&PublicKey::new_from_private_key(&account_key));
+    let mut state = V03State::new().with_test_programs();
+
+    assert_eq!(state.get_account_by_id(account_id), Account::default());
+
+    let message =
+        public_transaction::Message::try_new(program.id(), vec![account_id], vec![], 0_u128)
+            .unwrap();
+    let witness_set = public_transaction::WitnessSet::for_message(&message, &[]);
+    let tx = PublicTransaction::new(message, witness_set);
+
+    let result = state.transition_from_public_transaction(&tx, 2, 0);
+
+    assert!(matches!(result, Err(LeeError::InvalidProgramBehavior(_))));
+    assert_eq!(state.get_account_by_id(account_id), Account::default());
+}
+
+#[test]
+fn authorized_public_account_claiming_succeeds() {
+    let program = crate::test_methods::simple_balance_transfer();
+    let account_key = PrivateKey::try_new([10; 32]).unwrap();
+    let account_id = AccountId::from(&PublicKey::new_from_private_key(&account_key));
+    let mut state = V03State::new().with_test_programs();
+
+    assert_eq!(state.get_account_by_id(account_id), Account::default());
+
+    let message = public_transaction::Message::try_new(
+        program.id(),
+        vec![account_id],
+        vec![Nonce(0)],
+        0_u128,
+    )
+    .unwrap();
+    let witness_set = public_transaction::WitnessSet::for_message(&message, &[&account_key]);
+    let tx = PublicTransaction::new(message, witness_set);
+
+    state.transition_from_public_transaction(&tx, 1, 0).unwrap();
+
+    assert_eq!(
+        state.get_account_by_id(account_id),
+        Account {
+            program_owner: program.id(),
+            nonce: Nonce(1),
+            ..Account::default()
+        }
+    );
+}
+
+#[test]
+fn public_chained_call() {
+    let program = crate::test_methods::chain_caller();
+    let key = PrivateKey::try_new([1; 32]).unwrap();
+    let from = AccountId::from(&PublicKey::new_from_private_key(&key));
+    let to = AccountId::new([2; 32]);
+    let initial_balance = 1000;
+    let initial_data = [(from, initial_balance), (to, 0)];
+    let mut state = V03State::new()
+        .with_public_accounts(public_state_from_balances(&initial_data))
+        .with_test_programs();
+    let from_key = key;
+    let amount: u128 = 37;
+    let instruction: (u128, ProgramId, u32, Option<PdaSeed>) = (
+        amount,
+        crate::test_methods::simple_balance_transfer().id(),
+        2,
+        None,
+    );
+
+    let expected_to_post = Account {
+        program_owner: crate::test_methods::simple_balance_transfer().id(),
+        balance: amount * 2, // The `chain_caller` chains the program twice
+        ..Account::default()
+    };
+
+    let message = public_transaction::Message::try_new(
+        program.id(),
+        vec![to, from], // The chain_caller program permutes the account order in the chain
+        // call
+        vec![Nonce(0)],
+        instruction,
+    )
+    .unwrap();
+    let witness_set = public_transaction::WitnessSet::for_message(&message, &[&from_key]);
+    let tx = PublicTransaction::new(message, witness_set);
+
+    state.transition_from_public_transaction(&tx, 1, 0).unwrap();
+
+    let from_post = state.get_account_by_id(from);
+    let to_post = state.get_account_by_id(to);
+    // The `chain_caller` program calls the program twice
+    assert_eq!(from_post.balance, initial_balance - 2 * amount);
+    assert_eq!(to_post, expected_to_post);
+}
+
+#[test]
+fn execution_fails_if_chained_calls_exceeds_depth() {
+    let program = crate::test_methods::chain_caller();
+    let key = PrivateKey::try_new([1; 32]).unwrap();
+    let from = AccountId::from(&PublicKey::new_from_private_key(&key));
+    let to = AccountId::new([2; 32]);
+    let initial_balance = 100;
+    let initial_data = [(from, initial_balance), (to, 0)];
+    let mut state = V03State::new()
+        .with_public_accounts(public_state_from_balances(&initial_data))
+        .with_test_programs();
+    let from_key = key;
+    let amount: u128 = 0;
+    let instruction: (u128, ProgramId, u32, Option<PdaSeed>) = (
+        amount,
+        crate::test_methods::simple_balance_transfer().id(),
+        u32::try_from(MAX_NUMBER_CHAINED_CALLS).expect("MAX_NUMBER_CHAINED_CALLS fits in u32") + 1,
+        None,
+    );
+
+    let message = public_transaction::Message::try_new(
+        program.id(),
+        vec![to, from], // The chain_caller program permutes the account order in the chain
+        // call
+        vec![Nonce(0)],
+        instruction,
+    )
+    .unwrap();
+    let witness_set = public_transaction::WitnessSet::for_message(&message, &[&from_key]);
+    let tx = PublicTransaction::new(message, witness_set);
+
+    let result = state.transition_from_public_transaction(&tx, 1, 0);
+    assert!(matches!(
+        result,
+        Err(LeeError::MaxChainedCallsDepthExceeded)
+    ));
+}
+
+#[test]
+fn execution_that_requires_authentication_of_a_program_derived_account_id_succeeds() {
+    let chain_caller = crate::test_methods::chain_caller();
+    let pda_seed = PdaSeed::new([37; 32]);
+    let from = AccountId::for_public_pda(&chain_caller.id(), &pda_seed);
+    let to = AccountId::new([2; 32]);
+    let initial_balance = 1000;
+    let initial_data = [(from, initial_balance), (to, 0)];
+    let mut state = V03State::new()
+        .with_public_accounts(public_state_from_balances(&initial_data))
+        .with_test_programs();
+    let amount: u128 = 58;
+    let instruction: (u128, ProgramId, u32, Option<PdaSeed>) = (
+        amount,
+        crate::test_methods::simple_balance_transfer().id(),
+        1,
+        Some(pda_seed),
+    );
+
+    let expected_to_post = Account {
+        program_owner: crate::test_methods::simple_balance_transfer().id(),
+        balance: amount, // The `chain_caller` chains the program twice
+        ..Account::default()
+    };
+    let message = public_transaction::Message::try_new(
+        chain_caller.id(),
+        vec![to, from], // The chain_caller program permutes the account order in the chain
+        // call
+        vec![],
+        instruction,
+    )
+    .unwrap();
+    let witness_set = public_transaction::WitnessSet::for_message(&message, &[]);
+    let tx = PublicTransaction::new(message, witness_set);
+
+    state.transition_from_public_transaction(&tx, 1, 0).unwrap();
+
+    let from_post = state.get_account_by_id(from);
+    let to_post = state.get_account_by_id(to);
+    assert_eq!(from_post.balance, initial_balance - amount);
+    assert_eq!(to_post, expected_to_post);
+}
+
+#[test]
+fn claiming_mechanism_within_chain_call() {
+    // This test calls the authenticated transfer program through the chain_caller program.
+    // The transfer is made from an initialized sender to an uninitialized recipient. And
+    // it is expected that the recipient account is claimed by the authenticated transfer
+    // program and not the chained_caller program.
+    let chain_caller = crate::test_methods::chain_caller();
+    let simple_transfer = crate::test_methods::simple_balance_transfer();
+    let from_key = PrivateKey::try_new([1; 32]).unwrap();
+    let from = AccountId::from(&PublicKey::new_from_private_key(&from_key));
+    let initial_balance = 100;
+    let initial_data = [(from, initial_balance)];
+    let mut state = V03State::new()
+        .with_public_accounts(public_state_from_balances(&initial_data))
+        .with_test_programs();
+    let to_key = PrivateKey::try_new([2; 32]).unwrap();
+    let to = AccountId::from(&PublicKey::new_from_private_key(&to_key));
+    let amount: u128 = 37;
+
+    // Check the recipient is an uninitialized account
+    assert_eq!(state.get_account_by_id(to), Account::default());
+
+    let expected_to_post = Account {
+        // The expected program owner is the authenticated transfer program
+        program_owner: simple_transfer.id(),
+        balance: amount,
+        nonce: Nonce(1),
+        ..Account::default()
+    };
+
+    // The transaction executes the chain_caller program, which internally calls the
+    // authenticated_transfer program
+    let instruction: (u128, ProgramId, u32, Option<PdaSeed>) = (
+        amount,
+        crate::test_methods::simple_balance_transfer().id(),
+        1,
+        None,
+    );
+    let message = public_transaction::Message::try_new(
+        chain_caller.id(),
+        vec![to, from], // The chain_caller program permutes the account order in the chain
+        // call
+        vec![Nonce(0), Nonce(0)],
+        instruction,
+    )
+    .unwrap();
+    let witness_set = public_transaction::WitnessSet::for_message(&message, &[&from_key, &to_key]);
+    let tx = PublicTransaction::new(message, witness_set);
+
+    state.transition_from_public_transaction(&tx, 1, 0).unwrap();
+
+    let from_post = state.get_account_by_id(from);
+    let to_post = state.get_account_by_id(to);
+    assert_eq!(from_post.balance, initial_balance - amount);
+    assert_eq!(to_post, expected_to_post);
+}
+
+#[test]
+fn unauthorized_public_account_claiming_fails_when_executed_privately() {
+    let program = crate::test_methods::simple_balance_transfer();
+    let account_id = AccountId::new([11; 32]);
+    let public_account = AccountWithMetadata::new(Account::default(), false, account_id);
+
+    let result = execute_and_prove(
+        vec![public_account],
+        Program::serialize_instruction(0_u128).unwrap(),
+        vec![InputAccountIdentity::Public],
+        &program.into(),
+    );
+
+    assert!(matches!(result, Err(LeeError::CircuitProvingError(_))));
+}
+
+#[test]
+fn authorized_public_account_claiming_succeeds_when_executed_privately() {
+    let program = crate::test_methods::simple_balance_transfer();
+    let program_id = program.id();
+    let sender_keys = test_private_account_keys_1();
+    let sender_private_account = Account {
+        program_owner: program_id,
+        balance: 100,
+        ..Account::default()
+    };
+    let sender_account_id = AccountId::for_regular_private_account(&sender_keys.npk(), 0);
+    let sender_commitment = Commitment::new(&sender_account_id, &sender_private_account);
+    let sender_init_nullifier = Nullifier::for_account_initialization(&sender_account_id);
+    let mut state =
+        V03State::new().with_private_accounts([(sender_commitment.clone(), sender_init_nullifier)]);
+    let sender_pre =
+        AccountWithMetadata::new(sender_private_account, true, (&sender_keys.npk(), 0));
+    let recipient_private_key = PrivateKey::try_new([2; 32]).unwrap();
+    let recipient_account_id =
+        AccountId::from(&PublicKey::new_from_private_key(&recipient_private_key));
+    let recipient_pre = AccountWithMetadata::new(Account::default(), true, recipient_account_id);
+    let (shared_secret, epk) =
+        SharedSecretKey::encapsulate_deterministic(&sender_keys.vpk(), &[0_u8; 32], 0);
+
+    let balance = 37;
+
+    let (output, proof) = execute_and_prove(
+        vec![sender_pre, recipient_pre],
+        Program::serialize_instruction(balance).unwrap(),
+        vec![
+            InputAccountIdentity::PrivateAuthorizedUpdate {
+                epk,
+                view_tag: EncryptedAccountData::compute_view_tag(
+                    &sender_keys.npk(),
+                    &sender_keys.vpk(),
+                ),
+                ssk: shared_secret,
+                nsk: sender_keys.nsk,
+                membership_proof: state
+                    .get_proof_for_commitment(&sender_commitment)
+                    .expect("sender's commitment must be in state"),
+                identifier: 0,
+            },
+            InputAccountIdentity::Public,
+        ],
+        &program.into(),
+    )
+    .unwrap();
+
+    let message =
+        Message::try_from_circuit_output(vec![recipient_account_id], vec![Nonce(0)], output)
+            .unwrap();
+
+    let witness_set = WitnessSet::for_message(&message, proof, &[&recipient_private_key]);
+    let tx = PrivacyPreservingTransaction::new(message, witness_set);
+
+    state
+        .transition_from_privacy_preserving_transaction(&tx, 1, 0)
+        .unwrap();
+
+    let nullifier = Nullifier::for_account_update(&sender_commitment, &sender_keys.nsk);
+    assert!(state.private_state.1.contains(&nullifier));
+
+    assert_eq!(
+        state.get_account_by_id(recipient_account_id),
+        Account {
+            program_owner: program_id,
+            balance,
+            nonce: Nonce(1),
+            ..Account::default()
+        }
+    );
+}
+
+#[test_case::test_case(1; "single call")]
+#[test_case::test_case(2; "two calls")]
+fn private_chained_call(number_of_calls: u32) {
+    // Arrange
+    let chain_caller = crate::test_methods::chain_caller();
+    let simple_transfers = crate::test_methods::simple_balance_transfer();
+    let from_keys = test_private_account_keys_1();
+    let to_keys = test_private_account_keys_2();
+    let initial_balance = 100;
+    let from_account = AccountWithMetadata::new(
+        Account {
+            program_owner: simple_transfers.id(),
+            balance: initial_balance,
+            ..Account::default()
+        },
+        true,
+        (&from_keys.npk(), 0),
+    );
+    let to_account = AccountWithMetadata::new(
+        Account {
+            program_owner: simple_transfers.id(),
+            ..Account::default()
+        },
+        true,
+        (&to_keys.npk(), 0),
+    );
+
+    let from_account_id = AccountId::for_regular_private_account(&from_keys.npk(), 0);
+    let to_account_id = AccountId::for_regular_private_account(&to_keys.npk(), 0);
+    let from_commitment = Commitment::new(&from_account_id, &from_account.account);
+    let to_commitment = Commitment::new(&to_account_id, &to_account.account);
+    let from_init_nullifier = Nullifier::for_account_initialization(&from_account_id);
+    let to_init_nullifier = Nullifier::for_account_initialization(&to_account_id);
+    let mut state = V03State::new()
+        .with_private_accounts([
+            (from_commitment.clone(), from_init_nullifier),
+            (to_commitment.clone(), to_init_nullifier),
+        ])
+        .with_test_programs();
+    let amount: u128 = 37;
+    let instruction: (u128, ProgramId, u32, Option<PdaSeed>) = (
+        amount,
+        crate::test_methods::simple_balance_transfer().id(),
+        number_of_calls,
+        None,
+    );
+
+    let (from_ss, from_epk) =
+        SharedSecretKey::encapsulate_deterministic(&from_keys.vpk(), &[0_u8; 32], 0);
+
+    let (to_ss, to_epk) =
+        SharedSecretKey::encapsulate_deterministic(&to_keys.vpk(), &[0_u8; 32], 1);
+
+    let mut dependencies = HashMap::new();
+
+    dependencies.insert(simple_transfers.id(), simple_transfers);
+    let program_with_deps = ProgramWithDependencies::new(chain_caller, dependencies);
+
+    let from_new_nonce = Nonce::default().private_account_nonce_increment(&from_keys.nsk);
+    let to_new_nonce = Nonce::default().private_account_nonce_increment(&to_keys.nsk);
+
+    let from_expected_post = Account {
+        balance: initial_balance - u128::from(number_of_calls) * amount,
+        nonce: from_new_nonce,
+        ..from_account.account.clone()
+    };
+    let from_expected_commitment = Commitment::new(&from_account_id, &from_expected_post);
+
+    let to_expected_post = Account {
+        balance: u128::from(number_of_calls) * amount,
+        nonce: to_new_nonce,
+        ..to_account.account.clone()
+    };
+    let to_expected_commitment = Commitment::new(&to_account_id, &to_expected_post);
+
+    // Act
+    let (output, proof) = execute_and_prove(
+        vec![to_account, from_account],
+        Program::serialize_instruction(instruction).unwrap(),
+        vec![
+            InputAccountIdentity::PrivateAuthorizedUpdate {
+                epk: to_epk,
+                view_tag: EncryptedAccountData::compute_view_tag(&to_keys.npk(), &to_keys.vpk()),
+                ssk: to_ss,
+                nsk: from_keys.nsk,
+                membership_proof: state
+                    .get_proof_for_commitment(&from_commitment)
+                    .expect("from's commitment must be in state"),
+                identifier: 0,
+            },
+            InputAccountIdentity::PrivateAuthorizedUpdate {
+                epk: from_epk,
+                view_tag: EncryptedAccountData::compute_view_tag(
+                    &from_keys.npk(),
+                    &from_keys.vpk(),
+                ),
+                ssk: from_ss,
+                nsk: to_keys.nsk,
+                membership_proof: state
+                    .get_proof_for_commitment(&to_commitment)
+                    .expect("to's commitment must be in state"),
+                identifier: 0,
+            },
+        ],
+        &program_with_deps,
+    )
+    .unwrap();
+
+    let message = Message::try_from_circuit_output(vec![], vec![], output).unwrap();
+    let witness_set = WitnessSet::for_message(&message, proof, &[]);
+    let transaction = PrivacyPreservingTransaction::new(message, witness_set);
+
+    state
+        .transition_from_privacy_preserving_transaction(&transaction, 1, 0)
+        .unwrap();
+
+    // Assert
+    assert!(
+        state
+            .get_proof_for_commitment(&from_expected_commitment)
+            .is_some()
+    );
+    assert!(
+        state
+            .get_proof_for_commitment(&to_expected_commitment)
+            .is_some()
+    );
+}
+
+#[test]
+fn claiming_mechanism_cannot_claim_initialied_accounts() {
+    let claimer = crate::test_methods::claimer();
+    let mut state = V03State::new().with_test_programs();
+    let account_id = AccountId::new([2; 32]);
+
+    // Insert an account with non-default program owner
+    state.force_insert_account(
+        account_id,
+        Account {
+            program_owner: [1, 2, 3, 4, 5, 6, 7, 8],
+            ..Account::default()
+        },
+    );
+
+    let message =
+        public_transaction::Message::try_new(claimer.id(), vec![account_id], vec![], ()).unwrap();
+    let witness_set = public_transaction::WitnessSet::for_message(&message, &[]);
+    let tx = PublicTransaction::new(message, witness_set);
+
+    let result = state.transition_from_public_transaction(&tx, 1, 0);
+
+    assert!(matches!(
+        result,
+        Err(LeeError::InvalidProgramBehavior(
+            InvalidProgramBehaviorError::ClaimedNonDefaultAccount { account_id: err_account_id }
+        )) if err_account_id == account_id
+    ));
+}
+
+/// This test ensures that even if a malicious program tries to perform overflow of balances
+/// it will not be able to break the balance validation.
+#[test]
+fn malicious_program_cannot_break_balance_validation_if_not_in_genesis() {
+    let sender_key = PrivateKey::try_new([37; 32]).unwrap();
+    let sender_id = AccountId::from(&PublicKey::new_from_private_key(&sender_key));
+    let sender_init_balance: u128 = 10;
+
+    let recipient_key = PrivateKey::try_new([42; 32]).unwrap();
+    let recipient_id = AccountId::from(&PublicKey::new_from_private_key(&recipient_key));
+    let recipient_init_balance: u128 = 10;
+
+    let modified_transfer_id = crate::test_methods::modified_transfer_program().id();
+
+    let mut state = V03State::new()
+        .with_public_accounts([
+            (
+                sender_id,
+                Account {
+                    program_owner: modified_transfer_id,
+                    balance: sender_init_balance,
+                    ..Account::default()
+                },
+            ),
+            (
+                recipient_id,
+                Account {
+                    program_owner: modified_transfer_id,
+                    balance: recipient_init_balance,
+                    ..Account::default()
+                },
+            ),
+        ])
+        .with_test_programs();
+
+    let balance_to_move: u128 = 4;
+
+    let sender = AccountWithMetadata::new(state.get_account_by_id(sender_id), true, sender_id);
+
+    let sender_nonce = sender.account.nonce;
+
+    let _recipient =
+        AccountWithMetadata::new(state.get_account_by_id(recipient_id), false, sender_id);
+
+    let message = public_transaction::Message::try_new(
+        modified_transfer_id,
+        vec![sender_id, recipient_id],
+        vec![sender_nonce],
+        balance_to_move,
+    )
+    .unwrap();
+
+    let witness_set = public_transaction::WitnessSet::for_message(&message, &[&sender_key]);
+    let tx = PublicTransaction::new(message, witness_set);
+    let res = state.transition_from_public_transaction(&tx, 2, 0);
+    let expected_total_balance_pre_states =
+        WrappedBalanceSum::from_balances([sender_init_balance, recipient_init_balance].into_iter())
+            .unwrap();
+    let expected_total_balance_post_states = WrappedBalanceSum::from_balances(
+        [sender_init_balance, recipient_init_balance, u128::MAX, 1].into_iter(),
+    )
+    .unwrap();
+    assert!(matches!(
+        res,
+        Err(LeeError::InvalidProgramBehavior(
+            InvalidProgramBehaviorError::ExecutionValidationFailed(
+                ExecutionValidationError::MismatchedTotalBalance { total_balance_pre_states, total_balance_post_states }
+            )
+        )) if total_balance_pre_states == expected_total_balance_pre_states && total_balance_post_states == expected_total_balance_post_states
+    ));
+
+    let sender_post = state.get_account_by_id(sender_id);
+    let recipient_post = state.get_account_by_id(recipient_id);
+
+    let expected_sender_post = {
+        let mut this = state.get_account_by_id(sender_id);
+        this.balance = sender_init_balance;
+        this.nonce = Nonce(0);
+        this
+    };
+
+    let expected_recipient_post = {
+        let mut this = state.get_account_by_id(sender_id);
+        this.balance = recipient_init_balance;
+        this.nonce = Nonce(0);
+        this
+    };
+
+    assert_eq!(expected_sender_post, sender_post);
+    assert_eq!(expected_recipient_post, recipient_post);
+}

lee/state_machine/src/state/tests/flash_swap.rs

@@ -0,0 +1,235 @@
+use super::*;
+
+#[test]
+fn flash_swap_successful() {
+    let initiator = crate::test_methods::flash_swap_initiator();
+    let callback = crate::test_methods::flash_swap_callback();
+    let token = crate::test_methods::simple_balance_transfer();
+
+    let vault_id = AccountId::for_public_pda(&initiator.id(), &PdaSeed::new([0_u8; 32]));
+    let receiver_id = AccountId::for_public_pda(&callback.id(), &PdaSeed::new([1_u8; 32]));
+
+    let initial_balance: u128 = 1000;
+    let amount_out: u128 = 100;
+
+    let vault_account = Account {
+        program_owner: token.id(),
+        balance: initial_balance,
+        ..Account::default()
+    };
+    let receiver_account = Account {
+        program_owner: token.id(),
+        balance: 0,
+        ..Account::default()
+    };
+
+    let mut state = V03State::new().with_test_programs();
+    state.force_insert_account(vault_id, vault_account);
+    state.force_insert_account(receiver_id, receiver_account);
+
+    // Callback instruction: return funds
+    let cb_instruction = CallbackInstruction {
+        return_funds: true,
+        token_program_id: token.id(),
+        amount: amount_out,
+    };
+    let cb_data = Program::serialize_instruction(cb_instruction).unwrap();
+
+    let instruction = FlashSwapInstruction::Initiate {
+        token_program_id: token.id(),
+        callback_program_id: callback.id(),
+        amount_out,
+        callback_instruction_data: cb_data,
+    };
+
+    let tx = build_flash_swap_tx(&initiator, vault_id, receiver_id, instruction);
+    let result = state.transition_from_public_transaction(&tx, 1, 0);
+    assert!(result.is_ok(), "flash swap should succeed: {result:?}");
+
+    // Vault balance restored, receiver back to 0
+    assert_eq!(state.get_account_by_id(vault_id).balance, initial_balance);
+    assert_eq!(state.get_account_by_id(receiver_id).balance, 0);
+}
+
+#[test]
+fn flash_swap_callback_keeps_funds_rollback() {
+    let initiator = crate::test_methods::flash_swap_initiator();
+    let callback = crate::test_methods::flash_swap_callback();
+    let token = crate::test_methods::simple_balance_transfer();
+
+    let vault_id = AccountId::for_public_pda(&initiator.id(), &PdaSeed::new([0_u8; 32]));
+    let receiver_id = AccountId::for_public_pda(&callback.id(), &PdaSeed::new([1_u8; 32]));
+
+    let initial_balance: u128 = 1000;
+    let amount_out: u128 = 100;
+
+    let vault_account = Account {
+        program_owner: token.id(),
+        balance: initial_balance,
+        ..Account::default()
+    };
+    let receiver_account = Account {
+        program_owner: token.id(),
+        balance: 0,
+        ..Account::default()
+    };
+
+    let mut state = V03State::new().with_test_programs();
+    state.force_insert_account(vault_id, vault_account);
+    state.force_insert_account(receiver_id, receiver_account);
+
+    // Callback instruction: do NOT return funds
+    let cb_instruction = CallbackInstruction {
+        return_funds: false,
+        token_program_id: token.id(),
+        amount: amount_out,
+    };
+    let cb_data = Program::serialize_instruction(cb_instruction).unwrap();
+
+    let instruction = FlashSwapInstruction::Initiate {
+        token_program_id: token.id(),
+        callback_program_id: callback.id(),
+        amount_out,
+        callback_instruction_data: cb_data,
+    };
+
+    let tx = build_flash_swap_tx(&initiator, vault_id, receiver_id, instruction);
+    let result = state.transition_from_public_transaction(&tx, 1, 0);
+
+    // Invariant check fails → entire tx rolls back
+    assert!(
+        result.is_err(),
+        "flash swap should fail when callback keeps funds"
+    );
+
+    // State unchanged (rollback)
+    assert_eq!(state.get_account_by_id(vault_id).balance, initial_balance);
+    assert_eq!(state.get_account_by_id(receiver_id).balance, 0);
+}
+
+#[test]
+fn flash_swap_self_call_targets_correct_program() {
+    // Zero-amount flash swap: the invariant self-call still runs and succeeds
+    // because vault balance doesn't decrease.
+    let initiator = crate::test_methods::flash_swap_initiator();
+    let callback = crate::test_methods::flash_swap_callback();
+    let token = crate::test_methods::simple_balance_transfer();
+
+    let vault_id = AccountId::for_public_pda(&initiator.id(), &PdaSeed::new([0_u8; 32]));
+    let receiver_id = AccountId::for_public_pda(&callback.id(), &PdaSeed::new([1_u8; 32]));
+
+    let initial_balance: u128 = 1000;
+
+    let vault_account = Account {
+        program_owner: token.id(),
+        balance: initial_balance,
+        ..Account::default()
+    };
+    let receiver_account = Account {
+        program_owner: token.id(),
+        balance: 0,
+        ..Account::default()
+    };
+
+    let mut state = V03State::new().with_test_programs();
+    state.force_insert_account(vault_id, vault_account);
+    state.force_insert_account(receiver_id, receiver_account);
+
+    let cb_instruction = CallbackInstruction {
+        return_funds: true,
+        token_program_id: token.id(),
+        amount: 0,
+    };
+    let cb_data = Program::serialize_instruction(cb_instruction).unwrap();
+
+    let instruction = FlashSwapInstruction::Initiate {
+        token_program_id: token.id(),
+        callback_program_id: callback.id(),
+        amount_out: 0,
+        callback_instruction_data: cb_data,
+    };
+
+    let tx = build_flash_swap_tx(&initiator, vault_id, receiver_id, instruction);
+    let result = state.transition_from_public_transaction(&tx, 1, 0);
+    assert!(
+        result.is_ok(),
+        "zero-amount flash swap should succeed: {result:?}"
+    );
+}
+
+#[test]
+fn flash_swap_standalone_invariant_check_rejected() {
+    // Calling InvariantCheck directly (not as a chained self-call) should fail
+    // because caller_program_id will be None.
+    let initiator = crate::test_methods::flash_swap_initiator();
+    let token = crate::test_methods::simple_balance_transfer();
+
+    let vault_id = AccountId::for_public_pda(&initiator.id(), &PdaSeed::new([0_u8; 32]));
+
+    let vault_account = Account {
+        program_owner: token.id(),
+        balance: 1000,
+        ..Account::default()
+    };
+
+    let mut state = V03State::new().with_test_programs();
+    state.force_insert_account(vault_id, vault_account);
+
+    let instruction = FlashSwapInstruction::InvariantCheck {
+        min_vault_balance: 1000,
+    };
+
+    let message =
+        public_transaction::Message::try_new(initiator.id(), vec![vault_id], vec![], instruction)
+            .unwrap();
+    let witness_set = public_transaction::WitnessSet::for_message(&message, &[]);
+    let tx = PublicTransaction::new(message, witness_set);
+
+    let result = state.transition_from_public_transaction(&tx, 1, 0);
+    assert!(
+        result.is_err(),
+        "standalone InvariantCheck should be rejected (caller_program_id is None)"
+    );
+}
+
+#[test]
+fn malicious_self_program_id_rejected_in_public_execution() {
+    let program = crate::test_methods::malicious_self_program_id();
+    let acc_id = AccountId::new([99; 32]);
+    let account = Account::default();
+
+    let mut state = V03State::new().with_test_programs();
+    state.force_insert_account(acc_id, account);
+
+    let message =
+        public_transaction::Message::try_new(program.id(), vec![acc_id], vec![], ()).unwrap();
+    let witness_set = public_transaction::WitnessSet::for_message(&message, &[]);
+    let tx = PublicTransaction::new(message, witness_set);
+
+    let result = state.transition_from_public_transaction(&tx, 1, 0);
+    assert!(
+        result.is_err(),
+        "program with wrong self_program_id in output should be rejected"
+    );
+}
+
+#[test]
+fn malicious_caller_program_id_rejected_in_public_execution() {
+    let program = crate::test_methods::malicious_caller_program_id();
+    let acc_id = AccountId::new([99; 32]);
+    let account = Account::default();
+
+    let mut state = V03State::new().with_test_programs();
+    state.force_insert_account(acc_id, account);
+
+    let message =
+        public_transaction::Message::try_new(program.id(), vec![acc_id], vec![], ()).unwrap();
+    let witness_set = public_transaction::WitnessSet::for_message(&message, &[]);
+    let tx = PublicTransaction::new(message, witness_set);
+
+    let result = state.transition_from_public_transaction(&tx, 1, 0);
+    assert!(
+        result.is_err(),
+        "program with spoofed caller_program_id in output should be rejected"
+    );
+}

lee/state_machine/src/state/tests/genesis.rs

@@ -0,0 +1,128 @@
+use super::*;
+
+#[test]
+fn new_works() {
+    let key1 = PrivateKey::try_new([1; 32]).unwrap();
+    let key2 = PrivateKey::try_new([2; 32]).unwrap();
+    let addr1 = AccountId::from(&PublicKey::new_from_private_key(&key1));
+    let addr2 = AccountId::from(&PublicKey::new_from_private_key(&key2));
+    let expected_public_state = {
+        let mut this = HashMap::new();
+        this.insert(
+            addr1,
+            Account {
+                balance: 100,
+                ..Account::default()
+            },
+        );
+        this.insert(
+            addr2,
+            Account {
+                balance: 151,
+                ..Account::default()
+            },
+        );
+        this
+    };
+    let expected_builtin_programs = HashMap::new();
+
+    let state =
+        V03State::new().with_public_account_balances([(addr1, 100_u128), (addr2, 151_u128)]);
+
+    assert_eq!(state.public_state, expected_public_state);
+    assert_eq!(state.programs, expected_builtin_programs);
+}
+
+#[test]
+fn new_includes_nullifiers_for_private_accounts() {
+    let keys1 = test_private_account_keys_1();
+    let keys2 = test_private_account_keys_2();
+
+    let account = Account {
+        balance: 100,
+        ..Account::default()
+    };
+
+    let account_id1 = AccountId::for_regular_private_account(&keys1.npk(), 0);
+    let account_id2 = AccountId::for_regular_private_account(&keys2.npk(), 0);
+
+    let init_commitment1 = Commitment::new(&account_id1, &account);
+    let init_commitment2 = Commitment::new(&account_id2, &account);
+    let init_nullifier1 = Nullifier::for_account_initialization(&account_id1);
+    let init_nullifier2 = Nullifier::for_account_initialization(&account_id2);
+
+    let initial_private_accounts = vec![
+        (init_commitment1, init_nullifier1),
+        (init_commitment2, init_nullifier2),
+    ];
+
+    let state = V03State::new().with_private_accounts(initial_private_accounts);
+
+    assert!(state.private_state.1.contains(&init_nullifier1));
+    assert!(state.private_state.1.contains(&init_nullifier2));
+}
+
+#[test]
+fn insert_program() {
+    let mut state = V03State::new();
+    let program_to_insert = crate::test_methods::simple_balance_transfer();
+    let program_id = program_to_insert.id();
+    assert!(!state.programs.contains_key(&program_id));
+
+    state.insert_program(program_to_insert);
+
+    assert!(state.programs.contains_key(&program_id));
+}
+
+#[test]
+fn get_account_by_account_id_non_default_account() {
+    let key = PrivateKey::try_new([1; 32]).unwrap();
+    let account_id = AccountId::from(&PublicKey::new_from_private_key(&key));
+    let initial_data = [(
+        account_id,
+        Account {
+            program_owner: crate::test_methods::simple_balance_transfer().id(),
+            balance: 100,
+            ..Account::default()
+        },
+    )];
+    let state = V03State::new().with_public_accounts(initial_data);
+    let expected_account = &state.public_state[&account_id];
+
+    let account = state.get_account_by_id(account_id);
+
+    assert_eq!(&account, expected_account);
+}
+
+#[test]
+fn get_account_by_account_id_default_account() {
+    let addr2 = AccountId::new([0; 32]);
+    let state = V03State::new();
+    let expected_account = Account::default();
+
+    let account = state.get_account_by_id(addr2);
+
+    assert_eq!(account, expected_account);
+}
+
+#[test]
+fn builtin_programs_getter() {
+    let state = V03State::new();
+
+    let builtin_programs = state.programs();
+
+    assert_eq!(builtin_programs, &state.programs);
+}
+
+#[test]
+fn state_serialization_roundtrip() {
+    let account_id_1 = AccountId::new([1; 32]);
+    let account_id_2 = AccountId::new([2; 32]);
+    let initial_data = [(account_id_1, 100_u128), (account_id_2, 151_u128)];
+    let state = V03State::new()
+        .with_public_accounts(public_state_from_balances(&initial_data))
+        .with_test_programs();
+    let bytes = borsh::to_vec(&state).unwrap();
+    let state_from_bytes: V03State = borsh::from_slice(&bytes).unwrap();
+    assert_eq!(state, state_from_bytes);
+}

lee/state_machine/src/state/tests/mod.rs

@@ -0,0 +1,444 @@
+#![expect(
+    clippy::arithmetic_side_effects,
+    clippy::shadow_unrelated,
+    reason = "We don't care about it in tests"
+)]
+
+use std::collections::HashMap;
+
+use lee_core::{
+    BlockId, Commitment, DUMMY_COMMITMENT_HASH, EncryptedAccountData, InputAccountIdentity,
+    Nullifier, NullifierPublicKey, NullifierSecretKey, SharedSecretKey, Timestamp,
+    account::{Account, AccountId, AccountWithMetadata, Nonce, data::Data},
+    encryption::{EphemeralPublicKey, ML_KEM_768_CIPHERTEXT_LEN, ViewingPublicKey},
+    program::{
+        BlockValidityWindow, ExecutionValidationError, MAX_NUMBER_CHAINED_CALLS, PdaSeed,
+        ProgramId, TimestampValidityWindow, WrappedBalanceSum,
+    },
+};
+
+use crate::{
+    PublicKey, PublicTransaction, V03State,
+    error::{InvalidProgramBehaviorError, LeeError},
+    execute_and_prove,
+    privacy_preserving_transaction::{
+        PrivacyPreservingTransaction, circuit::ProgramWithDependencies, message::Message,
+        witness_set::WitnessSet,
+    },
+    program::Program,
+    public_transaction,
+    signature::PrivateKey,
+};
+
+mod authenticated_transfer;
+mod changer_claimer;
+mod circuit;
+mod claiming;
+mod flash_swap;
+mod genesis;
+mod privacy_preserving;
+mod public_program_rules;
+mod validity_window;
+
+impl V03State {
+    /// Include test programs in the builtin programs map.
+    #[must_use]
+    pub fn with_test_programs(mut self) -> Self {
+        self.insert_program(crate::test_methods::simple_balance_transfer());
+        self.insert_program(crate::test_methods::nonce_changer());
+        self.insert_program(crate::test_methods::extra_output());
+        self.insert_program(crate::test_methods::missing_output());
+        self.insert_program(crate::test_methods::program_owner_changer());
+        self.insert_program(crate::test_methods::data_changer());
+        self.insert_program(crate::test_methods::minter());
+        self.insert_program(crate::test_methods::burner());
+        self.insert_program(crate::test_methods::auth_asserting_noop());
+        self.insert_program(crate::test_methods::private_pda_delegator());
+        self.insert_program(crate::test_methods::pda_claimer());
+        self.insert_program(crate::test_methods::two_pda_claimer());
+        self.insert_program(crate::test_methods::noop());
+        self.insert_program(crate::test_methods::chain_caller());
+        self.insert_program(crate::test_methods::modified_transfer_program());
+        self.insert_program(crate::test_methods::malicious_authorization_changer());
+        self.insert_program(crate::test_methods::validity_window());
+        self.insert_program(crate::test_methods::flash_swap_initiator());
+        self.insert_program(crate::test_methods::flash_swap_callback());
+        self.insert_program(crate::test_methods::malicious_self_program_id());
+        self.insert_program(crate::test_methods::malicious_caller_program_id());
+        self.insert_program(crate::test_methods::pda_spend_proxy());
+        self.insert_program(crate::test_methods::claimer());
+        self.insert_program(crate::test_methods::changer_claimer());
+        self.insert_program(crate::test_methods::validity_window_chain_caller());
+        self.insert_program(crate::test_methods::simple_transfer_proxy());
+        self.insert_program(crate::test_methods::malicious_injector());
+        self.insert_program(crate::test_methods::malicious_launderer());
+        self.insert_program(crate::test_methods::modified_transfer_program());
+        self
+    }
+
+    #[must_use]
+    pub fn with_non_default_accounts_but_default_program_owners(mut self) -> Self {
+        let account_with_default_values_except_balance = Account {
+            balance: 100,
+            ..Account::default()
+        };
+        let account_with_default_values_except_nonce = Account {
+            nonce: Nonce(37),
+            ..Account::default()
+        };
+        let account_with_default_values_except_data = Account {
+            data: vec![0xca, 0xfe].try_into().unwrap(),
+            ..Account::default()
+        };
+        self.force_insert_account(
+            AccountId::new([255; 32]),
+            account_with_default_values_except_balance,
+        );
+        self.force_insert_account(
+            AccountId::new([254; 32]),
+            account_with_default_values_except_nonce,
+        );
+        self.force_insert_account(
+            AccountId::new([253; 32]),
+            account_with_default_values_except_data,
+        );
+        self
+    }
+
+    #[must_use]
+    pub fn with_account_owned_by_burner_program(mut self) -> Self {
+        let account = Account {
+            program_owner: crate::test_methods::burner().id(),
+            balance: 100,
+            ..Default::default()
+        };
+        self.force_insert_account(AccountId::new([252; 32]), account);
+        self
+    }
+
+    #[must_use]
+    pub fn with_private_account(mut self, keys: &TestPrivateKeys, account: &Account) -> Self {
+        let account_id = AccountId::for_regular_private_account(&keys.npk(), 0);
+        let commitment = Commitment::new(&account_id, account);
+        self.private_state.0.extend(&[commitment]);
+        self
+    }
+}
+
+pub struct TestPublicKeys {
+    pub signing_key: PrivateKey,
+}
+
+impl TestPublicKeys {
+    pub fn account_id(&self) -> AccountId {
+        AccountId::from(&PublicKey::new_from_private_key(&self.signing_key))
+    }
+}
+
+pub struct TestPrivateKeys {
+    pub nsk: NullifierSecretKey,
+    pub d: [u8; 32],
+    pub z: [u8; 32],
+}
+
+impl TestPrivateKeys {
+    pub fn npk(&self) -> NullifierPublicKey {
+        NullifierPublicKey::from(&self.nsk)
+    }
+
+    pub fn vpk(&self) -> ViewingPublicKey {
+        ViewingPublicKey::from_seed(&self.d, &self.z)
+    }
+}
+
+// ── Flash Swap types (mirrors of guest types for host-side serialisation) ──
+
+#[derive(serde::Serialize, serde::Deserialize)]
+struct CallbackInstruction {
+    return_funds: bool,
+    token_program_id: ProgramId,
+    amount: u128,
+}
+
+#[derive(serde::Serialize, serde::Deserialize)]
+enum FlashSwapInstruction {
+    Initiate {
+        token_program_id: ProgramId,
+        callback_program_id: ProgramId,
+        amount_out: u128,
+        callback_instruction_data: Vec<u32>,
+    },
+    InvariantCheck {
+        min_vault_balance: u128,
+    },
+}
+
+fn public_state_from_balances(initial_data: &[(AccountId, u128)]) -> HashMap<AccountId, Account> {
+    initial_data
+        .iter()
+        .copied()
+        .map(|(account_id, balance)| {
+            (
+                account_id,
+                Account {
+                    program_owner: crate::test_methods::simple_balance_transfer().id(),
+                    balance,
+                    ..Account::default()
+                },
+            )
+        })
+        .collect()
+}
+
+fn transfer_transaction(
+    from: AccountId,
+    from_key: &PrivateKey,
+    from_nonce: u128,
+    to: AccountId,
+    to_key: &PrivateKey,
+    to_nonce: u128,
+    balance: u128,
+) -> PublicTransaction {
+    let account_ids = vec![from, to];
+    let nonces = vec![Nonce(from_nonce), Nonce(to_nonce)];
+    let program_id = crate::test_methods::simple_balance_transfer().id();
+    let message =
+        public_transaction::Message::try_new(program_id, account_ids, nonces, balance).unwrap();
+    let witness_set = public_transaction::WitnessSet::for_message(&message, &[from_key, to_key]);
+    PublicTransaction::new(message, witness_set)
+}
+
+fn build_flash_swap_tx(
+    initiator: &Program,
+    vault_id: AccountId,
+    receiver_id: AccountId,
+    instruction: FlashSwapInstruction,
+) -> PublicTransaction {
+    let message = public_transaction::Message::try_new(
+        initiator.id(),
+        vec![vault_id, receiver_id],
+        vec![], // no signers — vault is PDA-authorised
+        instruction,
+    )
+    .unwrap();
+    let witness_set = public_transaction::WitnessSet::for_message(&message, &[]);
+    PublicTransaction::new(message, witness_set)
+}
+
+fn test_public_account_keys_1() -> TestPublicKeys {
+    TestPublicKeys {
+        signing_key: PrivateKey::try_new([37; 32]).unwrap(),
+    }
+}
+
+fn test_public_account_keys_2() -> TestPublicKeys {
+    TestPublicKeys {
+        signing_key: PrivateKey::try_new([38; 32]).unwrap(),
+    }
+}
+
+pub fn test_private_account_keys_1() -> TestPrivateKeys {
+    TestPrivateKeys {
+        nsk: [13; 32],
+        d: [31; 32],
+        z: [32; 32],
+    }
+}
+
+pub fn test_private_account_keys_2() -> TestPrivateKeys {
+    TestPrivateKeys {
+        nsk: [38; 32],
+        d: [83; 32],
+        z: [84; 32],
+    }
+}
+
+fn shielded_balance_transfer_for_tests(
+    sender_keys: &TestPublicKeys,
+    recipient_keys: &TestPrivateKeys,
+    balance_to_move: u128,
+    state: &V03State,
+) -> PrivacyPreservingTransaction {
+    let sender = AccountWithMetadata::new(
+        state.get_account_by_id(sender_keys.account_id()),
+        true,
+        sender_keys.account_id(),
+    );
+
+    let sender_nonce = sender.account.nonce;
+
+    let recipient = AccountWithMetadata::new(Account::default(), false, (&recipient_keys.npk(), 0));
+
+    let (shared_secret, epk) =
+        SharedSecretKey::encapsulate_deterministic(&recipient_keys.vpk(), &[0_u8; 32], 0);
+
+    let (output, proof) = crate::privacy_preserving_transaction::circuit::execute_and_prove(
+        vec![sender, recipient],
+        Program::serialize_instruction(balance_to_move).unwrap(),
+        vec![
+            InputAccountIdentity::Public,
+            InputAccountIdentity::PrivateUnauthorized {
+                epk,
+                view_tag: EncryptedAccountData::compute_view_tag(
+                    &recipient_keys.npk(),
+                    &recipient_keys.vpk(),
+                ),
+                npk: recipient_keys.npk(),
+                ssk: shared_secret,
+                identifier: 0,
+                commitment_root: DUMMY_COMMITMENT_HASH,
+            },
+        ],
+        &crate::test_methods::simple_balance_transfer().into(),
+    )
+    .unwrap();
+
+    let message = Message::try_from_circuit_output(
+        vec![sender_keys.account_id()],
+        vec![sender_nonce],
+        output,
+    )
+    .unwrap();
+
+    let witness_set = WitnessSet::for_message(&message, proof, &[&sender_keys.signing_key]);
+    PrivacyPreservingTransaction::new(message, witness_set)
+}
+
+fn private_balance_transfer_for_tests(
+    sender_keys: &TestPrivateKeys,
+    sender_private_account: &Account,
+    recipient_keys: &TestPrivateKeys,
+    balance_to_move: u128,
+    state: &V03State,
+) -> PrivacyPreservingTransaction {
+    let program = crate::test_methods::simple_balance_transfer();
+    let sender_account_id = AccountId::for_regular_private_account(&sender_keys.npk(), 0);
+    let sender_commitment = Commitment::new(&sender_account_id, sender_private_account);
+    let sender_pre = AccountWithMetadata::new(
+        sender_private_account.clone(),
+        true,
+        (&sender_keys.npk(), 0),
+    );
+    let recipient_pre =
+        AccountWithMetadata::new(Account::default(), false, (&recipient_keys.npk(), 0));
+
+    let (shared_secret_1, epk_1) =
+        SharedSecretKey::encapsulate_deterministic(&sender_keys.vpk(), &[0_u8; 32], 0);
+
+    let (shared_secret_2, epk_2) =
+        SharedSecretKey::encapsulate_deterministic(&recipient_keys.vpk(), &[0_u8; 32], 1);
+
+    let (output, proof) = crate::privacy_preserving_transaction::circuit::execute_and_prove(
+        vec![sender_pre, recipient_pre],
+        Program::serialize_instruction(balance_to_move).unwrap(),
+        vec![
+            InputAccountIdentity::PrivateAuthorizedUpdate {
+                epk: epk_1,
+                view_tag: EncryptedAccountData::compute_view_tag(
+                    &sender_keys.npk(),
+                    &sender_keys.vpk(),
+                ),
+                ssk: shared_secret_1,
+                nsk: sender_keys.nsk,
+                membership_proof: state
+                    .get_proof_for_commitment(&sender_commitment)
+                    .expect("sender's commitment must be in state"),
+                identifier: 0,
+            },
+            InputAccountIdentity::PrivateUnauthorized {
+                epk: epk_2,
+                view_tag: EncryptedAccountData::compute_view_tag(
+                    &recipient_keys.npk(),
+                    &recipient_keys.vpk(),
+                ),
+                npk: recipient_keys.npk(),
+                ssk: shared_secret_2,
+                identifier: 0,
+                commitment_root: DUMMY_COMMITMENT_HASH,
+            },
+        ],
+        &program.into(),
+    )
+    .unwrap();
+
+    let message = Message::try_from_circuit_output(vec![], vec![], output).unwrap();
+
+    let witness_set = WitnessSet::for_message(&message, proof, &[]);
+
+    PrivacyPreservingTransaction::new(message, witness_set)
+}
+
+fn deshielded_balance_transfer_for_tests(
+    sender_keys: &TestPrivateKeys,
+    sender_private_account: &Account,
+    recipient_account_id: &AccountId,
+    balance_to_move: u128,
+    state: &V03State,
+) -> PrivacyPreservingTransaction {
+    let program = crate::test_methods::simple_balance_transfer();
+    let sender_account_id = AccountId::for_regular_private_account(&sender_keys.npk(), 0);
+    let sender_commitment = Commitment::new(&sender_account_id, sender_private_account);
+    let sender_pre = AccountWithMetadata::new(
+        sender_private_account.clone(),
+        true,
+        (&sender_keys.npk(), 0),
+    );
+    let recipient_pre = AccountWithMetadata::new(
+        state.get_account_by_id(*recipient_account_id),
+        false,
+        *recipient_account_id,
+    );
+
+    let (shared_secret, epk) =
+        SharedSecretKey::encapsulate_deterministic(&sender_keys.vpk(), &[0_u8; 32], 0);
+
+    let (output, proof) = crate::privacy_preserving_transaction::circuit::execute_and_prove(
+        vec![sender_pre, recipient_pre],
+        Program::serialize_instruction(balance_to_move).unwrap(),
+        vec![
+            InputAccountIdentity::PrivateAuthorizedUpdate {
+                epk,
+                view_tag: EncryptedAccountData::compute_view_tag(
+                    &sender_keys.npk(),
+                    &sender_keys.vpk(),
+                ),
+                ssk: shared_secret,
+                nsk: sender_keys.nsk,
+                membership_proof: state
+                    .get_proof_for_commitment(&sender_commitment)
+                    .expect("sender's commitment must be in state"),
+                identifier: 0,
+            },
+            InputAccountIdentity::Public,
+        ],
+        &program.into(),
+    )
+    .unwrap();
+
+    let message =
+        Message::try_from_circuit_output(vec![*recipient_account_id], vec![], output).unwrap();
+
+    let witness_set = WitnessSet::for_message(&message, proof, &[]);
+
+    PrivacyPreservingTransaction::new(message, witness_set)
+}
+
+fn valid_private_transfer_tx_and_state() -> (V03State, PrivacyPreservingTransaction) {
+    let sender_keys = test_private_account_keys_1();
+    let sender_private_account = Account {
+        program_owner: crate::test_methods::simple_balance_transfer().id(),
+        balance: 100,
+        nonce: Nonce(0xdead_beef),
+        ..Account::default()
+    };
+    let recipient_keys = test_private_account_keys_2();
+    let state = V03State::new().with_private_account(&sender_keys, &sender_private_account);
+    let tx = private_balance_transfer_for_tests(
+        &sender_keys,
+        &sender_private_account,
+        &recipient_keys,
+        37,
+        &state,
+    );
+    (state, tx)
+}

lee/state_machine/src/state/tests/privacy_preserving.rs

@@ -0,0 +1,539 @@
+use super::*;
+
+#[test]
+fn transition_from_privacy_preserving_transaction_shielded() {
+    let sender_keys = test_public_account_keys_1();
+    let recipient_keys = test_private_account_keys_1();
+
+    let mut state = V03State::new().with_public_accounts([(
+        sender_keys.account_id(),
+        Account {
+            program_owner: crate::test_methods::simple_balance_transfer().id(),
+            balance: 200,
+            ..Account::default()
+        },
+    )]);
+
+    let balance_to_move = 37;
+
+    let tx =
+        shielded_balance_transfer_for_tests(&sender_keys, &recipient_keys, balance_to_move, &state);
+
+    let expected_sender_post = {
+        let mut this = state.get_account_by_id(sender_keys.account_id());
+        this.balance -= balance_to_move;
+        this.nonce.public_account_nonce_increment();
+        this
+    };
+
+    let [expected_new_commitment] = tx.message().new_commitments.clone().try_into().unwrap();
+    assert!(!state.private_state.0.contains(&expected_new_commitment));
+
+    state
+        .transition_from_privacy_preserving_transaction(&tx, 1, 0)
+        .unwrap();
+
+    let sender_post = state.get_account_by_id(sender_keys.account_id());
+    assert_eq!(sender_post, expected_sender_post);
+    assert!(state.private_state.0.contains(&expected_new_commitment));
+
+    assert_eq!(
+        state.get_account_by_id(sender_keys.account_id()).balance,
+        200 - balance_to_move
+    );
+}
+
+#[test]
+fn transition_from_privacy_preserving_transaction_private() {
+    let sender_keys = test_private_account_keys_1();
+    let sender_nonce = Nonce(0xdead_beef);
+
+    let sender_private_account = Account {
+        program_owner: crate::test_methods::simple_balance_transfer().id(),
+        balance: 100,
+        nonce: sender_nonce,
+        data: Data::default(),
+    };
+    let recipient_keys = test_private_account_keys_2();
+
+    let mut state = V03State::new().with_private_account(&sender_keys, &sender_private_account);
+
+    let balance_to_move = 37;
+
+    let tx = private_balance_transfer_for_tests(
+        &sender_keys,
+        &sender_private_account,
+        &recipient_keys,
+        balance_to_move,
+        &state,
+    );
+
+    let sender_account_id = AccountId::for_regular_private_account(&sender_keys.npk(), 0);
+    let recipient_account_id = AccountId::for_regular_private_account(&recipient_keys.npk(), 0);
+    let expected_new_commitment_1 = Commitment::new(
+        &sender_account_id,
+        &Account {
+            program_owner: crate::test_methods::simple_balance_transfer().id(),
+            nonce: sender_nonce.private_account_nonce_increment(&sender_keys.nsk),
+            balance: sender_private_account.balance - balance_to_move,
+            data: Data::default(),
+        },
+    );
+
+    let sender_pre_commitment = Commitment::new(&sender_account_id, &sender_private_account);
+    let expected_new_nullifier =
+        Nullifier::for_account_update(&sender_pre_commitment, &sender_keys.nsk);
+
+    let expected_new_commitment_2 = Commitment::new(
+        &recipient_account_id,
+        &Account {
+            program_owner: crate::test_methods::simple_balance_transfer().id(),
+            nonce: Nonce::private_account_nonce_init(&recipient_account_id),
+            balance: balance_to_move,
+            ..Account::default()
+        },
+    );
+
+    let previous_public_state = state.public_state.clone();
+    assert!(state.private_state.0.contains(&sender_pre_commitment));
+    assert!(!state.private_state.0.contains(&expected_new_commitment_1));
+    assert!(!state.private_state.0.contains(&expected_new_commitment_2));
+    assert!(!state.private_state.1.contains(&expected_new_nullifier));
+
+    state
+        .transition_from_privacy_preserving_transaction(&tx, 1, 0)
+        .unwrap();
+
+    assert_eq!(state.public_state, previous_public_state);
+    assert!(state.private_state.0.contains(&sender_pre_commitment));
+    assert!(state.private_state.0.contains(&expected_new_commitment_1));
+    assert!(state.private_state.0.contains(&expected_new_commitment_2));
+    assert!(state.private_state.1.contains(&expected_new_nullifier));
+}
+
+/// After a valid fully-private tx is proven, tampering with a note's epk should
+/// make the shielding proof invalid.
+#[test]
+fn privacy_tampered_epk_is_rejected() {
+    use crate::validated_state_diff::ValidatedStateDiff;
+
+    let (state, mut tx) = valid_private_transfer_tx_and_state();
+
+    // Baseline: the untampered tx verifies
+    assert!(
+        ValidatedStateDiff::from_privacy_preserving_transaction(&tx, &state, 1, 0).is_ok(),
+        "the unmodified private transfer must verify"
+    );
+
+    // Flip a byte of the first note's epk
+    tx.message.encrypted_private_post_states[0].epk.0[0] ^= 0xFF;
+
+    assert!(
+        matches!(
+            ValidatedStateDiff::from_privacy_preserving_transaction(&tx, &state, 1, 0),
+            Err(LeeError::InvalidPrivacyPreservingProof)
+        ),
+        "a tampered epk must be rejected by proof verification"
+    );
+}
+
+/// After a valid fully-private tx is proven, tampering with a note's view tag should
+/// make the shielding proof invalid.
+#[test]
+fn privacy_tampered_view_tag_is_rejected() {
+    use crate::validated_state_diff::ValidatedStateDiff;
+
+    let (state, mut tx) = valid_private_transfer_tx_and_state();
+
+    // Baseline: the untampered tx verifies.
+    assert!(
+        ValidatedStateDiff::from_privacy_preserving_transaction(&tx, &state, 1, 0).is_ok(),
+        "the unmodified private transfer must verify"
+    );
+
+    // Flip the first note's view_tag
+    tx.message.encrypted_private_post_states[0].view_tag ^= 0xFF;
+
+    assert!(
+        matches!(
+            ValidatedStateDiff::from_privacy_preserving_transaction(&tx, &state, 1, 0),
+            Err(LeeError::InvalidPrivacyPreservingProof)
+        ),
+        "a tampered view_tag must be rejected by proof verification"
+    );
+}
+
+#[test]
+fn transition_from_privacy_preserving_transaction_deshielded() {
+    let sender_keys = test_private_account_keys_1();
+    let sender_nonce = Nonce(0xdead_beef);
+
+    let sender_private_account = Account {
+        program_owner: crate::test_methods::simple_balance_transfer().id(),
+        balance: 100,
+        nonce: sender_nonce,
+        data: Data::default(),
+    };
+    let recipient_keys = test_public_account_keys_1();
+    let recipient_initial_balance = 400;
+    let mut state = V03State::new()
+        .with_public_accounts([(
+            recipient_keys.account_id(),
+            Account {
+                program_owner: crate::test_methods::simple_balance_transfer().id(),
+                balance: recipient_initial_balance,
+                ..Account::default()
+            },
+        )])
+        .with_private_account(&sender_keys, &sender_private_account);
+
+    let balance_to_move = 37;
+
+    let expected_recipient_post = {
+        let mut this = state.get_account_by_id(recipient_keys.account_id());
+        this.balance += balance_to_move;
+        this
+    };
+
+    let tx = deshielded_balance_transfer_for_tests(
+        &sender_keys,
+        &sender_private_account,
+        &recipient_keys.account_id(),
+        balance_to_move,
+        &state,
+    );
+
+    let sender_account_id = AccountId::for_regular_private_account(&sender_keys.npk(), 0);
+    let expected_new_commitment = Commitment::new(
+        &sender_account_id,
+        &Account {
+            program_owner: crate::test_methods::simple_balance_transfer().id(),
+            nonce: sender_nonce.private_account_nonce_increment(&sender_keys.nsk),
+            balance: sender_private_account.balance - balance_to_move,
+            data: Data::default(),
+        },
+    );
+
+    let sender_pre_commitment = Commitment::new(&sender_account_id, &sender_private_account);
+    let expected_new_nullifier =
+        Nullifier::for_account_update(&sender_pre_commitment, &sender_keys.nsk);
+
+    assert!(state.private_state.0.contains(&sender_pre_commitment));
+    assert!(!state.private_state.0.contains(&expected_new_commitment));
+    assert!(!state.private_state.1.contains(&expected_new_nullifier));
+
+    state
+        .transition_from_privacy_preserving_transaction(&tx, 1, 0)
+        .unwrap();
+
+    let recipient_post = state.get_account_by_id(recipient_keys.account_id());
+    assert_eq!(recipient_post, expected_recipient_post);
+    assert!(state.private_state.0.contains(&sender_pre_commitment));
+    assert!(state.private_state.0.contains(&expected_new_commitment));
+    assert!(state.private_state.1.contains(&expected_new_nullifier));
+    assert_eq!(
+        state.get_account_by_id(recipient_keys.account_id()).balance,
+        recipient_initial_balance + balance_to_move
+    );
+}
+
+#[test]
+fn burner_program_should_fail_in_privacy_preserving_circuit() {
+    let program = crate::test_methods::burner();
+    let public_account = AccountWithMetadata::new(
+        Account {
+            program_owner: program.id(),
+            balance: 100,
+            ..Account::default()
+        },
+        true,
+        AccountId::new([0; 32]),
+    );
+
+    let result = execute_and_prove(
+        vec![public_account],
+        Program::serialize_instruction(10_u128).unwrap(),
+        vec![InputAccountIdentity::Public],
+        &program.into(),
+    );
+
+    assert!(matches!(result, Err(LeeError::CircuitProvingError(_))));
+}
+
+#[test]
+fn minter_program_should_fail_in_privacy_preserving_circuit() {
+    let program = crate::test_methods::minter();
+    let public_account = AccountWithMetadata::new(
+        Account {
+            program_owner: program.id(),
+            balance: 0,
+            ..Account::default()
+        },
+        true,
+        AccountId::new([0; 32]),
+    );
+
+    let result = execute_and_prove(
+        vec![public_account],
+        Program::serialize_instruction(10_u128).unwrap(),
+        vec![InputAccountIdentity::Public],
+        &program.into(),
+    );
+
+    assert!(matches!(result, Err(LeeError::CircuitProvingError(_))));
+}
+
+#[test]
+fn nonce_changer_program_should_fail_in_privacy_preserving_circuit() {
+    let program = crate::test_methods::nonce_changer();
+    let public_account = AccountWithMetadata::new(
+        Account {
+            program_owner: program.id(),
+            balance: 0,
+            ..Account::default()
+        },
+        true,
+        AccountId::new([0; 32]),
+    );
+
+    let result = execute_and_prove(
+        vec![public_account],
+        Program::serialize_instruction(()).unwrap(),
+        vec![InputAccountIdentity::Public],
+        &program.into(),
+    );
+
+    assert!(matches!(result, Err(LeeError::CircuitProvingError(_))));
+}
+
+#[test]
+fn data_changer_program_should_fail_for_non_owned_account_in_privacy_preserving_circuit() {
+    let program = crate::test_methods::data_changer();
+    let public_account = AccountWithMetadata::new(
+        Account {
+            program_owner: [0, 1, 2, 3, 4, 5, 6, 7],
+            balance: 0,
+            ..Account::default()
+        },
+        true,
+        AccountId::new([0; 32]),
+    );
+
+    let result = execute_and_prove(
+        vec![public_account],
+        Program::serialize_instruction(vec![0]).unwrap(),
+        vec![InputAccountIdentity::Public],
+        &program.into(),
+    );
+
+    assert!(matches!(result, Err(LeeError::CircuitProvingError(_))));
+}
+
+#[test]
+fn data_changer_program_should_fail_for_too_large_data_in_privacy_preserving_circuit() {
+    let program = crate::test_methods::data_changer();
+    let public_account = AccountWithMetadata::new(
+        Account {
+            program_owner: program.id(),
+            balance: 0,
+            ..Account::default()
+        },
+        true,
+        AccountId::new([0; 32]),
+    );
+
+    let large_data: Vec<u8> =
+        vec![
+            0;
+            usize::try_from(lee_core::account::data::DATA_MAX_LENGTH.as_u64())
+                .expect("DATA_MAX_LENGTH fits in usize")
+                + 1
+        ];
+
+    let result = execute_and_prove(
+        vec![public_account],
+        Program::serialize_instruction(large_data).unwrap(),
+        vec![InputAccountIdentity::Public],
+        &program.into(),
+    );
+
+    assert!(matches!(result, Err(LeeError::ProgramProveFailed(_))));
+}
+
+#[test]
+fn extra_output_program_should_fail_in_privacy_preserving_circuit() {
+    let program = crate::test_methods::extra_output();
+    let public_account = AccountWithMetadata::new(
+        Account {
+            program_owner: program.id(),
+            balance: 0,
+            ..Account::default()
+        },
+        true,
+        AccountId::new([0; 32]),
+    );
+
+    let result = execute_and_prove(
+        vec![public_account],
+        Program::serialize_instruction(()).unwrap(),
+        vec![InputAccountIdentity::Public],
+        &program.into(),
+    );
+
+    assert!(matches!(result, Err(LeeError::CircuitProvingError(_))));
+}
+
+#[test]
+fn missing_output_program_should_fail_in_privacy_preserving_circuit() {
+    let program = crate::test_methods::missing_output();
+    let public_account_1 = AccountWithMetadata::new(
+        Account {
+            program_owner: program.id(),
+            balance: 0,
+            ..Account::default()
+        },
+        true,
+        AccountId::new([0; 32]),
+    );
+    let public_account_2 = AccountWithMetadata::new(
+        Account {
+            program_owner: program.id(),
+            balance: 0,
+            ..Account::default()
+        },
+        true,
+        AccountId::new([1; 32]),
+    );
+
+    let result = execute_and_prove(
+        vec![public_account_1, public_account_2],
+        Program::serialize_instruction(()).unwrap(),
+        vec![InputAccountIdentity::Public, InputAccountIdentity::Public],
+        &program.into(),
+    );
+
+    assert!(matches!(result, Err(LeeError::CircuitProvingError(_))));
+}
+
+#[test]
+fn program_owner_changer_should_fail_in_privacy_preserving_circuit() {
+    let program = crate::test_methods::program_owner_changer();
+    let public_account = AccountWithMetadata::new(
+        Account {
+            program_owner: program.id(),
+            balance: 0,
+            ..Account::default()
+        },
+        true,
+        AccountId::new([0; 32]),
+    );
+
+    let result = execute_and_prove(
+        vec![public_account],
+        Program::serialize_instruction(()).unwrap(),
+        vec![InputAccountIdentity::Public],
+        &program.into(),
+    );
+
+    assert!(matches!(result, Err(LeeError::CircuitProvingError(_))));
+}
+
+#[test]
+fn transfer_from_non_owned_account_should_fail_in_privacy_preserving_circuit() {
+    let program = crate::test_methods::simple_balance_transfer();
+    let public_account_1 = AccountWithMetadata::new(
+        Account {
+            program_owner: [0, 1, 2, 3, 4, 5, 6, 7],
+            balance: 100,
+            ..Account::default()
+        },
+        true,
+        AccountId::new([0; 32]),
+    );
+    let public_account_2 = AccountWithMetadata::new(
+        Account {
+            program_owner: program.id(),
+            balance: 0,
+            ..Account::default()
+        },
+        true,
+        AccountId::new([1; 32]),
+    );
+
+    let result = execute_and_prove(
+        vec![public_account_1, public_account_2],
+        Program::serialize_instruction(10_u128).unwrap(),
+        vec![InputAccountIdentity::Public, InputAccountIdentity::Public],
+        &program.into(),
+    );
+
+    assert!(matches!(result, Err(LeeError::CircuitProvingError(_))));
+}
+
+#[test]
+fn malicious_authorization_changer_should_fail_in_privacy_preserving_circuit() {
+    // Arrange
+    let malicious_program = crate::test_methods::malicious_authorization_changer();
+    let simple_transfers = crate::test_methods::simple_balance_transfer();
+    let sender_keys = test_public_account_keys_1();
+    let recipient_keys = test_private_account_keys_1();
+
+    let sender_account = AccountWithMetadata::new(
+        Account {
+            program_owner: simple_transfers.id(),
+            balance: 100,
+            ..Default::default()
+        },
+        false,
+        sender_keys.account_id(),
+    );
+    let recipient_account =
+        AccountWithMetadata::new(Account::default(), true, (&recipient_keys.npk(), 0));
+
+    let recipient_account_id = AccountId::for_regular_private_account(&recipient_keys.npk(), 0);
+    let recipient_commitment = Commitment::new(&recipient_account_id, &recipient_account.account);
+    let recipient_init_nullifier = Nullifier::for_account_initialization(&recipient_account_id);
+    let state = V03State::new()
+        .with_public_accounts(public_state_from_balances(&[(
+            sender_account.account_id,
+            sender_account.account.balance,
+        )]))
+        .with_private_accounts([(recipient_commitment.clone(), recipient_init_nullifier)])
+        .with_test_programs();
+
+    let balance_to_transfer = 10_u128;
+    let instruction = (balance_to_transfer, simple_transfers.id());
+
+    let recipient =
+        SharedSecretKey::encapsulate_deterministic(&recipient_keys.vpk(), &[0_u8; 32], 0).0;
+
+    let mut dependencies = HashMap::new();
+    dependencies.insert(simple_transfers.id(), simple_transfers);
+    let program_with_deps = ProgramWithDependencies::new(malicious_program, dependencies);
+
+    // Act - execute the malicious program - this should fail during proving
+    let result = execute_and_prove(
+        vec![sender_account, recipient_account],
+        Program::serialize_instruction(instruction).unwrap(),
+        vec![
+            InputAccountIdentity::Public,
+            InputAccountIdentity::PrivateAuthorizedUpdate {
+                epk: EphemeralPublicKey(Vec::new()),
+                view_tag: EncryptedAccountData::compute_view_tag(
+                    &recipient_keys.npk(),
+                    &recipient_keys.vpk(),
+                ),
+                ssk: recipient,
+                nsk: recipient_keys.nsk,
+                membership_proof: state
+                    .get_proof_for_commitment(&recipient_commitment)
+                    .expect("recipient's commitment must be in state"),
+                identifier: 0,
+            },
+        ],
+        &program_with_deps,
+    );
+
+    // Assert - should fail because the malicious program tries to manipulate is_authorized
+    assert!(matches!(result, Err(LeeError::CircuitProvingError(_))));
+}

lee/state_machine/src/state/tests/public_program_rules.rs

@@ -0,0 +1,325 @@
+use super::*;
+
+#[test]
+fn program_should_fail_if_modifies_nonces() {
+    let account_id = AccountId::new([1; 32]);
+    let mut state = V03State::new()
+        .with_public_account_balances([(account_id, 100)])
+        .with_test_programs();
+    let account_ids = vec![account_id];
+    let program_id = crate::test_methods::nonce_changer().id();
+    let message =
+        public_transaction::Message::try_new(program_id, account_ids, vec![], ()).unwrap();
+    let witness_set = public_transaction::WitnessSet::for_message(&message, &[]);
+    let tx = PublicTransaction::new(message, witness_set);
+
+    let result = state.transition_from_public_transaction(&tx, 1, 0);
+
+    assert!(matches!(
+        result,
+        Err(LeeError::InvalidProgramBehavior(
+            InvalidProgramBehaviorError::ExecutionValidationFailed(
+                ExecutionValidationError::ModifiedNonce { account_id: err_account_id }
+            )
+        )) if err_account_id == account_id
+    ));
+}
+
+#[test]
+fn program_should_fail_if_output_accounts_exceed_inputs() {
+    let mut state = V03State::new()
+        .with_public_account_balances([(AccountId::new([1; 32]), 0)])
+        .with_test_programs();
+    let account_ids = vec![AccountId::new([1; 32])];
+    let program_id = crate::test_methods::extra_output().id();
+    let message =
+        public_transaction::Message::try_new(program_id, account_ids, vec![], ()).unwrap();
+    let witness_set = public_transaction::WitnessSet::for_message(&message, &[]);
+    let tx = PublicTransaction::new(message, witness_set);
+
+    let result = state.transition_from_public_transaction(&tx, 1, 0);
+
+    assert!(matches!(
+        result,
+        Err(LeeError::InvalidProgramBehavior(
+            InvalidProgramBehaviorError::ExecutionValidationFailed(
+                ExecutionValidationError::MismatchedPreStatePostStateLength {
+                    pre_state_length,
+                    post_state_length
+                }
+            )
+        )) if pre_state_length == 1 && post_state_length == 2
+    ));
+}
+
+#[test]
+fn program_should_fail_with_missing_output_accounts() {
+    let mut state = V03State::new()
+        .with_public_account_balances([(AccountId::new([1; 32]), 100)])
+        .with_test_programs();
+    let account_ids = vec![AccountId::new([1; 32]), AccountId::new([2; 32])];
+    let program_id = crate::test_methods::missing_output().id();
+    let message =
+        public_transaction::Message::try_new(program_id, account_ids, vec![], ()).unwrap();
+    let witness_set = public_transaction::WitnessSet::for_message(&message, &[]);
+    let tx = PublicTransaction::new(message, witness_set);
+
+    let result = state.transition_from_public_transaction(&tx, 1, 0);
+
+    assert!(matches!(
+        result,
+        Err(LeeError::InvalidProgramBehavior(
+            InvalidProgramBehaviorError::ExecutionValidationFailed(
+                ExecutionValidationError::MismatchedPreStatePostStateLength {
+                    pre_state_length,
+                    post_state_length
+                }
+            )
+        )) if pre_state_length == 2 && post_state_length == 1
+    ));
+}
+
+#[test]
+fn program_should_fail_if_modifies_program_owner_with_only_non_default_program_owner() {
+    let initial_data = [(
+        AccountId::new([1; 32]),
+        Account {
+            program_owner: crate::test_methods::simple_balance_transfer().id(),
+            ..Account::default()
+        },
+    )];
+    let mut state = V03State::new()
+        .with_public_accounts(initial_data)
+        .with_test_programs();
+    let account_id = AccountId::new([1; 32]);
+    let account = state.get_account_by_id(account_id);
+    // Assert the target account only differs from the default account in the program owner
+    // field
+    assert_ne!(account.program_owner, Account::default().program_owner);
+    assert_eq!(account.balance, Account::default().balance);
+    assert_eq!(account.nonce, Account::default().nonce);
+    assert_eq!(account.data, Account::default().data);
+    let program_id = crate::test_methods::program_owner_changer().id();
+    let message =
+        public_transaction::Message::try_new(program_id, vec![account_id], vec![], ()).unwrap();
+    let witness_set = public_transaction::WitnessSet::for_message(&message, &[]);
+    let tx = PublicTransaction::new(message, witness_set);
+
+    let result = state.transition_from_public_transaction(&tx, 1, 0);
+
+    assert!(matches!(
+        result,
+        Err(LeeError::InvalidProgramBehavior(InvalidProgramBehaviorError::ExecutionValidationFailed(
+            ExecutionValidationError::ModifiedProgramOwner { account_id: err_account_id }
+        ))) if err_account_id == account_id
+    ));
+}
+
+#[test]
+fn program_should_fail_if_modifies_program_owner_with_only_non_default_balance() {
+    let initial_data = HashMap::new();
+    let mut state = V03State::new()
+        .with_public_accounts(initial_data)
+        .with_test_programs()
+        .with_non_default_accounts_but_default_program_owners();
+    let account_id = AccountId::new([255; 32]);
+    let account = state.get_account_by_id(account_id);
+    // Assert the target account only differs from the default account in balance field
+    assert_eq!(account.program_owner, Account::default().program_owner);
+    assert_ne!(account.balance, Account::default().balance);
+    assert_eq!(account.nonce, Account::default().nonce);
+    assert_eq!(account.data, Account::default().data);
+    let program_id = crate::test_methods::program_owner_changer().id();
+    let message =
+        public_transaction::Message::try_new(program_id, vec![account_id], vec![], ()).unwrap();
+    let witness_set = public_transaction::WitnessSet::for_message(&message, &[]);
+    let tx = PublicTransaction::new(message, witness_set);
+
+    let result = state.transition_from_public_transaction(&tx, 1, 0);
+
+    assert!(matches!(
+        result,
+        Err(LeeError::InvalidProgramBehavior(InvalidProgramBehaviorError::ExecutionValidationFailed(
+            ExecutionValidationError::ModifiedProgramOwner { account_id: err_account_id }
+        ))) if err_account_id == account_id
+    ));
+}
+
+#[test]
+fn program_should_fail_if_modifies_program_owner_with_only_non_default_nonce() {
+    let initial_data = HashMap::new();
+    let mut state = V03State::new()
+        .with_public_accounts(initial_data)
+        .with_test_programs()
+        .with_non_default_accounts_but_default_program_owners();
+    let account_id = AccountId::new([254; 32]);
+    let account = state.get_account_by_id(account_id);
+    // Assert the target account only differs from the default account in nonce field
+    assert_eq!(account.program_owner, Account::default().program_owner);
+    assert_eq!(account.balance, Account::default().balance);
+    assert_ne!(account.nonce, Account::default().nonce);
+    assert_eq!(account.data, Account::default().data);
+    let program_id = crate::test_methods::program_owner_changer().id();
+    let message =
+        public_transaction::Message::try_new(program_id, vec![account_id], vec![], ()).unwrap();
+    let witness_set = public_transaction::WitnessSet::for_message(&message, &[]);
+    let tx = PublicTransaction::new(message, witness_set);
+
+    let result = state.transition_from_public_transaction(&tx, 1, 0);
+
+    assert!(matches!(
+        result,
+        Err(LeeError::InvalidProgramBehavior(InvalidProgramBehaviorError::ExecutionValidationFailed(
+            ExecutionValidationError::ModifiedProgramOwner { account_id: err_account_id }
+        ))) if err_account_id == account_id
+    ));
+}
+
+#[test]
+fn program_should_fail_if_modifies_program_owner_with_only_non_default_data() {
+    let initial_data = HashMap::new();
+    let mut state = V03State::new()
+        .with_public_accounts(initial_data)
+        .with_test_programs()
+        .with_non_default_accounts_but_default_program_owners();
+    let account_id = AccountId::new([253; 32]);
+    let account = state.get_account_by_id(account_id);
+    // Assert the target account only differs from the default account in data field
+    assert_eq!(account.program_owner, Account::default().program_owner);
+    assert_eq!(account.balance, Account::default().balance);
+    assert_eq!(account.nonce, Account::default().nonce);
+    assert_ne!(account.data, Account::default().data);
+    let program_id = crate::test_methods::program_owner_changer().id();
+    let message =
+        public_transaction::Message::try_new(program_id, vec![account_id], vec![], ()).unwrap();
+    let witness_set = public_transaction::WitnessSet::for_message(&message, &[]);
+    let tx = PublicTransaction::new(message, witness_set);
+
+    let result = state.transition_from_public_transaction(&tx, 1, 0);
+
+    assert!(matches!(
+        result,
+        Err(LeeError::InvalidProgramBehavior(InvalidProgramBehaviorError::ExecutionValidationFailed(
+            ExecutionValidationError::ModifiedProgramOwner { account_id: err_account_id }
+        ))) if err_account_id == account_id
+    ));
+}
+
+#[test]
+fn program_should_fail_if_transfers_balance_from_non_owned_account() {
+    let sender_account_id = AccountId::new([1; 32]);
+    let receiver_account_id = AccountId::new([2; 32]);
+    let mut state = V03State::new()
+        .with_public_account_balances([(sender_account_id, 100)])
+        .with_test_programs();
+    let balance_to_move: u128 = 1;
+    let program_id = crate::test_methods::simple_balance_transfer().id();
+    assert_ne!(
+        state.get_account_by_id(sender_account_id).program_owner,
+        program_id
+    );
+    let message = public_transaction::Message::try_new(
+        program_id,
+        vec![sender_account_id, receiver_account_id],
+        vec![],
+        balance_to_move,
+    )
+    .unwrap();
+    let witness_set = public_transaction::WitnessSet::for_message(&message, &[]);
+    let tx = PublicTransaction::new(message, witness_set);
+
+    let result = state.transition_from_public_transaction(&tx, 1, 0);
+
+    assert!(matches!(
+        result,
+        Err(LeeError::InvalidProgramBehavior(InvalidProgramBehaviorError::ExecutionValidationFailed(
+            ExecutionValidationError::UnauthorizedBalanceDecrease { account_id: err_account_id, owner_program_id, executing_program_id }
+        ))) if err_account_id == sender_account_id && owner_program_id != program_id && executing_program_id == program_id
+    ));
+}
+
+#[test]
+fn program_should_fail_if_modifies_data_of_non_owned_account() {
+    let initial_data = HashMap::new();
+    let mut state = V03State::new()
+        .with_public_accounts(initial_data)
+        .with_test_programs()
+        .with_non_default_accounts_but_default_program_owners();
+    let account_id = AccountId::new([255; 32]);
+    let program_id = crate::test_methods::data_changer().id();
+
+    assert_ne!(state.get_account_by_id(account_id), Account::default());
+    assert_ne!(
+        state.get_account_by_id(account_id).program_owner,
+        program_id
+    );
+    let message =
+        public_transaction::Message::try_new(program_id, vec![account_id], vec![], vec![0])
+            .unwrap();
+    let witness_set = public_transaction::WitnessSet::for_message(&message, &[]);
+    let tx = PublicTransaction::new(message, witness_set);
+
+    let result = state.transition_from_public_transaction(&tx, 1, 0);
+
+    assert!(matches!(
+        result,
+        Err(LeeError::InvalidProgramBehavior(InvalidProgramBehaviorError::ExecutionValidationFailed(
+            ExecutionValidationError::UnauthorizedDataModification { account_id: err_account_id, executing_program_id }
+        ))) if err_account_id == account_id && executing_program_id == program_id
+    ));
+}
+
+#[test]
+fn program_should_fail_if_does_not_preserve_total_balance_by_minting() {
+    let initial_data = HashMap::new();
+    let mut state = V03State::new()
+        .with_public_accounts(initial_data)
+        .with_test_programs();
+    let account_id = AccountId::new([1; 32]);
+    let program_id = crate::test_methods::minter().id();
+
+    let message =
+        public_transaction::Message::try_new(program_id, vec![account_id], vec![], ()).unwrap();
+    let witness_set = public_transaction::WitnessSet::for_message(&message, &[]);
+    let tx = PublicTransaction::new(message, witness_set);
+
+    let result = state.transition_from_public_transaction(&tx, 2, 0);
+
+    assert!(matches!(
+        result,
+        Err(LeeError::InvalidProgramBehavior(InvalidProgramBehaviorError::ExecutionValidationFailed(
+            ExecutionValidationError::MismatchedTotalBalance { total_balance_pre_states, total_balance_post_states }
+        ))) if total_balance_pre_states == 0.into() && total_balance_post_states == 1.into()
+    ));
+}
+
+#[test]
+fn program_should_fail_if_does_not_preserve_total_balance_by_burning() {
+    let initial_data = HashMap::new();
+    let mut state = V03State::new()
+        .with_public_accounts(initial_data)
+        .with_test_programs()
+        .with_account_owned_by_burner_program();
+    let program_id = crate::test_methods::burner().id();
+    let account_id = AccountId::new([252; 32]);
+    assert_eq!(
+        state.get_account_by_id(account_id).program_owner,
+        program_id
+    );
+    let balance_to_burn: u128 = 1;
+    assert!(state.get_account_by_id(account_id).balance > balance_to_burn);
+
+    let message =
+        public_transaction::Message::try_new(program_id, vec![account_id], vec![], balance_to_burn)
+            .unwrap();
+    let witness_set = public_transaction::WitnessSet::for_message(&message, &[]);
+    let tx = PublicTransaction::new(message, witness_set);
+    let result = state.transition_from_public_transaction(&tx, 2, 0);
+
+    assert!(matches!(
+        result,
+        Err(LeeError::InvalidProgramBehavior(InvalidProgramBehaviorError::ExecutionValidationFailed(
+            ExecutionValidationError::MismatchedTotalBalance { total_balance_pre_states, total_balance_post_states }
+        ))) if total_balance_pre_states == 100.into() && total_balance_post_states == 99.into()
+    ));
+}

lee/state_machine/src/state/tests/validity_window.rs

@@ -0,0 +1,243 @@
+use super::*;
+
+#[test_case::test_case((Some(1), Some(3)), 3; "at upper bound")]
+#[test_case::test_case((Some(1), Some(3)), 2; "inside range")]
+#[test_case::test_case((Some(1), Some(3)), 0; "below range")]
+#[test_case::test_case((Some(1), Some(3)), 1; "at lower bound")]
+#[test_case::test_case((Some(1), Some(3)), 4; "above range")]
+#[test_case::test_case((Some(1), None), 1; "lower bound only - at bound")]
+#[test_case::test_case((Some(1), None), 10; "lower bound only - above")]
+#[test_case::test_case((Some(1), None), 0; "lower bound only - below")]
+#[test_case::test_case((None, Some(3)), 3; "upper bound only - at bound")]
+#[test_case::test_case((None, Some(3)), 0; "upper bound only - below")]
+#[test_case::test_case((None, Some(3)), 4; "upper bound only - above")]
+#[test_case::test_case((None, None), 0; "no bounds - always valid")]
+#[test_case::test_case((None, None), 100; "no bounds - always valid 2")]
+fn validity_window_works_in_public_transactions(
+    validity_window: (Option<BlockId>, Option<BlockId>),
+    block_id: BlockId,
+) {
+    let block_validity_window: BlockValidityWindow = validity_window.try_into().unwrap();
+    let validity_window_program = crate::test_methods::validity_window();
+    let account_keys = test_public_account_keys_1();
+    let pre = AccountWithMetadata::new(Account::default(), false, account_keys.account_id());
+    let mut state = V03State::new().with_test_programs();
+    let tx = {
+        let account_ids = vec![pre.account_id];
+        let nonces = vec![];
+        let program_id = validity_window_program.id();
+        let instruction = (
+            block_validity_window,
+            TimestampValidityWindow::new_unbounded(),
+        );
+        let message =
+            public_transaction::Message::try_new(program_id, account_ids, nonces, instruction)
+                .unwrap();
+        let witness_set = public_transaction::WitnessSet::for_message(&message, &[]);
+        PublicTransaction::new(message, witness_set)
+    };
+    let result = state.transition_from_public_transaction(&tx, block_id, 0);
+    let is_inside_validity_window =
+        match (block_validity_window.start(), block_validity_window.end()) {
+            (Some(s), Some(e)) => s <= block_id && block_id < e,
+            (Some(s), None) => s <= block_id,
+            (None, Some(e)) => block_id < e,
+            (None, None) => true,
+        };
+    if is_inside_validity_window {
+        assert!(result.is_ok());
+    } else {
+        assert!(matches!(result, Err(LeeError::OutOfValidityWindow)));
+    }
+}
+
+#[test_case::test_case((Some(1), Some(3)), 3; "at upper bound")]
+#[test_case::test_case((Some(1), Some(3)), 2; "inside range")]
+#[test_case::test_case((Some(1), Some(3)), 0; "below range")]
+#[test_case::test_case((Some(1), Some(3)), 1; "at lower bound")]
+#[test_case::test_case((Some(1), Some(3)), 4; "above range")]
+#[test_case::test_case((Some(1), None), 1; "lower bound only - at bound")]
+#[test_case::test_case((Some(1), None), 10; "lower bound only - above")]
+#[test_case::test_case((Some(1), None), 0; "lower bound only - below")]
+#[test_case::test_case((None, Some(3)), 3; "upper bound only - at bound")]
+#[test_case::test_case((None, Some(3)), 0; "upper bound only - below")]
+#[test_case::test_case((None, Some(3)), 4; "upper bound only - above")]
+#[test_case::test_case((None, None), 0; "no bounds - always valid")]
+#[test_case::test_case((None, None), 100; "no bounds - always valid 2")]
+fn timestamp_validity_window_works_in_public_transactions(
+    validity_window: (Option<Timestamp>, Option<Timestamp>),
+    timestamp: Timestamp,
+) {
+    let timestamp_validity_window: TimestampValidityWindow = validity_window.try_into().unwrap();
+    let validity_window_program = crate::test_methods::validity_window();
+    let account_keys = test_public_account_keys_1();
+    let pre = AccountWithMetadata::new(Account::default(), false, account_keys.account_id());
+    let mut state = V03State::new().with_test_programs();
+    let tx = {
+        let account_ids = vec![pre.account_id];
+        let nonces = vec![];
+        let program_id = validity_window_program.id();
+        let instruction = (
+            BlockValidityWindow::new_unbounded(),
+            timestamp_validity_window,
+        );
+        let message =
+            public_transaction::Message::try_new(program_id, account_ids, nonces, instruction)
+                .unwrap();
+        let witness_set = public_transaction::WitnessSet::for_message(&message, &[]);
+        PublicTransaction::new(message, witness_set)
+    };
+    let result = state.transition_from_public_transaction(&tx, 1, timestamp);
+    let is_inside_validity_window = match (
+        timestamp_validity_window.start(),
+        timestamp_validity_window.end(),
+    ) {
+        (Some(s), Some(e)) => s <= timestamp && timestamp < e,
+        (Some(s), None) => s <= timestamp,
+        (None, Some(e)) => timestamp < e,
+        (None, None) => true,
+    };
+    if is_inside_validity_window {
+        assert!(result.is_ok());
+    } else {
+        assert!(matches!(result, Err(LeeError::OutOfValidityWindow)));
+    }
+}
+
+#[test_case::test_case((Some(1), Some(3)), 3; "at upper bound")]
+#[test_case::test_case((Some(1), Some(3)), 2; "inside range")]
+#[test_case::test_case((Some(1), Some(3)), 0; "below range")]
+#[test_case::test_case((Some(1), Some(3)), 1; "at lower bound")]
+#[test_case::test_case((Some(1), Some(3)), 4; "above range")]
+#[test_case::test_case((Some(1), None), 1; "lower bound only - at bound")]
+#[test_case::test_case((Some(1), None), 10; "lower bound only - above")]
+#[test_case::test_case((Some(1), None), 0; "lower bound only - below")]
+#[test_case::test_case((None, Some(3)), 3; "upper bound only - at bound")]
+#[test_case::test_case((None, Some(3)), 0; "upper bound only - below")]
+#[test_case::test_case((None, Some(3)), 4; "upper bound only - above")]
+#[test_case::test_case((None, None), 0; "no bounds - always valid")]
+#[test_case::test_case((None, None), 100; "no bounds - always valid 2")]
+fn validity_window_works_in_privacy_preserving_transactions(
+    validity_window: (Option<BlockId>, Option<BlockId>),
+    block_id: BlockId,
+) {
+    let block_validity_window: BlockValidityWindow = validity_window.try_into().unwrap();
+    let validity_window_program = crate::test_methods::validity_window();
+    let account_keys = test_private_account_keys_1();
+    let pre = AccountWithMetadata::new(Account::default(), false, (&account_keys.npk(), 0));
+    let mut state = V03State::new().with_test_programs();
+    let tx = {
+        let (shared_secret, epk) =
+            SharedSecretKey::encapsulate_deterministic(&account_keys.vpk(), &[0_u8; 32], 0);
+
+        let instruction = (
+            block_validity_window,
+            TimestampValidityWindow::new_unbounded(),
+        );
+        let (output, proof) = crate::privacy_preserving_transaction::circuit::execute_and_prove(
+            vec![pre],
+            Program::serialize_instruction(instruction).unwrap(),
+            vec![InputAccountIdentity::PrivateUnauthorized {
+                epk,
+                view_tag: EncryptedAccountData::compute_view_tag(
+                    &account_keys.npk(),
+                    &account_keys.vpk(),
+                ),
+                npk: account_keys.npk(),
+                ssk: shared_secret,
+                identifier: 0,
+                commitment_root: DUMMY_COMMITMENT_HASH,
+            }],
+            &validity_window_program.into(),
+        )
+        .unwrap();
+
+        let message = Message::try_from_circuit_output(vec![], vec![], output).unwrap();
+
+        let witness_set = WitnessSet::for_message(&message, proof, &[]);
+        PrivacyPreservingTransaction::new(message, witness_set)
+    };
+    let result = state.transition_from_privacy_preserving_transaction(&tx, block_id, 0);
+    let is_inside_validity_window =
+        match (block_validity_window.start(), block_validity_window.end()) {
+            (Some(s), Some(e)) => s <= block_id && block_id < e,
+            (Some(s), None) => s <= block_id,
+            (None, Some(e)) => block_id < e,
+            (None, None) => true,
+        };
+    if is_inside_validity_window {
+        assert!(result.is_ok());
+    } else {
+        assert!(matches!(result, Err(LeeError::OutOfValidityWindow)));
+    }
+}
+
+#[test_case::test_case((Some(1), Some(3)), 3; "at upper bound")]
+#[test_case::test_case((Some(1), Some(3)), 2; "inside range")]
+#[test_case::test_case((Some(1), Some(3)), 0; "below range")]
+#[test_case::test_case((Some(1), Some(3)), 1; "at lower bound")]
+#[test_case::test_case((Some(1), Some(3)), 4; "above range")]
+#[test_case::test_case((Some(1), None), 1; "lower bound only - at bound")]
+#[test_case::test_case((Some(1), None), 10; "lower bound only - above")]
+#[test_case::test_case((Some(1), None), 0; "lower bound only - below")]
+#[test_case::test_case((None, Some(3)), 3; "upper bound only - at bound")]
+#[test_case::test_case((None, Some(3)), 0; "upper bound only - below")]
+#[test_case::test_case((None, Some(3)), 4; "upper bound only - above")]
+#[test_case::test_case((None, None), 0; "no bounds - always valid")]
+#[test_case::test_case((None, None), 100; "no bounds - always valid 2")]
+fn timestamp_validity_window_works_in_privacy_preserving_transactions(
+    validity_window: (Option<Timestamp>, Option<Timestamp>),
+    timestamp: Timestamp,
+) {
+    let timestamp_validity_window: TimestampValidityWindow = validity_window.try_into().unwrap();
+    let validity_window_program = crate::test_methods::validity_window();
+    let account_keys = test_private_account_keys_1();
+    let pre = AccountWithMetadata::new(Account::default(), false, (&account_keys.npk(), 0));
+    let mut state = V03State::new().with_test_programs();
+    let tx = {
+        let (shared_secret, epk) =
+            SharedSecretKey::encapsulate_deterministic(&account_keys.vpk(), &[0_u8; 32], 0);
+
+        let instruction = (
+            BlockValidityWindow::new_unbounded(),
+            timestamp_validity_window,
+        );
+        let (output, proof) = crate::privacy_preserving_transaction::circuit::execute_and_prove(
+            vec![pre],
+            Program::serialize_instruction(instruction).unwrap(),
+            vec![InputAccountIdentity::PrivateUnauthorized {
+                epk,
+                view_tag: EncryptedAccountData::compute_view_tag(
+                    &account_keys.npk(),
+                    &account_keys.vpk(),
+                ),
+                npk: account_keys.npk(),
+                ssk: shared_secret,
+                identifier: 0,
+                commitment_root: DUMMY_COMMITMENT_HASH,
+            }],
+            &validity_window_program.into(),
+        )
+        .unwrap();
+
+        let message = Message::try_from_circuit_output(vec![], vec![], output).unwrap();
+
+        let witness_set = WitnessSet::for_message(&message, proof, &[]);
+        PrivacyPreservingTransaction::new(message, witness_set)
+    };
+    let result = state.transition_from_privacy_preserving_transaction(&tx, 1, timestamp);
+    let is_inside_validity_window = match (
+        timestamp_validity_window.start(),
+        timestamp_validity_window.end(),
+    ) {
+        (Some(s), Some(e)) => s <= timestamp && timestamp < e,
+        (Some(s), None) => s <= timestamp,
+        (None, Some(e)) => timestamp < e,
+        (None, None) => true,
+    };
+    if is_inside_validity_window {
+        assert!(result.is_ok());
+    } else {
+        assert!(matches!(result, Err(LeeError::OutOfValidityWindow)));
+    }
+}

lee/state_machine/src/validated_state_diff/mod.rs

@@ -0,0 +1,521 @@
+use std::{
+    collections::{HashMap, HashSet, VecDeque},
+    hash::Hash,
+};
+
+use lee_core::{
+    BlockId, Commitment, Nullifier, PrivacyPreservingCircuitOutput, Timestamp,
+    account::{Account, AccountId, AccountWithMetadata},
+    program::{
+        ChainedCall, Claim, DEFAULT_PROGRAM_ID, ProgramId, compute_public_authorized_pdas,
+        validate_execution,
+    },
+};
+use log::debug;
+
+use crate::{
+    V03State, ensure,
+    error::{InvalidProgramBehaviorError, LeeError},
+    privacy_preserving_transaction::{
+        PrivacyPreservingTransaction, circuit::Proof, message::Message,
+    },
+    program::Program,
+    program_deployment_transaction::ProgramDeploymentTransaction,
+    public_transaction::PublicTransaction,
+    state::MAX_NUMBER_CHAINED_CALLS,
+};
+
+pub struct StateDiff {
+    pub signer_account_ids: Vec<AccountId>,
+    pub public_diff: HashMap<AccountId, Account>,
+    pub new_commitments: Vec<Commitment>,
+    pub new_nullifiers: Vec<Nullifier>,
+    pub program: Option<Program>,
+}
+
+/// The validated output of executing or verifying a transaction, ready to be applied to the state.
+///
+/// Can only be constructed by the transaction validation functions inside this crate, ensuring the
+/// diff has been checked before any state mutation occurs.
+pub struct ValidatedStateDiff(StateDiff);
+
+impl ValidatedStateDiff {
+    pub fn from_public_transaction(
+        tx: &PublicTransaction,
+        state: &V03State,
+        block_id: BlockId,
+        timestamp: Timestamp,
+    ) -> Result<Self, LeeError> {
+        let signer_account_ids = authenticate_public_transaction_signers(tx, state)?;
+        let message = tx.message();
+
+        ensure!(
+            !message.account_ids.is_empty(),
+            LeeError::InvalidInput("Public transaction must have at least one account".into())
+        );
+
+        // All account_ids must be different
+        ensure!(
+            message.account_ids.iter().collect::<HashSet<_>>().len() == message.account_ids.len(),
+            LeeError::InvalidInput("Duplicate account_ids found in message".into(),)
+        );
+
+        // Build pre_states for execution
+        let input_pre_states: Vec<_> = message
+            .account_ids
+            .iter()
+            .map(|account_id| {
+                AccountWithMetadata::new(
+                    state.get_account_by_id(*account_id),
+                    signer_account_ids.contains(account_id),
+                    *account_id,
+                )
+            })
+            .collect();
+
+        let mut state_diff: HashMap<AccountId, Account> = HashMap::new();
+
+        let initial_call = ChainedCall {
+            program_id: message.program_id,
+            instruction_data: message.instruction_data.clone(),
+            pre_states: input_pre_states,
+            pda_seeds: vec![],
+        };
+
+        let initial_caller_data = CallerData {
+            program_id: None,
+            authorized_accounts: signer_account_ids.iter().copied().collect(),
+        };
+
+        let mut chained_calls =
+            VecDeque::<(ChainedCall, CallerData)>::from_iter([(initial_call, initial_caller_data)]);
+        let mut chain_calls_counter = 0;
+
+        while let Some((chained_call, caller_data)) = chained_calls.pop_front() {
+            ensure!(
+                chain_calls_counter <= MAX_NUMBER_CHAINED_CALLS,
+                LeeError::MaxChainedCallsDepthExceeded
+            );
+
+            // Check that the `program_id` corresponds to a deployed program
+            let Some(program) = state.programs().get(&chained_call.program_id) else {
+                return Err(LeeError::InvalidInput("Unknown program".into()));
+            };
+
+            debug!(
+                "Program {:?} pre_states: {:?}, instruction_data: {:?}",
+                chained_call.program_id, chained_call.pre_states, chained_call.instruction_data
+            );
+            let mut program_output = program.execute(
+                caller_data.program_id,
+                &chained_call.pre_states,
+                &chained_call.instruction_data,
+            )?;
+            debug!(
+                "Program {:?} output: {:?}",
+                chained_call.program_id, program_output
+            );
+
+            let authorized_pdas =
+                compute_public_authorized_pdas(caller_data.program_id, &chained_call.pda_seeds);
+
+            // Account is authorized if it is either in the caller's authorized accounts or in the
+            // list of PDAs the caller has authorized.
+            let is_authorized = |account_id: &AccountId| {
+                authorized_pdas.contains(account_id)
+                    || caller_data.authorized_accounts.contains(account_id)
+            };
+
+            for pre in &program_output.pre_states {
+                let account_id = pre.account_id;
+                // Check that the program output pre_states coincide with the values in the public
+                // state or with any modifications to those values during the chain of calls.
+                let expected_pre = state_diff
+                    .get(&account_id)
+                    .cloned()
+                    .unwrap_or_else(|| state.get_account_by_id(account_id));
+                ensure!(
+                    pre.account == expected_pre,
+                    InvalidProgramBehaviorError::InconsistentAccountPreState {
+                        account_id,
+                        expected: Box::new(expected_pre),
+                        actual: Box::new(pre.account.clone())
+                    }
+                );
+
+                // Check that the program output pre_states marked as authorized are indeed
+                // authorized, and vice-versa.
+                let is_indeed_authorized = is_authorized(&account_id);
+                ensure!(
+                    !pre.is_authorized || is_indeed_authorized,
+                    InvalidProgramBehaviorError::InvalidAccountAuthorization { account_id }
+                );
+                ensure!(
+                    pre.is_authorized || !is_indeed_authorized,
+                    InvalidProgramBehaviorError::AuthorizedAccountMarkedAsNotAuthorized {
+                        account_id
+                    }
+                );
+            }
+
+            // Verify that the program output's self_program_id matches the expected program ID.
+            ensure!(
+                program_output.self_program_id == chained_call.program_id,
+                InvalidProgramBehaviorError::MismatchedProgramId {
+                    expected: chained_call.program_id,
+                    actual: program_output.self_program_id
+                }
+            );
+
+            // Verify that the program output's caller_program_id matches the actual caller.
+            ensure!(
+                program_output.caller_program_id == caller_data.program_id,
+                InvalidProgramBehaviorError::MismatchedCallerProgramId {
+                    expected: caller_data.program_id,
+                    actual: program_output.caller_program_id,
+                }
+            );
+
+            // Verify execution corresponds to a well-behaved program.
+            // See the # Programs section for the definition of the `validate_execution` method.
+            validate_execution(
+                &program_output.pre_states,
+                &program_output.post_states,
+                chained_call.program_id,
+            )
+            .map_err(InvalidProgramBehaviorError::ExecutionValidationFailed)?;
+
+            // Verify validity window
+            ensure!(
+                program_output.block_validity_window.is_valid_for(block_id)
+                    && program_output
+                        .timestamp_validity_window
+                        .is_valid_for(timestamp),
+                LeeError::OutOfValidityWindow
+            );
+
+            for (i, post) in program_output.post_states.iter_mut().enumerate() {
+                let Some(claim) = post.required_claim() else {
+                    continue;
+                };
+                let pre = &program_output.pre_states[i];
+                let account_id = pre.account_id;
+
+                // The invoked program can only claim accounts with default program id.
+                ensure!(
+                    post.account().program_owner == DEFAULT_PROGRAM_ID,
+                    InvalidProgramBehaviorError::ClaimedNonDefaultAccount { account_id }
+                );
+
+                match claim {
+                    Claim::Authorized => {
+                        // The program can only claim accounts that were authorized by the signer.
+                        ensure!(
+                            pre.is_authorized,
+                            InvalidProgramBehaviorError::ClaimedUnauthorizedAccount { account_id }
+                        );
+                    }
+                    Claim::Pda(seed) => {
+                        // The program can only claim accounts that correspond to the PDAs it is
+                        // authorized to claim. The public-execution path only sees public
+                        // accounts, so the public-PDA derivation is the correct formula here.
+                        let pda = AccountId::for_public_pda(&chained_call.program_id, &seed);
+                        ensure!(
+                            account_id == pda,
+                            InvalidProgramBehaviorError::MismatchedPdaClaim {
+                                expected: pda,
+                                actual: account_id
+                            }
+                        );
+                    }
+                }
+
+                post.account_mut().program_owner = chained_call.program_id;
+            }
+
+            // Update the state diff
+            for (pre, post) in program_output
+                .pre_states
+                .iter()
+                .zip(program_output.post_states.iter())
+            {
+                state_diff.insert(pre.account_id, post.account().clone());
+            }
+
+            // Source from `program_output.pre_states`, not `chained_call.pre_states`:
+            // the loop above already gates program_output's `is_authorized` via the
+            // `!pre.is_authorized || is_indeed_authorized` check, while `chained_call.
+            // pre_states` is caller-controlled and can be forged (audit-issue 91).
+            //
+            // Union with the caller's authorized set so that authorization is monotonically
+            // growing: once an account is authorized at any point in the chain it remains
+            // authorized for all subsequent calls.
+            let authorized_accounts: HashSet<_> = caller_data
+                .authorized_accounts
+                .into_iter()
+                .chain(
+                    program_output
+                        .pre_states
+                        .iter()
+                        .filter(|pre| pre.is_authorized)
+                        .map(|pre| pre.account_id),
+                )
+                .collect();
+            for new_call in program_output.chained_calls.into_iter().rev() {
+                chained_calls.push_front((
+                    new_call,
+                    CallerData {
+                        program_id: Some(chained_call.program_id),
+                        authorized_accounts: authorized_accounts.clone(),
+                    },
+                ));
+            }
+
+            chain_calls_counter = chain_calls_counter
+                .checked_add(1)
+                .expect("we check the max depth at the beginning of the loop");
+        }
+
+        // Check that all modified uninitialized accounts where claimed
+        for (account_id, post) in state_diff.iter().filter_map(|(account_id, post)| {
+            let pre = state.get_account_by_id(*account_id);
+            if pre.program_owner != DEFAULT_PROGRAM_ID {
+                return None;
+            }
+            if pre == *post {
+                return None;
+            }
+            Some((*account_id, post))
+        }) {
+            ensure!(
+                post.program_owner != DEFAULT_PROGRAM_ID,
+                InvalidProgramBehaviorError::DefaultAccountModifiedWithoutClaim { account_id }
+            );
+        }
+
+        Ok(Self(StateDiff {
+            signer_account_ids,
+            public_diff: state_diff,
+            new_commitments: vec![],
+            new_nullifiers: vec![],
+            program: None,
+        }))
+    }
+
+    pub fn from_privacy_preserving_transaction(
+        tx: &PrivacyPreservingTransaction,
+        state: &V03State,
+        block_id: BlockId,
+        timestamp: Timestamp,
+    ) -> Result<Self, LeeError> {
+        let message = &tx.message;
+        let witness_set = &tx.witness_set;
+
+        // 1. Commitments or nullifiers are non empty
+        ensure!(
+            !message.new_commitments.is_empty() || !message.new_nullifiers.is_empty(),
+            LeeError::InvalidInput(
+                "Empty commitments and empty nullifiers found in message".into(),
+            )
+        );
+
+        // 2. Check there are no duplicate account_ids in the public_account_ids list.
+        ensure!(
+            n_unique(&message.public_account_ids) == message.public_account_ids.len(),
+            LeeError::InvalidInput("Duplicate account_ids found in message".into())
+        );
+
+        // Check there are no duplicate nullifiers in the new_nullifiers list
+        ensure!(
+            n_unique(
+                &message
+                    .new_nullifiers
+                    .iter()
+                    .map(|(n, _)| n)
+                    .collect::<Vec<_>>()
+            ) == message.new_nullifiers.len(),
+            LeeError::InvalidInput("Duplicate nullifiers found in message".into())
+        );
+
+        // Check there are no duplicate commitments in the new_commitments list
+        ensure!(
+            n_unique(&message.new_commitments) == message.new_commitments.len(),
+            LeeError::InvalidInput("Duplicate commitments found in message".into())
+        );
+
+        // 3. Nonce checks and Valid signatures
+        // Check exactly one nonce is provided for each signature
+        ensure!(
+            message.nonces.len() == witness_set.signatures_and_public_keys.len(),
+            LeeError::InvalidInput(
+                "Mismatch between number of nonces and signatures/public keys".into(),
+            )
+        );
+
+        // Check the signatures are valid
+        ensure!(
+            witness_set.signatures_are_valid_for(message),
+            LeeError::InvalidInput("Invalid signature for given message and public key".into())
+        );
+
+        let signer_account_ids = tx.signer_account_ids();
+        // Check nonces corresponds to the current nonces on the public state.
+        for (account_id, nonce) in signer_account_ids.iter().zip(&message.nonces) {
+            let current_nonce = state.get_account_by_id(*account_id).nonce;
+            ensure!(
+                current_nonce == *nonce,
+                LeeError::InvalidInput("Nonce mismatch".into())
+            );
+        }
+
+        // Verify validity window
+        ensure!(
+            message.block_validity_window.is_valid_for(block_id)
+                && message.timestamp_validity_window.is_valid_for(timestamp),
+            LeeError::OutOfValidityWindow
+        );
+
+        // Build pre_states for proof verification
+        let public_pre_states: Vec<_> = message
+            .public_account_ids
+            .iter()
+            .map(|account_id| {
+                AccountWithMetadata::new(
+                    state.get_account_by_id(*account_id),
+                    signer_account_ids.contains(account_id),
+                    *account_id,
+                )
+            })
+            .collect();
+
+        // 4. Proof verification
+        check_privacy_preserving_circuit_proof_is_valid(
+            &witness_set.proof,
+            &public_pre_states,
+            message,
+        )?;
+
+        // 5. Commitment freshness
+        state.check_commitments_are_new(&message.new_commitments)?;
+
+        // 6. Nullifier uniqueness
+        state.check_nullifiers_are_valid(&message.new_nullifiers)?;
+
+        let public_diff = message
+            .public_account_ids
+            .iter()
+            .copied()
+            .zip(message.public_post_states.clone())
+            .collect();
+        let new_nullifiers = message
+            .new_nullifiers
+            .iter()
+            .copied()
+            .map(|(nullifier, _)| nullifier)
+            .collect();
+
+        Ok(Self(StateDiff {
+            signer_account_ids,
+            public_diff,
+            new_commitments: message.new_commitments.clone(),
+            new_nullifiers,
+            program: None,
+        }))
+    }
+
+    pub fn from_program_deployment_transaction(
+        tx: &ProgramDeploymentTransaction,
+        state: &V03State,
+    ) -> Result<Self, LeeError> {
+        // TODO: remove clone
+        let program = Program::new(tx.message.bytecode.clone().into())?;
+        if state.programs().contains_key(&program.id()) {
+            return Err(LeeError::ProgramAlreadyExists);
+        }
+        Ok(Self(StateDiff {
+            signer_account_ids: vec![],
+            public_diff: HashMap::new(),
+            new_commitments: vec![],
+            new_nullifiers: vec![],
+            program: Some(program),
+        }))
+    }
+
+    /// Returns the public account changes produced by this transaction.
+    ///
+    /// Used by callers (e.g. the sequencer) to inspect the diff before committing it, for example
+    /// to enforce that system accounts are not modified by user transactions.
+    #[must_use]
+    pub fn public_diff(&self) -> HashMap<AccountId, Account> {
+        self.0.public_diff.clone()
+    }
+
+    pub(crate) fn into_state_diff(self) -> StateDiff {
+        self.0
+    }
+}
+
+#[derive(Debug)]
+struct CallerData {
+    program_id: Option<ProgramId>,
+    authorized_accounts: HashSet<AccountId>,
+}
+
+fn authenticate_public_transaction_signers(
+    tx: &PublicTransaction,
+    state: &V03State,
+) -> Result<Vec<AccountId>, LeeError> {
+    let message = tx.message();
+    let witness_set = tx.witness_set();
+
+    ensure!(
+        message.nonces.len() == witness_set.signatures_and_public_keys.len(),
+        LeeError::InvalidInput(
+            "Mismatch between number of nonces and signatures/public keys".into(),
+        )
+    );
+
+    ensure!(
+        witness_set.is_valid_for(message),
+        LeeError::InvalidInput("Invalid signature for given message and public key".into())
+    );
+
+    let signer_account_ids = tx.signer_account_ids();
+    for (account_id, nonce) in signer_account_ids.iter().zip(&message.nonces) {
+        let current_nonce = state.get_account_by_id(*account_id).nonce;
+        ensure!(
+            current_nonce == *nonce,
+            LeeError::InvalidInput("Nonce mismatch".into())
+        );
+    }
+
+    Ok(signer_account_ids)
+}
+
+fn check_privacy_preserving_circuit_proof_is_valid(
+    proof: &Proof,
+    public_pre_states: &[AccountWithMetadata],
+    message: &Message,
+) -> Result<(), LeeError> {
+    let output = PrivacyPreservingCircuitOutput {
+        public_pre_states: public_pre_states.to_vec(),
+        public_post_states: message.public_post_states.clone(),
+        encrypted_private_post_states: message.encrypted_private_post_states.clone(),
+        new_commitments: message.new_commitments.clone(),
+        new_nullifiers: message.new_nullifiers.clone(),
+        block_validity_window: message.block_validity_window,
+        timestamp_validity_window: message.timestamp_validity_window,
+    };
+    proof
+        .is_valid_for(&output)
+        .then_some(())
+        .ok_or(LeeError::InvalidPrivacyPreservingProof)
+}
+
+fn n_unique<T: Eq + Hash>(data: &[T]) -> usize {
+    let set: HashSet<&T> = data.iter().collect();
+    set.len()
+}
+
+#[cfg(test)]
+mod tests;

lee/state_machine/src/validated_state_diff/tests.rs

@@ -0,0 +1,533 @@
+use std::collections::HashMap;
+
+use lee_core::account::{Account, AccountId, Nonce};
+
+use crate::{
+    PrivateKey, PublicKey, V03State,
+    error::{InvalidProgramBehaviorError, LeeError},
+    program::Program,
+    public_transaction::{Message, WitnessSet},
+    validated_state_diff::ValidatedStateDiff,
+};
+
+fn public_state_from_balances(initial_data: &[(AccountId, u128)]) -> HashMap<AccountId, Account> {
+    initial_data
+        .iter()
+        .copied()
+        .map(|(account_id, balance)| {
+            (
+                account_id,
+                Account {
+                    program_owner: crate::test_methods::simple_balance_transfer().id(),
+                    balance,
+                    ..Account::default()
+                },
+            )
+        })
+        .collect()
+}
+
+#[test]
+fn public_diff_reflects_a_successful_transfer() {
+    // A successful native transfer must record the debited sender in
+    // `public_diff()`.  Catches the mutation that replaces `public_diff` with
+    // `HashMap::new()` (which would hide every account change).
+    let from_key = PrivateKey::try_new([1_u8; 32]).unwrap();
+    let from = AccountId::from(&PublicKey::new_from_private_key(&from_key));
+    let to_key = PrivateKey::try_new([2_u8; 32]).unwrap();
+    let to = AccountId::from(&PublicKey::new_from_private_key(&to_key));
+
+    let state = V03State::new()
+        .with_public_accounts(public_state_from_balances(&[(from, 100)]))
+        .with_programs(std::iter::once(
+            crate::test_methods::simple_balance_transfer(),
+        ));
+    let program_id = crate::test_methods::simple_balance_transfer().id();
+    let message =
+        Message::try_new(program_id, vec![from, to], vec![Nonce(0), Nonce(0)], 5_u128).unwrap();
+    let witness_set = WitnessSet::for_message(&message, &[&from_key, &to_key]);
+    let tx = crate::PublicTransaction::new(message, witness_set);
+
+    let diff = ValidatedStateDiff::from_public_transaction(&tx, &state, 1, 0)
+        .expect("a valid native transfer must validate");
+    let public_diff = diff.public_diff();
+
+    assert!(
+        public_diff.contains_key(&from),
+        "public_diff must contain the debited sender",
+    );
+    assert_eq!(
+        public_diff[&from].balance, 95,
+        "sender balance in the diff must reflect the debit",
+    );
+}
+
+/// Privacy-path version of the authorization-injection attack. The test passes when the
+/// attack is rejected and the victim's balance is left untouched.
+///
+/// `execute_and_prove` succeeds because each inner receipt is individually valid and the
+/// outer circuit faithfully commits whatever the attacker's program output says, including
+/// `victim(is_authorized=true)`. The circuit has no access to chain state and cannot know
+/// the victim never signed.
+///
+/// The host-side validator is what catches the attack: it independently reconstructs
+/// `public_pre_states` from chain state using `signer_account_ids.contains(victim_id) = false`,
+/// so it expects `victim(is_authorized=false)`. The committed journal and the reconstructed
+/// expected output diverge, `receipt.verify` fails, and `from_privacy_preserving_transaction`
+/// returns an error before any state is applied.
+#[test]
+fn privacy_malicious_programs_cannot_drain_public_victim() {
+    use lee_core::{
+        Commitment, EncryptedAccountData, InputAccountIdentity, SharedSecretKey,
+        account::{Account, AccountWithMetadata},
+    };
+
+    use crate::{
+        PrivacyPreservingTransaction,
+        privacy_preserving_transaction::{
+            circuit::{ProgramWithDependencies, execute_and_prove},
+            message::Message,
+            witness_set::WitnessSet,
+        },
+        state::{CommitmentSet, tests::test_private_account_keys_1},
+    };
+
+    type InjectorInstruction = (
+        lee_core::program::ProgramId, // p2_id
+        lee_core::program::ProgramId, // simple_balance_transfer_id
+        [u8; 32],                     // victim_id_raw
+        u128,                         // victim_balance
+        u128,                         // victim_nonce
+        lee_core::program::ProgramId, // victim_program_owner
+        [u8; 32],                     // recipient_id_raw
+        u128,                         // amount
+    );
+
+    // Attacker controls a private account.
+    let attacker_keys = test_private_account_keys_1();
+    let attacker_id = AccountId::for_regular_private_account(&attacker_keys.npk(), 0);
+    let (attacker_ssk, attacker_epk) = SharedSecretKey::encapsulate(&attacker_keys.vpk());
+
+    let victim_id = AccountId::new([20_u8; 32]);
+    let recipient_id = AccountId::new([42_u8; 32]);
+    let victim_balance = 5_000_u128;
+
+    // genesis sets program_owner = simple_balance_transfer_program.id() on all accounts.
+    let state = V03State::new()
+        .with_public_accounts(public_state_from_balances(&[
+            (victim_id, victim_balance),
+            (recipient_id, 0),
+        ]))
+        .with_programs([
+            crate::test_methods::simple_balance_transfer(),
+            crate::test_methods::malicious_injector(),
+            crate::test_methods::malicious_launderer(),
+        ]);
+
+    // Build attacker's private account and its local commitment tree.
+    let attacker_account = Account {
+        program_owner: crate::test_methods::simple_balance_transfer().id(),
+        balance: 100,
+        ..Account::default()
+    };
+    let attacker_commitment = Commitment::new(&attacker_id, &attacker_account);
+    let mut commitment_set = CommitmentSet::with_capacity(1);
+    commitment_set.extend(std::slice::from_ref(&attacker_commitment));
+    let membership_proof = commitment_set
+        .get_proof_for(&attacker_commitment)
+        .expect("attacker commitment must be in the set");
+
+    let attacker_pre = AccountWithMetadata::new(attacker_account, true, attacker_id);
+
+    let victim_account = state.get_account_by_id(victim_id);
+    let instruction: InjectorInstruction = (
+        crate::test_methods::malicious_launderer().id(),
+        crate::test_methods::simple_balance_transfer().id(),
+        *victim_id.value(),
+        victim_account.balance,
+        victim_account.nonce.0,
+        victim_account.program_owner,
+        *recipient_id.value(),
+        victim_balance,
+    );
+    let instruction_data = Program::serialize_instruction(instruction).unwrap();
+
+    let p2 = crate::test_methods::malicious_launderer();
+    let at = crate::test_methods::simple_balance_transfer();
+    let program_with_deps = ProgramWithDependencies::new(
+        crate::test_methods::malicious_injector(),
+        [(p2.id(), p2), (at.id(), at)].into(),
+    );
+
+    // account_identities order must match self.pre_states as built by the circuit:
+    //   [0] attacker — first seen in P1's program_output.pre_states
+    //   [1] victim   — first seen in simple_balance_transfer's program_output.pre_states
+    //   [2] recipient — first seen in simple_balance_transfer's program_output.pre_states
+    let account_identities = vec![
+        InputAccountIdentity::PrivateAuthorizedUpdate {
+            epk: attacker_epk,
+            view_tag: EncryptedAccountData::compute_view_tag(
+                &attacker_keys.npk(),
+                &attacker_keys.vpk(),
+            ),
+            ssk: attacker_ssk,
+            nsk: attacker_keys.nsk,
+            membership_proof,
+            identifier: 0,
+        },
+        InputAccountIdentity::Public, // victim
+        InputAccountIdentity::Public, // recipient
+    ];
+
+    // execute_and_prove succeeds: all inner receipts are valid.
+    // The outer circuit commits victim(is_authorized=true) to its journal.
+    let (circuit_output, proof) = execute_and_prove(
+        vec![attacker_pre],
+        instruction_data,
+        account_identities,
+        &program_with_deps,
+    )
+    .expect("execute_and_prove should succeed \u{2014} the programs execute correctly");
+
+    // public_account_ids lists the Public entries from account_identities, in order.
+    // The single ciphertext belongs to attacker's private account update.
+    let message = Message::try_from_circuit_output(
+        vec![victim_id, recipient_id],
+        vec![], // no public signers, no nonces
+        circuit_output,
+    )
+    .unwrap();
+
+    let witness_set = WitnessSet::for_message(&message, proof, &[]); // no signatures
+    let tx = PrivacyPreservingTransaction::new(message, witness_set);
+
+    let result = ValidatedStateDiff::from_privacy_preserving_transaction(&tx, &state, 1, 0);
+
+    assert!(
+        matches!(result, Err(LeeError::InvalidPrivacyPreservingProof)),
+        "attack privacy transaction should be rejected with InvalidPrivacyPreservingProof"
+    );
+    assert_eq!(state.get_account_by_id(victim_id).balance, victim_balance);
+    assert_eq!(state.get_account_by_id(recipient_id).balance, 0);
+}
+
+/// Private-victim variant of the authorization-injection attack. The test passes when the
+/// attack is rejected and the recipient's balance remains zero.
+///
+/// After the circuit's Vacant branch accepts the injected `victim(is_authorized=true)`
+/// verbatim, the attacker must choose how to declare the victim in `account_identities`.
+/// There are two routes, both closed:
+///
+/// - **mask=1 (`PrivateAuthorizedUpdate`)**: the circuit derives `account_id =
+///   AccountId::for_regular_private_account(&npk_from(nsk), identifier)` and asserts it matches
+///   `pre_state.account_id`. Passing this check requires the victim's `nsk`, which the attacker
+///   does not have. `execute_and_prove` panics inside the ZKVM and no proof is produced.
+///
+/// - **mask=0 (`Public`)**: the circuit places the account in `public_pre_states` and
+///   `execute_and_prove` succeeds. The host-side validator then reconstructs `public_pre_states`
+///   from chain state; `state.get_account_by_id(victim_id)` returns the default account (balance=0)
+///   because the victim has no public state entry. The committed journal and the reconstructed
+///   expected output diverge, `receipt.verify` fails, and `from_privacy_preserving_transaction`
+///   returns an error before any state is applied. This test exercises this route.
+#[test]
+fn privacy_malicious_programs_cannot_drain_private_victim() {
+    use lee_core::{
+        Commitment, EncryptedAccountData, InputAccountIdentity, SharedSecretKey,
+        account::{Account, AccountWithMetadata},
+    };
+
+    use crate::{
+        PrivacyPreservingTransaction,
+        privacy_preserving_transaction::{
+            circuit::{ProgramWithDependencies, execute_and_prove},
+            message::Message,
+            witness_set::WitnessSet,
+        },
+        state::{
+            CommitmentSet,
+            tests::{test_private_account_keys_1, test_private_account_keys_2},
+        },
+    };
+
+    type InjectorInstruction = (
+        lee_core::program::ProgramId, // p2_id
+        lee_core::program::ProgramId, // simple_balance_transfer_id
+        [u8; 32],                     // victim_id_raw
+        u128,                         // victim_balance
+        u128,                         // victim_nonce
+        lee_core::program::ProgramId, // victim_program_owner
+        [u8; 32],                     // recipient_id_raw
+        u128,                         // amount
+    );
+
+    // Attacker controls a private account.
+    let attacker_keys = test_private_account_keys_1();
+    let attacker_id = AccountId::for_regular_private_account(&attacker_keys.npk(), 0);
+    let (attacker_ssk, attacker_epk) = SharedSecretKey::encapsulate(&attacker_keys.vpk());
+
+    // Victim is a private account — not registered in public chain state.
+    let victim_keys = test_private_account_keys_2();
+    let victim_id = AccountId::for_regular_private_account(&victim_keys.npk(), 0);
+    let victim_balance = 5_000_u128;
+
+    let recipient_id = AccountId::new([42_u8; 32]);
+
+    // Victim has no public state entry; only recipient is registered at genesis.
+    let state = V03State::new()
+        .with_public_accounts(public_state_from_balances(&[(recipient_id, 0)]))
+        .with_programs([
+            crate::test_methods::simple_balance_transfer(),
+            crate::test_methods::malicious_injector(),
+            crate::test_methods::malicious_launderer(),
+        ]);
+
+    // Build attacker's private account and its local commitment tree.
+    let attacker_account = Account {
+        program_owner: crate::test_methods::simple_balance_transfer().id(),
+        balance: 100,
+        ..Account::default()
+    };
+    let attacker_commitment = Commitment::new(&attacker_id, &attacker_account);
+    let mut commitment_set = CommitmentSet::with_capacity(1);
+    commitment_set.extend(std::slice::from_ref(&attacker_commitment));
+    let membership_proof = commitment_set
+        .get_proof_for(&attacker_commitment)
+        .expect("attacker commitment must be in the set");
+
+    let attacker_pre = AccountWithMetadata::new(attacker_account, true, attacker_id);
+
+    // The attacker supplies the victim's account data directly — it cannot be read from
+    // public state. The injected balance and program_owner allow simple_balance_transfer
+    // to succeed inside the circuit, which has no access to chain state and cannot detect
+    // that these values are fabricated.
+    let instruction: InjectorInstruction = (
+        crate::test_methods::malicious_launderer().id(),
+        crate::test_methods::simple_balance_transfer().id(),
+        *victim_id.value(),
+        victim_balance,
+        0_u128,                                              // nonce
+        crate::test_methods::simple_balance_transfer().id(), // program_owner
+        *recipient_id.value(),
+        victim_balance,
+    );
+    let instruction_data = Program::serialize_instruction(instruction).unwrap();
+
+    let p2 = crate::test_methods::malicious_launderer();
+    let at = crate::test_methods::simple_balance_transfer();
+    let program_with_deps = ProgramWithDependencies::new(
+        crate::test_methods::malicious_injector(),
+        [(p2.id(), p2), (at.id(), at)].into(),
+    );
+
+    // account_identities order must match self.pre_states as built by the circuit:
+    //   [0] attacker  — first seen in P1's program_output.pre_states
+    //   [1] victim    — first seen in simple_balance_transfer's program_output.pre_states
+    //   [2] recipient — first seen in simple_balance_transfer's program_output.pre_states
+    //
+    // Victim is marked Public: the attacker has no nsk for the victim's private account,
+    // so PrivateAuthorizedUpdate is not an option.
+    let account_identities = vec![
+        InputAccountIdentity::PrivateAuthorizedUpdate {
+            epk: attacker_epk,
+            view_tag: EncryptedAccountData::compute_view_tag(
+                &attacker_keys.npk(),
+                &attacker_keys.vpk(),
+            ),
+            ssk: attacker_ssk,
+            nsk: attacker_keys.nsk,
+            membership_proof,
+            identifier: 0,
+        },
+        InputAccountIdentity::Public, // victim — attacker lacks victim's nsk
+        InputAccountIdentity::Public, // recipient
+    ];
+
+    // execute_and_prove succeeds: simple_balance_transfer runs against the injected
+    // victim(balance=5000, is_authorized=true) and produces valid inner receipts.
+    // The outer circuit commits victim(is_authorized=true) to public_pre_states.
+    let (circuit_output, proof) = execute_and_prove(
+        vec![attacker_pre],
+        instruction_data,
+        account_identities,
+        &program_with_deps,
+    )
+    .expect("execute_and_prove should succeed \u{2014} the programs execute correctly");
+
+    // public_account_ids lists the Public entries from account_identities, in order.
+    // The single ciphertext belongs to attacker's private account update.
+    let message = Message::try_from_circuit_output(
+        vec![victim_id, recipient_id],
+        vec![], // no public signers, no nonces
+        circuit_output,
+    )
+    .unwrap();
+
+    let witness_set = WitnessSet::for_message(&message, proof, &[]); // no signatures
+    let tx = PrivacyPreservingTransaction::new(message, witness_set);
+
+    let result = ValidatedStateDiff::from_privacy_preserving_transaction(&tx, &state, 1, 0);
+
+    assert!(
+        matches!(result, Err(LeeError::InvalidPrivacyPreservingProof)),
+        "attack on private victim should be rejected with InvalidPrivacyPreservingProof"
+    );
+    // Victim has no public balance to check; confirming the recipient received nothing
+    // is sufficient to show no funds moved.
+    assert_eq!(state.get_account_by_id(recipient_id).balance, 0);
+}
+
+/// Two malicious programs (injector + launderer) attempt to drain a victim's balance
+/// without the victim signing anything. The test passes when the attack is rejected
+/// and the victim's balance is left untouched.
+///
+/// Attack flow:
+///   Transaction (attacker signs) → P1 (`malicious_injector`)
+///     → injects `victim(is_authorized=true)` into chained-call `pre_states` for P2
+///   P2 (`malicious_launderer`)
+///     → outputs empty pre/post states, forwarding the forged flag to `simple_balance_transfer`
+///     → if `authorized_accounts` were built from the injected `pre_states`,
+///       `{victim}.contains(victim)` would pass and the transfer would execute.
+///
+/// The validator must reject this: `authorized_accounts` must be derived from the
+/// parent program's own validated `program_output.pre_states`, not from the chained-call
+/// input, so a forged `is_authorized=true` flag is never trusted.
+#[test]
+fn malicious_programs_cannot_drain_victim_without_signature() {
+    // p2_id, simple_balance_transfer_id, victim_id_raw, victim_balance, victim_nonce,
+    // victim_program_owner, recipient_id_raw, amount.
+    // Primitives only — AccountId/Account cannot round-trip through instruction_data
+    // via risc0_zkvm::serde (SerializeDisplay issue).
+    type InjectorInstruction = (
+        lee_core::program::ProgramId, // p2_id
+        lee_core::program::ProgramId, // simple_balance_transfer_id
+        [u8; 32],                     // victim_id_raw
+        u128,                         // victim_balance
+        u128,                         // victim_nonce
+        lee_core::program::ProgramId, // victim_program_owner
+        [u8; 32],                     // recipient_id_raw
+        u128,                         // amount
+    );
+
+    let attacker_key = PrivateKey::try_new([10; 32]).unwrap();
+    let attacker_id = AccountId::from(&PublicKey::new_from_private_key(&attacker_key));
+
+    let victim_key = PrivateKey::try_new([20; 32]).unwrap();
+    let victim_id = AccountId::from(&PublicKey::new_from_private_key(&victim_key));
+
+    let recipient_id = AccountId::new([42; 32]);
+
+    let victim_balance = 5_000_u128;
+    let state = V03State::new()
+        .with_public_accounts(public_state_from_balances(&[
+            (attacker_id, 100),
+            (victim_id, victim_balance),
+            (recipient_id, 0),
+        ]))
+        .with_programs([
+            crate::test_methods::simple_balance_transfer(),
+            crate::test_methods::malicious_injector(),
+            crate::test_methods::malicious_launderer(),
+        ]);
+
+    // Read victim state from chain, exactly as the attacker would.
+    let victim_account = state.get_account_by_id(victim_id);
+
+    let instruction: InjectorInstruction = (
+        crate::test_methods::malicious_launderer().id(),
+        crate::test_methods::simple_balance_transfer().id(),
+        *victim_id.value(),
+        victim_account.balance,
+        victim_account.nonce.0,
+        victim_account.program_owner,
+        *recipient_id.value(),
+        victim_balance,
+    );
+
+    let message = Message::try_new(
+        crate::test_methods::malicious_injector().id(),
+        vec![attacker_id],
+        vec![Nonce(0)],
+        instruction,
+    )
+    .unwrap();
+
+    let witness_set = WitnessSet::for_message(&message, &[&attacker_key]);
+    let tx = crate::PublicTransaction::new(message, witness_set);
+
+    let result = ValidatedStateDiff::from_public_transaction(&tx, &state, 1, 0);
+
+    assert!(
+        matches!(
+            result,
+            Err(LeeError::InvalidProgramBehavior(
+                InvalidProgramBehaviorError::InvalidAccountAuthorization { account_id }
+            )) if account_id == victim_id
+        ),
+        "attack transaction should be rejected with InvalidAccountAuthorization for the victim"
+    );
+
+    // Confirm the victim's balance is untouched.
+    let victim_balance_after = state.get_account_by_id(victim_id).balance;
+    let recipient_balance_after = state.get_account_by_id(recipient_id).balance;
+
+    assert_eq!(
+        victim_balance_after, victim_balance,
+        "victim balance should be unchanged"
+    );
+    assert_eq!(
+        recipient_balance_after, 0,
+        "recipient should receive nothing"
+    );
+}
+
+/// Regression test: a `PrivacyPreservingTransaction` carrying a structurally invalid
+/// proof must be rejected with a clean `Err`.
+#[test]
+fn privacy_garbage_proof_is_rejected() {
+    use lee_core::{
+        Commitment,
+        account::Account,
+        program::{BlockValidityWindow, TimestampValidityWindow},
+    };
+
+    use crate::{
+        PrivacyPreservingTransaction,
+        privacy_preserving_transaction::{
+            circuit::Proof, message::Message, witness_set::WitnessSet,
+        },
+    };
+
+    let state = V03State::new();
+
+    // Minimal message that passes every check up to proof verification: a single
+    // commitment satisfies the non-empty requirement, no signers makes the
+    // nonce/signature checks vacuously true, and unbounded validity windows are valid
+    // for any block/timestamp.
+    let account_id = AccountId::from(&PublicKey::new_from_private_key(
+        &PrivateKey::try_new([1_u8; 32]).unwrap(),
+    ));
+    let commitment = Commitment::new(&account_id, &Account::default());
+    let message = Message {
+        public_account_ids: vec![],
+        nonces: vec![],
+        public_post_states: vec![],
+        encrypted_private_post_states: vec![],
+        new_commitments: vec![commitment],
+        new_nullifiers: vec![],
+        block_validity_window: BlockValidityWindow::new_unbounded(),
+        timestamp_validity_window: TimestampValidityWindow::new_unbounded(),
+    };
+
+    // Garbage proof bytes: not a valid borsh-encoded `InnerReceipt`.
+    let garbage_proof = Proof::from_inner(vec![0xff_u8; 64]);
+    let witness_set = WitnessSet::for_message(&message, garbage_proof, &[]);
+    let tx = PrivacyPreservingTransaction::new(message, witness_set);
+
+    let result = ValidatedStateDiff::from_privacy_preserving_transaction(&tx, &state, 1, 0);
+
+    match result {
+        Err(LeeError::InvalidPrivacyPreservingProof) => {}
+        Err(other) => panic!("expected InvalidPrivacyPreservingProof, got {other:?}"),
+        Ok(_) => panic!("garbage proof was accepted instead of rejected"),
+    }
+}

lez/common/src/transaction.rs

@@ -173,7 +173,7 @@ impl LeeTransaction {
                 balance: pre.balance,
                 ..post.clone()
             };
-            (expected_pre == pre) && (pre.balance <= post.balance)
+            (expected_pre == pre) && (pre.balance < post.balance)
         };
 
         if only_balance_increased {

lez/explorer_service/src/components/account_nonce_list.rs

@@ -0,0 +1,58 @@
+use indexer_service_protocol::AccountId;
+use itertools::{EitherOrBoth, Itertools as _};
+use leptos::prelude::*;
+use leptos_router::components::A;
+
+#[component]
+pub fn AccountNonceList(account_ids: Vec<AccountId>, nonces: Vec<u128>) -> impl IntoView {
+    view! {
+        <div class="accounts-list">
+            {account_ids
+                .into_iter()
+                .zip_longest(nonces.into_iter())
+                .map(|maybe_pair| {
+                    match maybe_pair {
+                        EitherOrBoth::Both(account_id, nonce) => {
+                            let account_id_str = account_id.to_string();
+                            view! {
+                                <div class="account-item">
+                                    <A href=format!("/account/{}", account_id_str)>
+                                        <span class="hash">{account_id_str}</span>
+                                    </A>
+                                    <span class="nonce">
+                                        " (nonce: " {nonce.to_string()} ")"
+                                    </span>
+                                </div>
+                            }
+                        }
+                        EitherOrBoth::Left(account_id) => {
+                            let account_id_str = account_id.to_string();
+                            view! {
+                                <div class="account-item">
+                                    <A href=format!("/account/{}", account_id_str)>
+                                        <span class="hash">{account_id_str}</span>
+                                    </A>
+                                    <span class="nonce">
+                                        " (nonce: "{"Not affected by this transaction".to_owned()}" )"
+                                    </span>
+                                </div>
+                            }
+                        }
+                        EitherOrBoth::Right(_) => {
+                            view! {
+                                <div class="account-item">
+                                    <A href=format!("/account/{}", "Account not found")>
+                                        <span class="hash">{"Account not found"}</span>
+                                    </A>
+                                    <span class="nonce">
+                                        " (nonce: "{"Account not found".to_owned()}" )"
+                                    </span>
+                                </div>
+                            }
+                        }
+                    }
+                })
+                .collect::<Vec<_>>()}
+        </div>
+    }
+}

lez/explorer_service/src/components/mod.rs

@@ -1,7 +1,15 @@
+pub use account_nonce_list::AccountNonceList;
 pub use account_preview::AccountPreview;
 pub use block_preview::BlockPreview;
+pub use search_results::SearchResultsView;
+pub use transaction_details::{
+    PrivacyPreservingTxDetails, ProgramDeploymentTxDetails, PublicTxDetails,
+};
 pub use transaction_preview::TransactionPreview;
 
+pub mod account_nonce_list;
 pub mod account_preview;
 pub mod block_preview;
+pub mod search_results;
+pub mod transaction_details;
 pub mod transaction_preview;

lez/explorer_service/src/components/search_results.rs

@@ -0,0 +1,93 @@
+use leptos::prelude::*;
+
+use super::{AccountPreview, BlockPreview, TransactionPreview};
+use crate::api::SearchResults;
+
+/// Search results view component
+#[component]
+pub fn SearchResultsView(results: SearchResults) -> impl IntoView {
+    let SearchResults {
+        blocks,
+        transactions,
+        accounts,
+    } = results;
+    let has_results = !blocks.is_empty() || !transactions.is_empty() || !accounts.is_empty();
+
+    view! {
+        <div class="search-results">
+            <h2>"Search Results"</h2>
+            {if has_results {
+                view! {
+                    <div class="results-container">
+                        {if blocks.is_empty() {
+                            ().into_any()
+                        } else {
+                            view! {
+                                <div class="results-section">
+                                    <h3>"Blocks"</h3>
+                                    <div class="results-list">
+                                        {blocks
+                                            .into_iter()
+                                            .map(|block| {
+                                                view! { <BlockPreview block=block /> }
+                                            })
+                                            .collect::<Vec<_>>()}
+                                    </div>
+                                </div>
+                            }
+                                .into_any()
+                        }}
+
+                        {if transactions.is_empty() {
+                            ().into_any()
+                        } else {
+                            view! {
+                                <div class="results-section">
+                                    <h3>"Transactions"</h3>
+                                    <div class="results-list">
+                                        {transactions
+                                            .into_iter()
+                                            .map(|tx| {
+                                                view! { <TransactionPreview transaction=tx /> }
+                                            })
+                                            .collect::<Vec<_>>()}
+                                    </div>
+                                </div>
+                            }
+                                .into_any()
+                        }}
+
+                        {if accounts.is_empty() {
+                            ().into_any()
+                        } else {
+                            view! {
+                                <div class="results-section">
+                                    <h3>"Accounts"</h3>
+                                    <div class="results-list">
+                                        {accounts
+                                            .into_iter()
+                                            .map(|(id, account)| {
+                                                view! {
+                                                    <AccountPreview
+                                                        account_id=id
+                                                        account=account
+                                                    />
+                                                }
+                                            })
+                                            .collect::<Vec<_>>()}
+                                    </div>
+                                </div>
+                            }
+                                .into_any()
+                        }}
+
+                    </div>
+                }
+                    .into_any()
+            } else {
+                    view! { <div class="not-found">"No results found"</div> }
+                    .into_any()
+            }}
+        </div>
+    }
+}

lez/explorer_service/src/components/transaction_details.rs

@@ -0,0 +1,150 @@
+use indexer_service_protocol::{
+    PrivacyPreservingMessage, PrivacyPreservingTransaction, ProgramDeploymentMessage,
+    ProgramDeploymentTransaction, PublicMessage, PublicTransaction, WitnessSet,
+};
+use leptos::prelude::*;
+
+use super::AccountNonceList;
+
+/// Public transaction details component
+#[component]
+pub fn PublicTxDetails(tx: PublicTransaction) -> impl IntoView {
+    let PublicTransaction {
+        hash: _,
+        message,
+        witness_set,
+    } = tx;
+    let PublicMessage {
+        program_id,
+        account_ids,
+        nonces,
+        instruction_data,
+    } = message;
+    let WitnessSet {
+        signatures_and_public_keys,
+        proof,
+    } = witness_set;
+
+    let program_id_str = program_id.to_string();
+    let proof_len = proof.map_or(0, |p| p.0.len());
+    let signatures_count = signatures_and_public_keys.len();
+
+    view! {
+        <div class="transaction-details">
+            <h2>"Public Transaction Details"</h2>
+            <div class="info-grid">
+                <div class="info-row">
+                    <span class="info-label">"Program ID:"</span>
+                    <span class="info-value hash">{program_id_str}</span>
+                </div>
+                <div class="info-row">
+                    <span class="info-label">"Instruction Data:"</span>
+                    <span class="info-value">
+                        {format!("{} u32 values", instruction_data.len())}
+                    </span>
+                </div>
+                <div class="info-row">
+                    <span class="info-label">"Proof Size:"</span>
+                    <span class="info-value">{format!("{proof_len} bytes")}</span>
+                </div>
+                <div class="info-row">
+                    <span class="info-label">"Signatures:"</span>
+                    <span class="info-value">{signatures_count.to_string()}</span>
+                </div>
+            </div>
+
+            <h3>"Accounts"</h3>
+            <AccountNonceList account_ids=account_ids nonces=nonces />
+        </div>
+    }
+}
+
+/// Privacy-preserving transaction details component
+#[component]
+pub fn PrivacyPreservingTxDetails(tx: PrivacyPreservingTransaction) -> impl IntoView {
+    let PrivacyPreservingTransaction {
+        hash: _,
+        message,
+        witness_set,
+    } = tx;
+    let PrivacyPreservingMessage {
+        public_account_ids,
+        nonces,
+        public_post_states: _,
+        encrypted_private_post_states,
+        new_commitments,
+        new_nullifiers,
+        block_validity_window,
+        timestamp_validity_window,
+    } = message;
+    let WitnessSet {
+        signatures_and_public_keys: _,
+        proof,
+    } = witness_set;
+    let proof_len = proof.map_or(0, |p| p.0.len());
+
+    view! {
+        <div class="transaction-details">
+            <h2>"Privacy-Preserving Transaction Details"</h2>
+            <div class="info-grid">
+                <div class="info-row">
+                    <span class="info-label">"Public Accounts:"</span>
+                    <span class="info-value">
+                        {public_account_ids.len().to_string()}
+                    </span>
+                </div>
+                <div class="info-row">
+                    <span class="info-label">"New Commitments:"</span>
+                    <span class="info-value">{new_commitments.len().to_string()}</span>
+                </div>
+                <div class="info-row">
+                    <span class="info-label">"Nullifiers:"</span>
+                    <span class="info-value">{new_nullifiers.len().to_string()}</span>
+                </div>
+                <div class="info-row">
+                    <span class="info-label">"Encrypted States:"</span>
+                    <span class="info-value">
+                        {encrypted_private_post_states.len().to_string()}
+                    </span>
+                </div>
+                <div class="info-row">
+                    <span class="info-label">"Proof Size:"</span>
+                    <span class="info-value">{format!("{proof_len} bytes")}</span>
+                </div>
+                <div class="info-row">
+                    <span class="info-label">"Block Validity Window:"</span>
+                    <span class="info-value">{block_validity_window.to_string()}</span>
+                </div>
+                <div class="info-row">
+                    <span class="info-label">"Timestamp Validity Window:"</span>
+                    <span class="info-value">{timestamp_validity_window.to_string()}</span>
+                </div>
+            </div>
+
+            <h3>"Public Accounts"</h3>
+            <AccountNonceList account_ids=public_account_ids nonces=nonces />
+        </div>
+    }
+}
+
+/// Program deployment transaction details component
+#[component]
+pub fn ProgramDeploymentTxDetails(tx: ProgramDeploymentTransaction) -> impl IntoView {
+    let ProgramDeploymentTransaction { hash: _, message } = tx;
+    let ProgramDeploymentMessage { bytecode } = message;
+
+    let bytecode_len = bytecode.len();
+    view! {
+        <div class="transaction-details">
+            <h2>"Program Deployment Transaction Details"</h2>
+            <div class="info-grid">
+                <div class="info-row">
+                    <span class="info-label">"Bytecode Size:"</span>
+                    <span class="info-value">
+                        {format!("{bytecode_len} bytes")}
+                    </span>
+                </div>
+            </div>
+        </div>
+    }
+}

lez/explorer_service/src/pages/main_page.rs

@@ -6,8 +6,8 @@ use leptos_router::{
 use web_sys::SubmitEvent;
 
 use crate::{
-    api::{self, SearchResults},
-    components::{AccountPreview, BlockPreview, TransactionPreview},
+    api,
+    components::{BlockPreview, SearchResultsView},
 };
 
 const RECENT_BLOCKS_LIMIT: u64 = 10;
@@ -138,93 +138,8 @@ pub fn MainPage() -> impl IntoView {
                             .get()
                             .and_then(|opt_results| opt_results)
                             .map(|results| {
-                                let SearchResults {
-                                    blocks,
-                                    transactions,
-                                    accounts,
-                                } = results;
-                                let has_results = !blocks.is_empty()
-                                    || !transactions.is_empty()
-                                    || !accounts.is_empty();
-                                view! {
-                                    <div class="search-results">
-                                        <h2>"Search Results"</h2>
-                                        {if has_results {
-                                            view! {
-                                                <div class="results-container">
-                                                    {if blocks.is_empty() {
-                                                        ().into_any()
-                                                    } else {
-                                                        view! {
-                                                            <div class="results-section">
-                                                                <h3>"Blocks"</h3>
-                                                                <div class="results-list">
-                                                                    {blocks
-                                                                        .into_iter()
-                                                                        .map(|block| {
-                                                                            view! { <BlockPreview block=block /> }
-                                                                        })
-                                                                        .collect::<Vec<_>>()}
-                                                                </div>
-                                                            </div>
-                                                        }
-                                                            .into_any()
-                                                    }}
-
-                                                    {if transactions.is_empty() {
-                                                        ().into_any()
-                                                    } else {
-                                                        view! {
-                                                            <div class="results-section">
-                                                                <h3>"Transactions"</h3>
-                                                                <div class="results-list">
-                                                                    {transactions
-                                                                        .into_iter()
-                                                                        .map(|tx| {
-                                                                            view! { <TransactionPreview transaction=tx /> }
-                                                                        })
-                                                                        .collect::<Vec<_>>()}
-                                                                </div>
-                                                            </div>
-                                                        }
-                                                            .into_any()
-                                                    }}
-
-                                                    {if accounts.is_empty() {
-                                                        ().into_any()
-                                                    } else {
-                                                        view! {
-                                                            <div class="results-section">
-                                                                <h3>"Accounts"</h3>
-                                                                <div class="results-list">
-                                                                    {accounts
-                                                                        .into_iter()
-                                                                        .map(|(id, account)| {
-                                                                            view! {
-                                                                                <AccountPreview
-                                                                                    account_id=id
-                                                                                    account=account
-                                                                                />
-                                                                            }
-                                                                        })
-                                                                        .collect::<Vec<_>>()}
-                                                                </div>
-                                                            </div>
-                                                        }
-                                                            .into_any()
-                                                    }}
-
-                                                </div>
-                                            }
-                                                .into_any()
-                                        } else {
-                                                view! { <div class="not-found">"No results found"</div> }
-                                                .into_any()
-                                        }}
-                                </div>
-                            }
-                                .into_any()
-                        })
+                                view! { <SearchResultsView results=results /> }.into_any()
+                            })
                     }}
 
                 </Suspense>

lez/explorer_service/src/pages/transaction_page.rs

@@ -1,14 +1,13 @@
 use std::str::FromStr as _;
 
-use indexer_service_protocol::{
-    HashType, PrivacyPreservingMessage, PrivacyPreservingTransaction, ProgramDeploymentMessage,
-    ProgramDeploymentTransaction, PublicMessage, PublicTransaction, Transaction, WitnessSet,
-};
-use itertools::{EitherOrBoth, Itertools as _};
+use indexer_service_protocol::{HashType, Transaction};
 use leptos::prelude::*;
-use leptos_router::{components::A, hooks::use_params_map};
+use leptos_router::hooks::use_params_map;
 
-use crate::api;
+use crate::{
+    api,
+    components::{PrivacyPreservingTxDetails, ProgramDeploymentTxDetails, PublicTxDetails},
+};
 
 /// Transaction page component
 #[component]
@@ -66,244 +65,21 @@ pub fn TransactionPage() -> impl IntoView {
 
                                         {
                                             match tx {
-                                Transaction::Public(ptx) => {
-                                    let PublicTransaction {
-                                        hash: _,
-                                        message,
-                                        witness_set,
-                                    } = ptx;
-                                    let PublicMessage {
-                                        program_id,
-                                        account_ids,
-                                        nonces,
-                                        instruction_data,
-                                    } = message;
-                                    let WitnessSet {
-                                        signatures_and_public_keys,
-                                        proof,
-                                    } = witness_set;
+                                                Transaction::Public(ptx) => {
+                                                    view! { <PublicTxDetails tx=ptx /> }.into_any()
+                                                }
+                                                Transaction::PrivacyPreserving(pptx) => {
+                                                    view! { <PrivacyPreservingTxDetails tx=pptx /> }.into_any()
+                                                }
+                                                Transaction::ProgramDeployment(pdtx) => {
+                                                    view! { <ProgramDeploymentTxDetails tx=pdtx /> }.into_any()
+                                                }
+                                            }
+                                        }
 
-                                    let program_id_str = program_id.to_string();
-                                    let proof_len = proof.map_or(0, |p| p.0.len());
-                                    let signatures_count = signatures_and_public_keys.len();
-
-                                    view! {
-                                        <div class="transaction-details">
-                                            <h2>"Public Transaction Details"</h2>
-                                            <div class="info-grid">
-                                                <div class="info-row">
-                                                    <span class="info-label">"Program ID:"</span>
-                                                    <span class="info-value hash">{program_id_str}</span>
-                                                </div>
-                                                <div class="info-row">
-                                                    <span class="info-label">"Instruction Data:"</span>
-                                                    <span class="info-value">
-                                                        {format!("{} u32 values", instruction_data.len())}
-                                                    </span>
-                                                </div>
-                                                <div class="info-row">
-                                                    <span class="info-label">"Proof Size:"</span>
-                                                    <span class="info-value">{format!("{proof_len} bytes")}</span>
-                                                </div>
-                                                <div class="info-row">
-                                                    <span class="info-label">"Signatures:"</span>
-                                                    <span class="info-value">{signatures_count.to_string()}</span>
-                                                </div>
-                                            </div>
-
-                                            <h3>"Accounts"</h3>
-                                            <div class="accounts-list">
-                                                {account_ids
-                                                    .into_iter()
-                                                    .zip_longest(nonces.into_iter())
-                                                    .map(|maybe_pair| {
-                                                        match maybe_pair {
-                                                            EitherOrBoth::Both(account_id, nonce) => {
-                                                                let account_id_str = account_id.to_string();
-                                                        view! {
-                                                            <div class="account-item">
-                                                                <A href=format!("/account/{}", account_id_str)>
-                                                                    <span class="hash">{account_id_str}</span>
-                                                                </A>
-                                                                <span class="nonce">
-                                                                    " (nonce: " {nonce.to_string()} ")"
-                                                                </span>
-                                                            </div>
-                                                        }
-                                                            }
-                                                            EitherOrBoth::Left(account_id) => {
-                                                                let account_id_str = account_id.to_string();
-                                                        view! {
-                                                            <div class="account-item">
-                                                                <A href=format!("/account/{}", account_id_str)>
-                                                                    <span class="hash">{account_id_str}</span>
-                                                                </A>
-                                                                <span class="nonce">
-                                                                    " (nonce: "{"Not affected by this transaction".to_owned()}" )"
-                                                                </span>
-                                                            </div>
-                                                        }
-                                                            }
-                                                            EitherOrBoth::Right(_) => {
-                                                                view! {
-                                                            <div class="account-item">
-                                                                <A href=format!("/account/{}", "Account not found")>
-                                                                    <span class="hash">{"Account not found"}</span>
-                                                                </A>
-                                                                <span class="nonce">
-                                                                    " (nonce: "{"Account not found".to_owned()}" )"
-                                                                </span>
-                                                            </div>
-                                                        }
-                                                            }
-                                                        }
-                                                    })
-                                                    .collect::<Vec<_>>()}
-                                            </div>
-                                        </div>
-                                    }
-                                        .into_any()
+                                    </div>
                                 }
-                                Transaction::PrivacyPreserving(pptx) => {
-                                    let PrivacyPreservingTransaction {
-                                        hash: _,
-                                        message,
-                                        witness_set,
-                                    } = pptx;
-                                    let PrivacyPreservingMessage {
-                                        public_account_ids,
-                                        nonces,
-                                        public_post_states: _,
-                                        encrypted_private_post_states,
-                                        new_commitments,
-                                        new_nullifiers,
-                                        block_validity_window,
-                                        timestamp_validity_window,
-                                    } = message;
-                                    let WitnessSet {
-                                        signatures_and_public_keys: _,
-                                        proof,
-                                    } = witness_set;
-                                    let proof_len = proof.map_or(0, |p| p.0.len());
-                                    view! {
-                                        <div class="transaction-details">
-                                            <h2>"Privacy-Preserving Transaction Details"</h2>
-                                            <div class="info-grid">
-                                                <div class="info-row">
-                                                    <span class="info-label">"Public Accounts:"</span>
-                                                    <span class="info-value">
-                                                        {public_account_ids.len().to_string()}
-                                                    </span>
-                                                </div>
-                                                <div class="info-row">
-                                                    <span class="info-label">"New Commitments:"</span>
-                                                    <span class="info-value">{new_commitments.len().to_string()}</span>
-                                                </div>
-                                                <div class="info-row">
-                                                    <span class="info-label">"Nullifiers:"</span>
-                                                    <span class="info-value">{new_nullifiers.len().to_string()}</span>
-                                                </div>
-                                                <div class="info-row">
-                                                    <span class="info-label">"Encrypted States:"</span>
-                                                    <span class="info-value">
-                                                        {encrypted_private_post_states.len().to_string()}
-                                                    </span>
-                                                </div>
-                                                <div class="info-row">
-                                                    <span class="info-label">"Proof Size:"</span>
-                                                    <span class="info-value">{format!("{proof_len} bytes")}</span>
-                                                </div>
-                                                <div class="info-row">
-                                                    <span class="info-label">"Block Validity Window:"</span>
-                                                    <span class="info-value">{block_validity_window.to_string()}</span>
-                                                </div>
-                                                <div class="info-row">
-                                                    <span class="info-label">"Timestamp Validity Window:"</span>
-                                                    <span class="info-value">{timestamp_validity_window.to_string()}</span>
-                                                </div>
-                                            </div>
-
-                                            <h3>"Public Accounts"</h3>
-                                            <div class="accounts-list">
-                                                {public_account_ids
-                                                    .into_iter()
-                                                    .zip_longest(nonces.into_iter())
-                                                    .map(|maybe_pair| {
-                                                        match maybe_pair {
-                                                            EitherOrBoth::Both(account_id, nonce) => {
-                                                                let account_id_str = account_id.to_string();
-                                                        view! {
-                                                            <div class="account-item">
-                                                                <A href=format!("/account/{}", account_id_str)>
-                                                                    <span class="hash">{account_id_str}</span>
-                                                                </A>
-                                                                <span class="nonce">
-                                                                    " (nonce: " {nonce.to_string()} ")"
-                                                                </span>
-                                                            </div>
-                                                        }
-                                                            }
-                                                            EitherOrBoth::Left(account_id) => {
-                                                                let account_id_str = account_id.to_string();
-                                                        view! {
-                                                            <div class="account-item">
-                                                                <A href=format!("/account/{}", account_id_str)>
-                                                                    <span class="hash">{account_id_str}</span>
-                                                                </A>
-                                                                <span class="nonce">
-                                                                    " (nonce: "{"Not affected by this transaction".to_owned()}" )"
-                                                                </span>
-                                                            </div>
-                                                        }
-                                                            }
-                                                            EitherOrBoth::Right(_) => {
-                                                                view! {
-                                                            <div class="account-item">
-                                                                <A href=format!("/account/{}", "Account not found")>
-                                                                    <span class="hash">{"Account not found"}</span>
-                                                                </A>
-                                                                <span class="nonce">
-                                                                    " (nonce: "{"Account not found".to_owned()}" )"
-                                                                </span>
-                                                            </div>
-                                                        }
-                                                            }
-                                                        }
-                                                    })
-                                                    .collect::<Vec<_>>()}
-                                            </div>
-                                        </div>
-                                    }
-                                        .into_any()
-                                }
-                                Transaction::ProgramDeployment(pdtx) => {
-                                    let ProgramDeploymentTransaction {
-                                        hash: _,
-                                        message,
-                                    } = pdtx;
-                                    let ProgramDeploymentMessage { bytecode } = message;
-
-                                    let bytecode_len = bytecode.len();
-                                    view! {
-                                        <div class="transaction-details">
-                                            <h2>"Program Deployment Transaction Details"</h2>
-                                            <div class="info-grid">
-                                                <div class="info-row">
-                                                    <span class="info-label">"Bytecode Size:"</span>
-                                                    <span class="info-value">
-                                                        {format!("{bytecode_len} bytes")}
-                                                    </span>
-                                                </div>
-                                            </div>
-                                        </div>
-                                    }
-                                        .into_any()
-                                }
-                            }}
-
-                        </div>
-                    }
-                        .into_any()
+                                    .into_any()
                             }
                             Err(e) => {
                                 view! {

lez/indexer/core/src/block_store.rs

@@ -217,6 +217,60 @@ mod tests {
 
     use super::*;
 
+    struct TestFixture {
+        storage: IndexerStore,
+        from: AccountId,
+        to: AccountId,
+        _home: tempfile::TempDir,
+    }
+
+    #[expect(
+        clippy::arithmetic_side_effects,
+        reason = "test helper with bounded inputs"
+    )]
+    async fn store_with_transfer_blocks(
+        block_count: u64,
+        prev_hash: Option<common::HashType>,
+    ) -> TestFixture {
+        let home = tempdir().unwrap();
+        let storage = IndexerStore::open_db(home.path(), Vec::new()).unwrap();
+
+        let initial_accounts = initial_pub_accounts_private_keys();
+        let from = initial_accounts[0].account_id;
+        let to = initial_accounts[1].account_id;
+        let sign_key = initial_accounts[0].pub_sign_key.clone();
+
+        let mut prev_hash = prev_hash;
+        for i in 0..block_count {
+            let tx = common::test_utils::create_transaction_native_token_transfer(
+                from,
+                u128::from(i),
+                to,
+                10,
+                &sign_key,
+            );
+            let block_id = i + 1;
+
+            let next_block = common::test_utils::produce_dummy_block(block_id, prev_hash, vec![tx]);
+            prev_hash = Some(next_block.header.hash);
+
+            storage
+                .put_block(
+                    next_block,
+                    HeaderId::from([u8::try_from(i + 1).unwrap(); 32]),
+                )
+                .await
+                .unwrap();
+        }
+
+        TestFixture {
+            storage,
+            from,
+            to,
+            _home: home,
+        }
+    }
+
     #[test]
     fn correct_startup() {
         let home = tempdir().unwrap();
@@ -228,29 +282,9 @@ mod tests {
         assert_eq!(final_id, None);
     }
 
-    #[tokio::test]
-    async fn seeds_bridge_lock_holding_into_genesis_state() {
-        // The holding is force-inserted, not produced by a transaction, so the
-        // indexer must seed it to match the sequencer. Use the same builder both
-        // sides use, so this also guards against the two drifting.
-        let home = tempdir().unwrap();
-        let holder = AccountId::new([5; 32]);
-        let (id, account) = cross_zone::build_holding_account(holder, 42);
-
-        let storage = IndexerStore::open_db(home.as_ref(), vec![(id, account)]).unwrap();
-
-        let seeded = storage.account_current_state(&holder).await.unwrap();
-        assert_eq!(
-            bridge_lock_core::read_balance(&seeded.data.into_inner()),
-            42,
-            "indexer genesis must hold the seeded bridge-lock balance"
-        );
-    }
-
     #[tokio::test]
     async fn state_transition() {
         let home = tempdir().unwrap();
-
         let storage = IndexerStore::open_db(home.as_ref(), Vec::new()).unwrap();
 
         let initial_accounts = initial_pub_accounts_private_keys();
@@ -258,7 +292,6 @@ mod tests {
         let to = initial_accounts[1].account_id;
         let sign_key = initial_accounts[0].pub_sign_key.clone();
 
-        // Submit genesis block
         let clock_tx = LeeTransaction::Public(clock_invocation(0));
         let genesis_block_data = HashableBlockData {
             block_id: 1,
@@ -274,15 +307,13 @@ mod tests {
             .await
             .unwrap();
 
-        for i in 0..10 {
+        for i in 0..10_u128 {
             let tx = common::test_utils::create_transaction_native_token_transfer(
                 from, i, to, 10, &sign_key,
             );
             let block_id = u64::try_from(i + 1).unwrap();
-
             let next_block = common::test_utils::produce_dummy_block(block_id, prev_hash, vec![tx]);
             prev_hash = Some(next_block.header.hash);
-
             storage
                 .put_block(
                     next_block,
@@ -301,48 +332,23 @@ mod tests {
 
     #[tokio::test]
     async fn account_state_at_block() {
-        let home = tempdir().unwrap();
+        let TestFixture {
+            storage,
+            from,
+            to,
+            _home,
+        } = store_with_transfer_blocks(10, None).await;
 
-        let storage = IndexerStore::open_db(home.as_ref(), Vec::new()).unwrap();
-
-        let mut prev_hash = None;
-
-        let initial_accounts = initial_pub_accounts_private_keys();
-        let from = initial_accounts[0].account_id;
-        let to = initial_accounts[1].account_id;
-        let sign_key = initial_accounts[0].pub_sign_key.clone();
-
-        for i in 0..10 {
-            let tx = common::test_utils::create_transaction_native_token_transfer(
-                from, i, to, 10, &sign_key,
-            );
-            let block_id = u64::try_from(i + 1).unwrap();
-
-            let next_block = common::test_utils::produce_dummy_block(block_id, prev_hash, vec![tx]);
-            prev_hash = Some(next_block.header.hash);
-
-            storage
-                .put_block(
-                    next_block,
-                    HeaderId::from([u8::try_from(i + 1).unwrap(); 32]),
-                )
-                .await
-                .unwrap();
-        }
-
-        // Genesis block: no transfers applied yet.
         let acc1_at_1 = storage.account_state_at_block(&from, 1).unwrap();
         let acc2_at_1 = storage.account_state_at_block(&to, 1).unwrap();
         assert_eq!(acc1_at_1.balance, 9990);
         assert_eq!(acc2_at_1.balance, 20010);
 
-        // After block 5: 4 transfers of 10 applied (one each in blocks 2..=5).
         let acc1_at_5 = storage.account_state_at_block(&from, 5).unwrap();
         let acc2_at_5 = storage.account_state_at_block(&to, 5).unwrap();
         assert_eq!(acc1_at_5.balance, 9950);
         assert_eq!(acc2_at_5.balance, 20050);
 
-        // After final block 9: 8 transfers applied; should match current state.
         let acc1_at_9 = storage.account_state_at_block(&from, 9).unwrap();
         let acc2_at_9 = storage.account_state_at_block(&to, 9).unwrap();
         assert_eq!(acc1_at_9.balance, 9910);

lez/indexer/service/src/mock_service.rs

@@ -330,6 +330,76 @@ impl indexer_service_rpc::RpcServer for MockIndexerService {
     }
 }
 
+fn mock_public_tx(
+    tx_hash: HashType,
+    block_id: BlockId,
+    tx_idx: u64,
+    account_ids: &[AccountId],
+) -> Transaction {
+    Transaction::Public(PublicTransaction {
+        hash: tx_hash,
+        message: PublicMessage {
+            program_id: ProgramId([1_u32; 8]),
+            account_ids: vec![
+                account_ids[tx_idx as usize % account_ids.len()],
+                account_ids[(tx_idx as usize + 1) % account_ids.len()],
+            ],
+            nonces: vec![block_id as u128, (block_id + 1) as u128],
+            instruction_data: vec![1, 2, 3, 4],
+        },
+        witness_set: WitnessSet {
+            signatures_and_public_keys: vec![],
+            proof: None,
+        },
+    })
+}
+
+fn mock_privacy_preserving_tx(
+    tx_hash: HashType,
+    block_id: BlockId,
+    tx_idx: u64,
+    account_ids: &[AccountId],
+) -> Transaction {
+    Transaction::PrivacyPreserving(PrivacyPreservingTransaction {
+        hash: tx_hash,
+        message: PrivacyPreservingMessage {
+            public_account_ids: vec![account_ids[tx_idx as usize % account_ids.len()]],
+            nonces: vec![block_id as u128],
+            public_post_states: vec![Account {
+                program_owner: ProgramId([1_u32; 8]),
+                balance: 500,
+                data: Data(vec![0xdd, 0xee]),
+                nonce: block_id as u128,
+            }],
+            encrypted_private_post_states: vec![EncryptedAccountData {
+                ciphertext: indexer_service_protocol::Ciphertext(vec![0x01, 0x02, 0x03, 0x04]),
+                epk: indexer_service_protocol::EphemeralPublicKey(vec![0xaa; 32]),
+                view_tag: 42,
+            }],
+            new_commitments: vec![Commitment([block_id as u8; 32])],
+            new_nullifiers: vec![(
+                indexer_service_protocol::Nullifier([tx_idx as u8; 32]),
+                CommitmentSetDigest([0xff; 32]),
+            )],
+            block_validity_window: ValidityWindow((None, None)),
+            timestamp_validity_window: ValidityWindow((None, None)),
+        },
+        witness_set: WitnessSet {
+            signatures_and_public_keys: vec![],
+            proof: Some(indexer_service_protocol::Proof(vec![0; 32])),
+        },
+    })
+}
+
+fn mock_program_deployment_tx(tx_hash: HashType) -> Transaction {
+    Transaction::ProgramDeployment(ProgramDeploymentTransaction {
+        hash: tx_hash,
+        message: ProgramDeploymentMessage {
+            bytecode: vec![0x00, 0x61, 0x73, 0x6d, 0x01, 0x00, 0x00, 0x00],
+        },
+    })
+}
+
 fn build_mock_block(
     block_id: BlockId,
     prev_hash: HashType,
@@ -344,7 +414,6 @@ fn build_mock_block(
         HashType(hash)
     };
 
-    // Create 2-4 transactions per block (mix of Public, PrivacyPreserving, and ProgramDeployment)
     let num_txs = 2 + (block_id % 3);
     let mut block_transactions = Vec::new();
 
@@ -356,65 +425,10 @@ fn build_mock_block(
             HashType(hash)
         };
 
-        // Vary transaction types: Public, PrivacyPreserving, or ProgramDeployment
         let tx = match (block_id + tx_idx) % 5 {
-            // Public transactions (most common)
-            0 | 1 => Transaction::Public(PublicTransaction {
-                hash: tx_hash,
-                message: PublicMessage {
-                    program_id: ProgramId([1_u32; 8]),
-                    account_ids: vec![
-                        account_ids[tx_idx as usize % account_ids.len()],
-                        account_ids[(tx_idx as usize + 1) % account_ids.len()],
-                    ],
-                    nonces: vec![block_id as u128, (block_id + 1) as u128],
-                    instruction_data: vec![1, 2, 3, 4],
-                },
-                witness_set: WitnessSet {
-                    signatures_and_public_keys: vec![],
-                    proof: None,
-                },
-            }),
-            // PrivacyPreserving transactions
-            2 | 3 => Transaction::PrivacyPreserving(PrivacyPreservingTransaction {
-                hash: tx_hash,
-                message: PrivacyPreservingMessage {
-                    public_account_ids: vec![account_ids[tx_idx as usize % account_ids.len()]],
-                    nonces: vec![block_id as u128],
-                    public_post_states: vec![Account {
-                        program_owner: ProgramId([1_u32; 8]),
-                        balance: 500,
-                        data: Data(vec![0xdd, 0xee]),
-                        nonce: block_id as u128,
-                    }],
-                    encrypted_private_post_states: vec![EncryptedAccountData {
-                        ciphertext: indexer_service_protocol::Ciphertext(vec![
-                            0x01, 0x02, 0x03, 0x04,
-                        ]),
-                        epk: indexer_service_protocol::EphemeralPublicKey(vec![0xaa; 32]),
-                        view_tag: 42,
-                    }],
-                    new_commitments: vec![Commitment([block_id as u8; 32])],
-                    new_nullifiers: vec![(
-                        indexer_service_protocol::Nullifier([tx_idx as u8; 32]),
-                        CommitmentSetDigest([0xff; 32]),
-                    )],
-                    block_validity_window: ValidityWindow((None, None)),
-                    timestamp_validity_window: ValidityWindow((None, None)),
-                },
-                witness_set: WitnessSet {
-                    signatures_and_public_keys: vec![],
-                    proof: Some(indexer_service_protocol::Proof(vec![0; 32])),
-                },
-            }),
-            // ProgramDeployment transactions (rare)
-            _ => Transaction::ProgramDeployment(ProgramDeploymentTransaction {
-                hash: tx_hash,
-                message: ProgramDeploymentMessage {
-                    bytecode: vec![0x00, 0x61, 0x73, 0x6d, 0x01, 0x00, 0x00, 0x00], /* WASM magic
-                                                                                     * number */
-                },
-            }),
+            0 | 1 => mock_public_tx(tx_hash, block_id, tx_idx, account_ids),
+            2 | 3 => mock_privacy_preserving_tx(tx_hash, block_id, tx_idx, account_ids),
+            _ => mock_program_deployment_tx(tx_hash),
         };
 
         block_transactions.push(tx);

lez/keycard_wallet/src/lib.rs

@@ -134,10 +134,6 @@ impl KeycardWallet {
         })
     }
 
-    #[expect(
-        clippy::arithmetic_side_effects,
-        reason = "64 - s_stripped.len() is safe: s_stripped.len() ≤ 31 because py_signature.len() is in [32, 63]"
-    )]
     pub fn sign_message_for_path(
         &self,
         py: Python,
@@ -150,33 +146,9 @@ impl KeycardWallet {
             .call_method1("sign_message_for_path", (message, path))?
             .extract()?;
 
-        // The keycard Python library strips leading zeros from S when S < 2^(8k) for some k.
-        // Left-pad S back to 32 bytes so the full signature is always 64 bytes (R || S).
-        let py_signature = if py_signature.len() < 64 {
-            if py_signature.len() < 32 {
-                return Err(PyErr::new::<pyo3::exceptions::PyValueError, _>(format!(
-                    "signature from keycard too short: {} bytes",
-                    py_signature.len()
-                )));
-            }
-            let s_stripped = &py_signature[32..];
-            let mut padded = [0_u8; 64];
-            padded[..32].copy_from_slice(&py_signature[..32]);
-            padded[(64 - s_stripped.len())..].copy_from_slice(s_stripped);
-            padded.to_vec()
-        } else {
-            py_signature
+        let sig = Signature {
+            value: normalize_keycard_signature(py_signature)?,
         };
-
-        let signature: [u8; 64] = py_signature.try_into().map_err(|vec: Vec<u8>| {
-            PyErr::new::<pyo3::exceptions::PyValueError, _>(format!(
-                "Invalid signature length: expected 64 bytes, got {} (bytes: {:02x?})",
-                vec.len(),
-                vec
-            ))
-        })?;
-
-        let sig = Signature { value: signature };
         let pub_key = self.get_public_key_for_path(py, path)?;
         if !sig.is_valid_for(message, &pub_key) {
             return Err(PyErr::new::<pyo3::exceptions::PyValueError, _>(
@@ -224,32 +196,8 @@ impl KeycardWallet {
             .call_method1("get_private_keys_for_path", (path,))?
             .extract()?;
 
-        let raw_nsk = Zeroizing::new(raw_nsk);
-        let raw_vsk = Zeroizing::new(raw_vsk);
-
-        let nsk = {
-            if raw_nsk.len() != 32 {
-                return Err(PyErr::new::<pyo3::exceptions::PyValueError, _>(format!(
-                    "expected 32-byte NSK from keycard, got {} bytes",
-                    raw_nsk.len()
-                )));
-            }
-            let mut arr = Zeroizing::new([0_u8; 32]);
-            arr.copy_from_slice(&raw_nsk);
-            arr
-        };
-
-        let vsk = {
-            if raw_vsk.len() != 64 {
-                return Err(PyErr::new::<pyo3::exceptions::PyValueError, _>(format!(
-                    "expected 64-byte VSK from keycard, got {} bytes",
-                    raw_vsk.len()
-                )));
-            }
-            let mut arr = Zeroizing::new([0_u8; 64]);
-            arr.copy_from_slice(&raw_vsk);
-            arr
-        };
+        let nsk = zeroizing_fixed_bytes::<32>("nullifier secret key", Zeroizing::new(raw_nsk))?;
+        let vsk = zeroizing_fixed_bytes::<64>("view secret key", Zeroizing::new(raw_vsk))?;
 
         Ok((nsk, vsk))
     }
@@ -269,6 +217,55 @@ impl KeycardWallet {
     }
 }
 
+/// The keycard Python library strips leading zeros from S when S < 2^(8k) for some k.
+/// Left-pad S back to 32 bytes so the full signature is always 64 bytes (R || S).
+#[expect(
+    clippy::arithmetic_side_effects,
+    reason = "64 - s_stripped.len() is safe: s_stripped.len() ≤ 31 because py_signature.len() is in [32, 63]"
+)]
+fn normalize_keycard_signature(py_signature: Vec<u8>) -> PyResult<[u8; 64]> {
+    if py_signature.len() < 64 {
+        if py_signature.len() < 32 {
+            return Err(PyErr::new::<pyo3::exceptions::PyValueError, _>(format!(
+                "signature from keycard too short: {} bytes",
+                py_signature.len()
+            )));
+        }
+        let s_stripped = &py_signature[32..];
+        let mut padded = [0_u8; 64];
+        padded[..32].copy_from_slice(&py_signature[..32]);
+        padded[(64 - s_stripped.len())..].copy_from_slice(s_stripped);
+        Ok(padded)
+    } else {
+        py_signature.try_into().map_err(|vec: Vec<u8>| {
+            PyErr::new::<pyo3::exceptions::PyValueError, _>(format!(
+                "Invalid signature length: expected 64 bytes, got {} (bytes: {:02x?})",
+                vec.len(),
+                vec
+            ))
+        })
+    }
+}
+
+#[expect(
+    clippy::needless_pass_by_value,
+    reason = "Zeroizing<Vec<u8>> is consumed to ensure the source is zeroed on drop"
+)]
+fn zeroizing_fixed_bytes<const N: usize>(
+    label: &str,
+    raw: Zeroizing<Vec<u8>>,
+) -> PyResult<Zeroizing<[u8; N]>> {
+    if raw.len() != N {
+        return Err(PyErr::new::<pyo3::exceptions::PyValueError, _>(format!(
+            "expected {N}-byte {label} from keycard, got {} bytes",
+            raw.len()
+        )));
+    }
+    let mut arr = Zeroizing::new([0_u8; N]);
+    arr.copy_from_slice(&raw);
+    Ok(arr)
+}
+
 fn pairing_file_path() -> Option<PathBuf> {
     let home = std::env::var("LEE_WALLET_HOME_DIR")
         .map(PathBuf::from)

lez/keycard_wallet/src/python_path.rs

@@ -2,8 +2,7 @@ use std::{env, path::PathBuf};
 
 use pyo3::{prelude::*, types::PyList};
 
-/// Adds the project's `python/` directory and venv site-packages to Python's sys.path.
-pub fn add_python_path(py: Python<'_>) -> PyResult<()> {
+fn collect_python_paths() -> Vec<PathBuf> {
     let current_dir = env::current_dir().expect("Failed to get current working directory");
 
     let python_base = env::var("VIRTUAL_ENV")
@@ -11,7 +10,7 @@ pub fn add_python_path(py: Python<'_>) -> PyResult<()> {
         .and_then(|v| PathBuf::from(v).parent().map(PathBuf::from))
         .unwrap_or_else(|| current_dir.clone());
 
-    let mut paths_to_add: Vec<PathBuf> = vec![
+    let mut paths = vec![
         python_base
             .join("lez")
             .join("keycard_wallet")
@@ -23,24 +22,28 @@ pub fn add_python_path(py: Python<'_>) -> PyResult<()> {
             .join("keycard-py"),
     ];
 
-    // If a virtualenv is active, add its site-packages so that dependencies
-    // installed in the venv (e.g. smartcard, ecdsa) are importable by the
-    // pyo3 embedded interpreter, which does not inherit sys.path from the
-    // shell's `python3` executable.
+    // pyo3's embedded interpreter does not inherit sys.path from the shell,
+    // so venv site-packages must be added explicitly.
     if let Ok(venv) = env::var("VIRTUAL_ENV") {
         let lib = PathBuf::from(&venv).join("lib");
         if let Ok(entries) = std::fs::read_dir(&lib) {
             for entry in entries.flatten() {
                 let site_packages = entry.path().join("site-packages");
                 if site_packages.exists() {
-                    paths_to_add.push(site_packages);
+                    paths.push(site_packages);
                 }
             }
         }
     }
 
-    // Sanity check — warns early if a path doesn't exist
-    for path in &paths_to_add {
+    paths
+}
+
+/// Adds the project's `python/` directory and venv site-packages to Python's sys.path.
+pub fn add_python_path(py: Python<'_>) -> PyResult<()> {
+    let paths = collect_python_paths();
+
+    for path in &paths {
         if !path.exists() {
             log::info!("Warning: Python path does not exist: {}", path.display());
         }
@@ -50,10 +53,9 @@ pub fn add_python_path(py: Python<'_>) -> PyResult<()> {
     let binding = sys.getattr("path")?;
     let sys_path = binding.cast::<PyList>()?;
 
-    for path in &paths_to_add {
+    for path in &paths {
         let path_str = path.to_str().expect("Invalid path");
 
-        // Avoid duplicating the path
         let already_present = sys_path
             .iter()
             .any(|p| p.extract::<&str>().is_ok_and(|s| s == path_str));

lez/sequencer/core/src/block_publisher.rs

@@ -4,7 +4,7 @@ use anyhow::{Context as _, Result, anyhow};
 use common::block::Block;
 use log::{info, warn};
 pub use logos_blockchain_core::mantle::ops::channel::MsgId;
-use logos_blockchain_core::mantle::ops::channel::inscribe::Inscription;
+use logos_blockchain_core::mantle::ops::channel::{ChannelId, inscribe::Inscription};
 pub use logos_blockchain_key_management_system_service::keys::{Ed25519Key, ZkKey};
 pub use logos_blockchain_zone_sdk::sequencer::SequencerCheckpoint;
 use logos_blockchain_zone_sdk::{
@@ -61,11 +61,14 @@ pub trait BlockPublisherTrait: Clone {
     /// Fire-and-forget publish. Zone-sdk drives the actual submission and
     /// retries internally; this just hands the payload off.
     async fn publish_block(&self, block: &Block, withdrawals: Vec<WithdrawArg>) -> Result<()>;
+
+    fn channel_id(&self) -> ChannelId;
 }
 
 /// Real block publisher backed by zone-sdk's `ZoneSequencer`.
 #[derive(Clone)]
 pub struct ZoneSdkPublisher {
+    channel_id: ChannelId,
     publish_tx: mpsc::Sender<(Inscription, Vec<WithdrawArg>)>,
     // Aborts the drive task when the last clone is dropped.
     _drive_task: Arc<DriveTaskGuard>,
@@ -192,6 +195,7 @@ impl BlockPublisherTrait for ZoneSdkPublisher {
             .context("Zone-sdk readiness channel closed before becoming ready")?;
 
         Ok(Self {
+            channel_id: config.channel_id,
             publish_tx,
             _drive_task: Arc::new(DriveTaskGuard(drive_task)),
         })
@@ -210,6 +214,10 @@ impl BlockPublisherTrait for ZoneSdkPublisher {
 
         Ok(())
     }
+
+    fn channel_id(&self) -> ChannelId {
+        self.channel_id
+    }
 }
 
 /// Deserialize inscription payload as a `Block` and return it's`block_id`.

lez/sequencer/core/src/mock.rs

@@ -2,6 +2,7 @@ use std::time::Duration;
 
 use anyhow::Result;
 use common::block::Block;
+use logos_blockchain_core::mantle::ops::channel::ChannelId;
 use logos_blockchain_key_management_system_service::keys::Ed25519Key;
 use logos_blockchain_zone_sdk::sequencer::WithdrawArg;
 
@@ -16,11 +17,13 @@ use crate::{
 pub type SequencerCoreWithMockClients = crate::SequencerCore<MockBlockPublisher>;
 
 #[derive(Clone)]
-pub struct MockBlockPublisher;
+pub struct MockBlockPublisher {
+    channel_id: ChannelId,
+}
 
 impl BlockPublisherTrait for MockBlockPublisher {
     async fn new(
-        _config: &BedrockConfig,
+        config: &BedrockConfig,
         _bedrock_signing_key: Ed25519Key,
         _resubmit_interval: Duration,
         _initial_checkpoint: Option<SequencerCheckpoint>,
@@ -29,7 +32,9 @@ impl BlockPublisherTrait for MockBlockPublisher {
         _on_deposit_event: OnDepositEventSink,
         _on_withdraw_event: OnWithdrawEventSink,
     ) -> Result<Self> {
-        Ok(Self)
+        Ok(Self {
+            channel_id: config.channel_id,
+        })
     }
 
     async fn publish_block(
@@ -39,4 +44,8 @@ impl BlockPublisherTrait for MockBlockPublisher {
     ) -> Result<()> {
         Ok(())
     }
+
+    fn channel_id(&self) -> ChannelId {
+        self.channel_id
+    }
 }

lez/sequencer/service/protocol/Cargo.toml

@@ -11,3 +11,6 @@ workspace = true
 common.workspace = true
 lee.workspace = true
 lee_core.workspace = true
+
+hex.workspace = true
+serde_with.workspace = true

lez/sequencer/service/protocol/src/lib.rs

@@ -1,5 +1,28 @@
 //! Reexports of types used by sequencer rpc specification.
 
+use std::{fmt::Display, str::FromStr};
+
 pub use common::{HashType, block::Block, transaction::LeeTransaction};
 pub use lee::{Account, AccountId, ProgramId};
 pub use lee_core::{BlockId, Commitment, MembershipProof, account::Nonce};
+use serde_with::{DeserializeFromStr, SerializeDisplay};
+
+#[derive(Debug, Clone, PartialEq, Eq, Hash, SerializeDisplay, DeserializeFromStr)]
+pub struct ChannelId(pub [u8; 32]);
+
+impl Display for ChannelId {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        let hex_string = hex::encode(self.0);
+        write!(f, "{hex_string}")
+    }
+}
+
+impl FromStr for ChannelId {
+    type Err = hex::FromHexError;
+
+    fn from_str(s: &str) -> Result<Self, Self::Err> {
+        let mut bytes = [0_u8; 32];
+        hex::decode_to_slice(s, &mut bytes)?;
+        Ok(Self(bytes))
+    }
+}

lez/sequencer/service/rpc/src/lib.rs

@@ -6,8 +6,8 @@ use jsonrpsee::types::ErrorObjectOwned;
 #[cfg(feature = "client")]
 pub use jsonrpsee::{core::ClientError, http_client::HttpClientBuilder as SequencerClientBuilder};
 use sequencer_service_protocol::{
-    Account, AccountId, Block, BlockId, Commitment, HashType, LeeTransaction, MembershipProof,
-    Nonce, ProgramId,
+    Account, AccountId, Block, BlockId, ChannelId, Commitment, HashType, LeeTransaction,
+    MembershipProof, Nonce, ProgramId,
 };
 
 #[cfg(all(not(feature = "server"), not(feature = "client")))]
@@ -88,5 +88,8 @@ pub trait Rpc {
     #[method(name = "getProgramIds")]
     async fn get_program_ids(&self) -> Result<BTreeMap<String, ProgramId>, ErrorObjectOwned>;
 
+    #[method(name = "getChannelId")]
+    async fn get_channel_id(&self) -> Result<ChannelId, ErrorObjectOwned>;
+
     // =============================================================================================
 }

lez/sequencer/service/src/service.rs

@@ -12,7 +12,8 @@ use sequencer_core::{
     DbError, SequencerCore, TransactionOrigin, block_publisher::BlockPublisherTrait,
 };
 use sequencer_service_protocol::{
-    Account, AccountId, Block, BlockId, Commitment, HashType, MembershipProof, Nonce, ProgramId,
+    Account, AccountId, Block, BlockId, ChannelId, Commitment, HashType, MembershipProof, Nonce,
+    ProgramId,
 };
 use tokio::sync::Mutex;
 
@@ -190,6 +191,11 @@ impl<BC: BlockPublisherTrait + Send + 'static> sequencer_service_rpc::RpcServer
         );
         Ok(program_ids)
     }
+
+    async fn get_channel_id(&self) -> Result<ChannelId, ErrorObjectOwned> {
+        let channel_id = self.sequencer.lock().await.block_publisher().channel_id();
+        Ok(ChannelId(*channel_id.as_ref()))
+    }
 }
 
 fn internal_error(err: &DbError) -> ErrorObjectOwned {

lez/storage/src/indexer/mod.rs

@@ -153,97 +153,19 @@ impl RocksDBIO {
         }
 
         let br_id = closest_breakpoint_id(block_id);
-        let mut breakpoint = self.get_breakpoint(br_id)?;
+        let mut state = self.get_breakpoint(br_id)?;
 
         let start = u64::from(BREAKPOINT_INTERVAL)
             .checked_mul(br_id)
             .expect("Reached maximum breakpoint id");
 
-        for mut block in self.get_block_batch_seq(
+        for block in self.get_block_batch_seq(
             start.checked_add(1).expect("Will be lesser that u64::MAX")..=block_id,
         )? {
-            let expected_clock = LeeTransaction::Public(clock_invocation(block.header.timestamp));
-
-            let clock_tx = block.body.transactions.pop().ok_or_else(|| {
-                DbError::db_interaction_error(
-                    "Block must contain clock transaction at the end".to_owned(),
-                )
-            })?;
-            let user_txs = block.body.transactions;
-
-            if clock_tx != expected_clock {
-                return Err(DbError::db_interaction_error(
-                        "Last transaction in block must be the clock invocation for the block timestamp"
-                            .to_owned(),
-                    ));
-            }
-            for transaction in user_txs {
-                let is_genesis = block.header.block_id == GENESIS_BLOCK_ID;
-                if is_genesis {
-                    let genesis_tx = match transaction {
-                        LeeTransaction::Public(public_tx) => public_tx,
-                        LeeTransaction::PrivacyPreserving(_)
-                        | LeeTransaction::ProgramDeployment(_) => {
-                            return Err(DbError::db_interaction_error(
-                                "Genesis block should contain only public transactions".to_owned(),
-                            ));
-                        }
-                    };
-                    breakpoint
-                        .transition_from_public_transaction(
-                            &genesis_tx,
-                            block.header.block_id,
-                            block.header.timestamp,
-                        )
-                        .map_err(|err| {
-                            DbError::db_interaction_error(format!(
-                                "genesis transaction execution failed with err {err:?}"
-                            ))
-                        })?;
-                } else {
-                    transaction
-                        .transaction_stateless_check()
-                        .map_err(|err| {
-                            DbError::db_interaction_error(format!(
-                                "transaction pre check failed with err {err:?}"
-                            ))
-                        })?
-                        // FIXME: HOT FIX (testnet v0.2): does not check for system account updates due to
-                        // sequencer-generated deposit tx'es;
-                        // CHANGE ME back to `execute_check_on_state` when the indexer can authenticate deposit transactions
-                        .execute_without_system_accounts_check_on_state(
-                            &mut breakpoint,
-                            block.header.block_id,
-                            block.header.timestamp,
-                        )
-                        .map_err(|err| {
-                            DbError::db_interaction_error(format!(
-                                "transaction execution failed with err {err:?}"
-                            ))
-                        })?;
-                }
-            }
-
-            let LeeTransaction::Public(clock_public_tx) = clock_tx else {
-                return Err(DbError::db_interaction_error(
-                    "Clock invocation must be a public transaction".to_owned(),
-                ));
-            };
-
-            breakpoint
-                .transition_from_public_transaction(
-                    &clock_public_tx,
-                    block.header.block_id,
-                    block.header.timestamp,
-                )
-                .map_err(|err| {
-                    DbError::db_interaction_error(format!(
-                        "clock transaction execution failed with err {err:?}"
-                    ))
-                })?;
+            apply_block_transactions(block, &mut state)?;
         }
 
-        Ok(breakpoint)
+        Ok(state)
     }
 
     pub fn final_state(&self) -> DbResult<V03State> {
@@ -252,6 +174,86 @@ impl RocksDBIO {
     }
 }
 
+fn apply_block_transactions(mut block: Block, state: &mut V03State) -> DbResult<()> {
+    let expected_clock = LeeTransaction::Public(clock_invocation(block.header.timestamp));
+
+    let clock_tx = block.body.transactions.pop().ok_or_else(|| {
+        DbError::db_interaction_error("Block must contain clock transaction at the end".to_owned())
+    })?;
+
+    if clock_tx != expected_clock {
+        return Err(DbError::db_interaction_error(
+            "Last transaction in block must be the clock invocation for the block timestamp"
+                .to_owned(),
+        ));
+    }
+
+    for transaction in block.body.transactions {
+        if block.header.block_id == GENESIS_BLOCK_ID {
+            let genesis_tx = match transaction {
+                LeeTransaction::Public(public_tx) => public_tx,
+                LeeTransaction::PrivacyPreserving(_) | LeeTransaction::ProgramDeployment(_) => {
+                    return Err(DbError::db_interaction_error(
+                        "Genesis block should contain only public transactions".to_owned(),
+                    ));
+                }
+            };
+            state
+                .transition_from_public_transaction(
+                    &genesis_tx,
+                    block.header.block_id,
+                    block.header.timestamp,
+                )
+                .map_err(|err| {
+                    DbError::db_interaction_error(format!(
+                        "genesis transaction execution failed with err {err:?}"
+                    ))
+                })?;
+        } else {
+            transaction
+                .transaction_stateless_check()
+                .map_err(|err| {
+                    DbError::db_interaction_error(format!(
+                        "transaction pre check failed with err {err:?}"
+                    ))
+                })?
+                // FIXME: HOT FIX (testnet v0.2): does not check for system account updates due to
+                // sequencer-generated deposit tx'es;
+                // CHANGE ME back to `execute_check_on_state` when the indexer can authenticate deposit transactions
+                .execute_without_system_accounts_check_on_state(
+                    state,
+                    block.header.block_id,
+                    block.header.timestamp,
+                )
+                .map_err(|err| {
+                    DbError::db_interaction_error(format!(
+                        "transaction execution failed with err {err:?}"
+                    ))
+                })?;
+        }
+    }
+
+    let LeeTransaction::Public(clock_public_tx) = clock_tx else {
+        return Err(DbError::db_interaction_error(
+            "Clock invocation must be a public transaction".to_owned(),
+        ));
+    };
+
+    state
+        .transition_from_public_transaction(
+            &clock_public_tx,
+            block.header.block_id,
+            block.header.timestamp,
+        )
+        .map_err(|err| {
+            DbError::db_interaction_error(format!(
+                "clock transaction execution failed with err {err:?}"
+            ))
+        })?;
+
+    Ok(())
+}
+
 fn closest_breakpoint_id(block_id: u64) -> u64 {
     block_id
         .saturating_sub(1)
@@ -261,439 +263,4 @@ fn closest_breakpoint_id(block_id: u64) -> u64 {
 
 #[expect(clippy::shadow_unrelated, reason = "Fine for tests")]
 #[cfg(test)]
-mod tests {
-    use common::test_utils::produce_dummy_block;
-    use lee::{Account, AccountId, PublicKey};
-    use tempfile::tempdir;
-
-    use super::*;
-
-    fn genesis_block() -> Block {
-        produce_dummy_block(1, None, vec![])
-    }
-
-    fn acc1_sign_key() -> lee::PrivateKey {
-        lee::PrivateKey::try_new([1; 32]).unwrap()
-    }
-
-    fn acc2_sign_key() -> lee::PrivateKey {
-        lee::PrivateKey::try_new([2; 32]).unwrap()
-    }
-
-    fn acc1() -> AccountId {
-        AccountId::from(&PublicKey::new_from_private_key(&acc1_sign_key()))
-    }
-
-    fn acc2() -> AccountId {
-        AccountId::from(&PublicKey::new_from_private_key(&acc2_sign_key()))
-    }
-
-    fn initial_state() -> lee::V03State {
-        let mut public_accounts = [(acc1(), 10000), (acc2(), 20000)]
-            .into_iter()
-            .map(|(id, balance)| {
-                (
-                    id,
-                    Account {
-                        program_owner: programs::authenticated_transfer().id(),
-                        balance,
-                        ..Account::default()
-                    },
-                )
-            })
-            .collect::<Vec<_>>();
-        for clock_id in system_accounts::clock_account_ids() {
-            public_accounts.push((clock_id, system_accounts::clock_account()));
-        }
-
-        lee::V03State::new()
-            .with_public_accounts(public_accounts)
-            .with_programs([programs::authenticated_transfer(), programs::clock()])
-    }
-
-    #[test]
-    fn start_db() {
-        let temp_dir = tempdir().unwrap();
-        let temdir_path = temp_dir.path();
-
-        let dbio = RocksDBIO::open_or_create(temdir_path, &initial_state()).unwrap();
-
-        let last_id = dbio.get_meta_last_block_id_in_db().unwrap();
-        let first_id = dbio.get_meta_first_block_id_in_db().unwrap();
-        let is_first_set = dbio.get_meta_is_first_block_set().unwrap();
-        let last_observed_l1_header = dbio.get_meta_last_observed_l1_lib_header_in_db().unwrap();
-        let last_br_id = dbio.get_meta_last_breakpoint_id().unwrap();
-        let last_block = dbio.get_block(1).unwrap();
-        let breakpoint = dbio.get_breakpoint(0).unwrap();
-        let final_state = dbio.final_state().unwrap();
-
-        assert_eq!(last_id, None);
-        assert_eq!(first_id, None);
-        assert_eq!(last_observed_l1_header, None);
-        assert!(!is_first_set);
-        assert_eq!(last_br_id, Some(0)); // TODO: Will be None after we remove hardcoded testnet state
-        assert!(last_block.is_none());
-        assert_eq!(
-            breakpoint.get_account_by_id(acc1()),
-            final_state.get_account_by_id(acc1())
-        );
-        assert_eq!(
-            breakpoint.get_account_by_id(acc2()),
-            final_state.get_account_by_id(acc2())
-        );
-    }
-
-    #[test]
-    fn one_block_insertion() {
-        let temp_dir = tempdir().unwrap();
-        let temdir_path = temp_dir.path();
-
-        let dbio = RocksDBIO::open_or_create(temdir_path, &initial_state()).unwrap();
-
-        let genesis_block = genesis_block();
-        dbio.put_block(&genesis_block, [0; 32]).unwrap();
-
-        let prev_hash = genesis_block.header.hash;
-        let from = acc1();
-        let to = acc2();
-        let sign_key = acc1_sign_key();
-
-        let transfer_tx =
-            common::test_utils::create_transaction_native_token_transfer(from, 0, to, 1, &sign_key);
-        let block = produce_dummy_block(2, Some(prev_hash), vec![transfer_tx]);
-
-        dbio.put_block(&block, [1; 32]).unwrap();
-
-        let last_id = dbio.get_meta_last_block_id_in_db().unwrap().unwrap();
-        let first_id = dbio.get_meta_first_block_id_in_db().unwrap();
-        let last_observed_l1_header = dbio
-            .get_meta_last_observed_l1_lib_header_in_db()
-            .unwrap()
-            .unwrap();
-        let is_first_set = dbio.get_meta_is_first_block_set().unwrap();
-        let last_br_id = dbio.get_meta_last_breakpoint_id().unwrap();
-        let last_block = dbio.get_block(last_id).unwrap().unwrap();
-        let breakpoint = dbio.get_breakpoint(0).unwrap();
-        let final_state = dbio.final_state().unwrap();
-
-        assert_eq!(last_id, 2);
-        assert_eq!(first_id, Some(1));
-        assert_eq!(last_observed_l1_header, [1; 32]);
-        assert!(is_first_set);
-        assert_eq!(last_br_id, Some(0));
-        assert_eq!(last_block.header.hash, block.header.hash);
-        assert_eq!(
-            breakpoint.get_account_by_id(acc1()).balance
-                - final_state.get_account_by_id(acc1()).balance,
-            1
-        );
-        assert_eq!(
-            final_state.get_account_by_id(acc2()).balance
-                - breakpoint.get_account_by_id(acc2()).balance,
-            1
-        );
-    }
-
-    #[test]
-    fn new_breakpoint() {
-        let temp_dir = tempdir().unwrap();
-        let temdir_path = temp_dir.path();
-
-        let dbio = RocksDBIO::open_or_create(temdir_path, &initial_state()).unwrap();
-
-        let from = acc1();
-        let to = acc2();
-        let sign_key = acc1_sign_key();
-
-        for i in 1..=BREAKPOINT_INTERVAL + 1 {
-            let prev_hash = dbio.get_meta_last_block_id_in_db().unwrap().map(|last_id| {
-                let last_block = dbio.get_block(last_id).unwrap().unwrap();
-                last_block.header.hash
-            });
-
-            let transfer_tx = common::test_utils::create_transaction_native_token_transfer(
-                from,
-                (i - 1).into(),
-                to,
-                1,
-                &sign_key,
-            );
-            let block = produce_dummy_block(i.into(), prev_hash, vec![transfer_tx]);
-            dbio.put_block(&block, [i; 32]).unwrap();
-        }
-
-        let last_id = dbio.get_meta_last_block_id_in_db().unwrap().unwrap();
-        let first_id = dbio.get_meta_first_block_id_in_db().unwrap();
-        let is_first_set = dbio.get_meta_is_first_block_set().unwrap();
-        let last_br_id = dbio.get_meta_last_breakpoint_id().unwrap();
-        let last_block = dbio.get_block(last_id).unwrap().unwrap();
-        let prev_breakpoint = dbio.get_breakpoint(0).unwrap();
-        let breakpoint = dbio.get_breakpoint(1).unwrap();
-        let final_state = dbio.final_state().unwrap();
-
-        assert_eq!(last_id, 101);
-        assert_eq!(first_id, Some(1));
-        assert!(is_first_set);
-        assert_eq!(last_br_id, Some(1));
-        assert_ne!(last_block.header.hash, genesis_block().header.hash);
-        assert_eq!(
-            prev_breakpoint.get_account_by_id(acc1()).balance
-                - final_state.get_account_by_id(acc1()).balance,
-            101
-        );
-        assert_eq!(
-            final_state.get_account_by_id(acc2()).balance
-                - prev_breakpoint.get_account_by_id(acc2()).balance,
-            101
-        );
-        assert_eq!(
-            breakpoint.get_account_by_id(acc1()).balance
-                - final_state.get_account_by_id(acc1()).balance,
-            1
-        );
-        assert_eq!(
-            final_state.get_account_by_id(acc2()).balance
-                - breakpoint.get_account_by_id(acc2()).balance,
-            1
-        );
-    }
-
-    #[test]
-    fn simple_maps() {
-        let temp_dir = tempdir().unwrap();
-        let temdir_path = temp_dir.path();
-
-        let dbio = RocksDBIO::open_or_create(temdir_path, &initial_state()).unwrap();
-
-        let from = acc1();
-        let to = acc2();
-        let sign_key = acc1_sign_key();
-
-        let transfer_tx =
-            common::test_utils::create_transaction_native_token_transfer(from, 0, to, 1, &sign_key);
-        let block = produce_dummy_block(1, None, vec![transfer_tx]);
-
-        let control_hash1 = block.header.hash;
-
-        dbio.put_block(&block, [1; 32]).unwrap();
-
-        let last_id = dbio.get_meta_last_block_id_in_db().unwrap().unwrap();
-        let last_block = dbio.get_block(last_id).unwrap().unwrap();
-
-        let prev_hash = last_block.header.hash;
-        let transfer_tx =
-            common::test_utils::create_transaction_native_token_transfer(from, 1, to, 1, &sign_key);
-        let block = produce_dummy_block(2, Some(prev_hash), vec![transfer_tx]);
-
-        let control_hash2 = block.header.hash;
-
-        dbio.put_block(&block, [2; 32]).unwrap();
-
-        let last_id = dbio.get_meta_last_block_id_in_db().unwrap().unwrap();
-        let last_block = dbio.get_block(last_id).unwrap().unwrap();
-
-        let prev_hash = last_block.header.hash;
-        let transfer_tx =
-            common::test_utils::create_transaction_native_token_transfer(from, 2, to, 1, &sign_key);
-
-        let control_tx_hash1 = transfer_tx.hash();
-
-        let block = produce_dummy_block(3, Some(prev_hash), vec![transfer_tx]);
-        dbio.put_block(&block, [3; 32]).unwrap();
-
-        let last_id = dbio.get_meta_last_block_id_in_db().unwrap().unwrap();
-        let last_block = dbio.get_block(last_id).unwrap().unwrap();
-
-        let prev_hash = last_block.header.hash;
-        let transfer_tx =
-            common::test_utils::create_transaction_native_token_transfer(from, 3, to, 1, &sign_key);
-
-        let control_tx_hash2 = transfer_tx.hash();
-
-        let block = produce_dummy_block(4, Some(prev_hash), vec![transfer_tx]);
-        dbio.put_block(&block, [4; 32]).unwrap();
-
-        let control_block_id1 = dbio.get_block_id_by_hash(control_hash1.0).unwrap().unwrap();
-        let control_block_id2 = dbio.get_block_id_by_hash(control_hash2.0).unwrap().unwrap();
-        let control_block_id3 = dbio
-            .get_block_id_by_tx_hash(control_tx_hash1.0)
-            .unwrap()
-            .unwrap();
-        let control_block_id4 = dbio
-            .get_block_id_by_tx_hash(control_tx_hash2.0)
-            .unwrap()
-            .unwrap();
-
-        assert_eq!(control_block_id1, 1);
-        assert_eq!(control_block_id2, 2);
-        assert_eq!(control_block_id3, 3);
-        assert_eq!(control_block_id4, 4);
-    }
-
-    #[test]
-    fn block_batch() {
-        let temp_dir = tempdir().unwrap();
-        let temdir_path = temp_dir.path();
-
-        let mut block_res = vec![];
-
-        let dbio = RocksDBIO::open_or_create(temdir_path, &initial_state()).unwrap();
-
-        let from = acc1();
-        let to = acc2();
-        let sign_key = acc1_sign_key();
-
-        let transfer_tx =
-            common::test_utils::create_transaction_native_token_transfer(from, 0, to, 1, &sign_key);
-        let block = produce_dummy_block(1, None, vec![transfer_tx]);
-
-        block_res.push(block.clone());
-        dbio.put_block(&block, [1; 32]).unwrap();
-
-        let last_id = dbio.get_meta_last_block_id_in_db().unwrap().unwrap();
-        let last_block = dbio.get_block(last_id).unwrap().unwrap();
-
-        let prev_hash = last_block.header.hash;
-        let transfer_tx =
-            common::test_utils::create_transaction_native_token_transfer(from, 1, to, 1, &sign_key);
-        let block = produce_dummy_block(2, Some(prev_hash), vec![transfer_tx]);
-
-        block_res.push(block.clone());
-        dbio.put_block(&block, [2; 32]).unwrap();
-
-        let last_id = dbio.get_meta_last_block_id_in_db().unwrap().unwrap();
-        let last_block = dbio.get_block(last_id).unwrap().unwrap();
-
-        let prev_hash = last_block.header.hash;
-        let transfer_tx =
-            common::test_utils::create_transaction_native_token_transfer(from, 2, to, 1, &sign_key);
-
-        let block = produce_dummy_block(3, Some(prev_hash), vec![transfer_tx]);
-        block_res.push(block.clone());
-        dbio.put_block(&block, [3; 32]).unwrap();
-
-        let last_id = dbio.get_meta_last_block_id_in_db().unwrap().unwrap();
-        let last_block = dbio.get_block(last_id).unwrap().unwrap();
-
-        let prev_hash = last_block.header.hash;
-        let transfer_tx =
-            common::test_utils::create_transaction_native_token_transfer(from, 3, to, 1, &sign_key);
-
-        let block = produce_dummy_block(4, Some(prev_hash), vec![transfer_tx]);
-        block_res.push(block.clone());
-        dbio.put_block(&block, [4; 32]).unwrap();
-
-        let block_hashes_mem: Vec<[u8; 32]> =
-            block_res.into_iter().map(|bl| bl.header.hash.0).collect();
-
-        // Get blocks before ID 5 (i.e., starting from 4 going backwards), limit 4
-        // This should return blocks 4, 3, 2, 1 in descending order
-        let mut batch_res = dbio.get_block_batch(Some(5), 4).unwrap();
-        batch_res.reverse(); // Reverse to match ascending order for comparison
-
-        let block_hashes_db: Vec<[u8; 32]> =
-            batch_res.into_iter().map(|bl| bl.header.hash.0).collect();
-
-        assert_eq!(block_hashes_mem, block_hashes_db);
-
-        let block_hashes_mem_limited = &block_hashes_mem[1..];
-
-        // Get blocks before ID 5, limit 3
-        // This should return blocks 4, 3, 2 in descending order
-        let mut batch_res_limited = dbio.get_block_batch(Some(5), 3).unwrap();
-        batch_res_limited.reverse(); // Reverse to match ascending order for comparison
-
-        let block_hashes_db_limited: Vec<[u8; 32]> = batch_res_limited
-            .into_iter()
-            .map(|bl| bl.header.hash.0)
-            .collect();
-
-        assert_eq!(block_hashes_mem_limited, block_hashes_db_limited.as_slice());
-
-        let block_batch_seq = dbio.get_block_batch_seq(1..=5).unwrap();
-        let block_batch_ids = block_batch_seq
-            .into_iter()
-            .map(|block| block.header.block_id)
-            .collect::<Vec<_>>();
-
-        assert_eq!(block_batch_ids, vec![1, 2, 3, 4]);
-    }
-
-    #[test]
-    fn account_map() {
-        let temp_dir = tempdir().unwrap();
-        let temdir_path = temp_dir.path();
-
-        let dbio = RocksDBIO::open_or_create(temdir_path, &initial_state()).unwrap();
-
-        let from = acc1();
-        let to = acc2();
-        let sign_key = acc1_sign_key();
-
-        let mut tx_hash_res = vec![];
-
-        let transfer_tx1 =
-            common::test_utils::create_transaction_native_token_transfer(from, 0, to, 1, &sign_key);
-        let transfer_tx2 =
-            common::test_utils::create_transaction_native_token_transfer(from, 1, to, 1, &sign_key);
-        tx_hash_res.push(transfer_tx1.hash().0);
-        tx_hash_res.push(transfer_tx2.hash().0);
-
-        let block = produce_dummy_block(1, None, vec![transfer_tx1, transfer_tx2]);
-
-        dbio.put_block(&block, [1; 32]).unwrap();
-
-        let last_id = dbio.get_meta_last_block_id_in_db().unwrap().unwrap();
-        let last_block = dbio.get_block(last_id).unwrap().unwrap();
-
-        let prev_hash = last_block.header.hash;
-        let transfer_tx1 =
-            common::test_utils::create_transaction_native_token_transfer(from, 2, to, 1, &sign_key);
-        let transfer_tx2 =
-            common::test_utils::create_transaction_native_token_transfer(from, 3, to, 1, &sign_key);
-        tx_hash_res.push(transfer_tx1.hash().0);
-        tx_hash_res.push(transfer_tx2.hash().0);
-
-        let block = produce_dummy_block(2, Some(prev_hash), vec![transfer_tx1, transfer_tx2]);
-
-        dbio.put_block(&block, [2; 32]).unwrap();
-
-        let last_id = dbio.get_meta_last_block_id_in_db().unwrap().unwrap();
-        let last_block = dbio.get_block(last_id).unwrap().unwrap();
-
-        let prev_hash = last_block.header.hash;
-        let transfer_tx1 =
-            common::test_utils::create_transaction_native_token_transfer(from, 4, to, 1, &sign_key);
-        let transfer_tx2 =
-            common::test_utils::create_transaction_native_token_transfer(from, 5, to, 1, &sign_key);
-        tx_hash_res.push(transfer_tx1.hash().0);
-        tx_hash_res.push(transfer_tx2.hash().0);
-
-        let block = produce_dummy_block(3, Some(prev_hash), vec![transfer_tx1, transfer_tx2]);
-
-        dbio.put_block(&block, [3; 32]).unwrap();
-
-        let last_id = dbio.get_meta_last_block_id_in_db().unwrap().unwrap();
-        let last_block = dbio.get_block(last_id).unwrap().unwrap();
-
-        let prev_hash = last_block.header.hash;
-        let transfer_tx =
-            common::test_utils::create_transaction_native_token_transfer(from, 6, to, 1, &sign_key);
-        tx_hash_res.push(transfer_tx.hash().0);
-
-        let block = produce_dummy_block(4, Some(prev_hash), vec![transfer_tx]);
-
-        dbio.put_block(&block, [4; 32]).unwrap();
-
-        let acc1_tx = dbio.get_acc_transactions(*acc1().value(), 0, 7).unwrap();
-        let acc1_tx_hashes: Vec<[u8; 32]> = acc1_tx.into_iter().map(|tx| tx.hash().0).collect();
-
-        assert_eq!(acc1_tx_hashes, tx_hash_res);
-
-        let acc1_tx_limited = dbio.get_acc_transactions(*acc1().value(), 1, 4).unwrap();
-        let acc1_tx_limited_hashes: Vec<[u8; 32]> =
-            acc1_tx_limited.into_iter().map(|tx| tx.hash().0).collect();
-
-        assert_eq!(acc1_tx_limited_hashes.as_slice(), &tx_hash_res[1..5]);
-    }
-}
+mod tests;

lez/storage/src/indexer/tests.rs

@@ -0,0 +1,433 @@
+use common::test_utils::produce_dummy_block;
+use lee::{Account, AccountId, PublicKey};
+use tempfile::tempdir;
+
+use super::*;
+
+fn genesis_block() -> Block {
+    produce_dummy_block(1, None, vec![])
+}
+
+fn acc1_sign_key() -> lee::PrivateKey {
+    lee::PrivateKey::try_new([1; 32]).unwrap()
+}
+
+fn acc2_sign_key() -> lee::PrivateKey {
+    lee::PrivateKey::try_new([2; 32]).unwrap()
+}
+
+fn acc1() -> AccountId {
+    AccountId::from(&PublicKey::new_from_private_key(&acc1_sign_key()))
+}
+
+fn acc2() -> AccountId {
+    AccountId::from(&PublicKey::new_from_private_key(&acc2_sign_key()))
+}
+
+fn initial_state() -> lee::V03State {
+    let mut public_accounts = [(acc1(), 10000), (acc2(), 20000)]
+        .into_iter()
+        .map(|(id, balance)| {
+            (
+                id,
+                Account {
+                    program_owner: programs::authenticated_transfer().id(),
+                    balance,
+                    ..Account::default()
+                },
+            )
+        })
+        .collect::<Vec<_>>();
+    for clock_id in system_accounts::clock_account_ids() {
+        public_accounts.push((clock_id, system_accounts::clock_account()));
+    }
+
+    lee::V03State::new()
+        .with_public_accounts(public_accounts)
+        .with_programs([programs::authenticated_transfer(), programs::clock()])
+}
+
+#[test]
+fn start_db() {
+    let temp_dir = tempdir().unwrap();
+    let temdir_path = temp_dir.path();
+
+    let dbio = RocksDBIO::open_or_create(temdir_path, &initial_state()).unwrap();
+
+    let last_id = dbio.get_meta_last_block_id_in_db().unwrap();
+    let first_id = dbio.get_meta_first_block_id_in_db().unwrap();
+    let is_first_set = dbio.get_meta_is_first_block_set().unwrap();
+    let last_observed_l1_header = dbio.get_meta_last_observed_l1_lib_header_in_db().unwrap();
+    let last_br_id = dbio.get_meta_last_breakpoint_id().unwrap();
+    let last_block = dbio.get_block(1).unwrap();
+    let breakpoint = dbio.get_breakpoint(0).unwrap();
+    let final_state = dbio.final_state().unwrap();
+
+    assert_eq!(last_id, None);
+    assert_eq!(first_id, None);
+    assert_eq!(last_observed_l1_header, None);
+    assert!(!is_first_set);
+    assert_eq!(last_br_id, Some(0)); // TODO: Will be None after we remove hardcoded testnet state
+    assert!(last_block.is_none());
+    assert_eq!(
+        breakpoint.get_account_by_id(acc1()),
+        final_state.get_account_by_id(acc1())
+    );
+    assert_eq!(
+        breakpoint.get_account_by_id(acc2()),
+        final_state.get_account_by_id(acc2())
+    );
+}
+
+#[test]
+fn one_block_insertion() {
+    let temp_dir = tempdir().unwrap();
+    let temdir_path = temp_dir.path();
+
+    let dbio = RocksDBIO::open_or_create(temdir_path, &initial_state()).unwrap();
+
+    let genesis_block = genesis_block();
+    dbio.put_block(&genesis_block, [0; 32]).unwrap();
+
+    let prev_hash = genesis_block.header.hash;
+    let from = acc1();
+    let to = acc2();
+    let sign_key = acc1_sign_key();
+
+    let transfer_tx =
+        common::test_utils::create_transaction_native_token_transfer(from, 0, to, 1, &sign_key);
+    let block = produce_dummy_block(2, Some(prev_hash), vec![transfer_tx]);
+
+    dbio.put_block(&block, [1; 32]).unwrap();
+
+    let last_id = dbio.get_meta_last_block_id_in_db().unwrap().unwrap();
+    let first_id = dbio.get_meta_first_block_id_in_db().unwrap();
+    let last_observed_l1_header = dbio
+        .get_meta_last_observed_l1_lib_header_in_db()
+        .unwrap()
+        .unwrap();
+    let is_first_set = dbio.get_meta_is_first_block_set().unwrap();
+    let last_br_id = dbio.get_meta_last_breakpoint_id().unwrap();
+    let last_block = dbio.get_block(last_id).unwrap().unwrap();
+    let breakpoint = dbio.get_breakpoint(0).unwrap();
+    let final_state = dbio.final_state().unwrap();
+
+    assert_eq!(last_id, 2);
+    assert_eq!(first_id, Some(1));
+    assert_eq!(last_observed_l1_header, [1; 32]);
+    assert!(is_first_set);
+    assert_eq!(last_br_id, Some(0));
+    assert_eq!(last_block.header.hash, block.header.hash);
+    assert_eq!(
+        breakpoint.get_account_by_id(acc1()).balance
+            - final_state.get_account_by_id(acc1()).balance,
+        1
+    );
+    assert_eq!(
+        final_state.get_account_by_id(acc2()).balance
+            - breakpoint.get_account_by_id(acc2()).balance,
+        1
+    );
+}
+
+#[test]
+fn new_breakpoint() {
+    let temp_dir = tempdir().unwrap();
+    let temdir_path = temp_dir.path();
+
+    let dbio = RocksDBIO::open_or_create(temdir_path, &initial_state()).unwrap();
+
+    let from = acc1();
+    let to = acc2();
+    let sign_key = acc1_sign_key();
+
+    for i in 1..=BREAKPOINT_INTERVAL + 1 {
+        let prev_hash = dbio.get_meta_last_block_id_in_db().unwrap().map(|last_id| {
+            let last_block = dbio.get_block(last_id).unwrap().unwrap();
+            last_block.header.hash
+        });
+
+        let transfer_tx = common::test_utils::create_transaction_native_token_transfer(
+            from,
+            (i - 1).into(),
+            to,
+            1,
+            &sign_key,
+        );
+        let block = produce_dummy_block(i.into(), prev_hash, vec![transfer_tx]);
+        dbio.put_block(&block, [i; 32]).unwrap();
+    }
+
+    let last_id = dbio.get_meta_last_block_id_in_db().unwrap().unwrap();
+    let first_id = dbio.get_meta_first_block_id_in_db().unwrap();
+    let is_first_set = dbio.get_meta_is_first_block_set().unwrap();
+    let last_br_id = dbio.get_meta_last_breakpoint_id().unwrap();
+    let last_block = dbio.get_block(last_id).unwrap().unwrap();
+    let prev_breakpoint = dbio.get_breakpoint(0).unwrap();
+    let breakpoint = dbio.get_breakpoint(1).unwrap();
+    let final_state = dbio.final_state().unwrap();
+
+    assert_eq!(last_id, 101);
+    assert_eq!(first_id, Some(1));
+    assert!(is_first_set);
+    assert_eq!(last_br_id, Some(1));
+    assert_ne!(last_block.header.hash, genesis_block().header.hash);
+    assert_eq!(
+        prev_breakpoint.get_account_by_id(acc1()).balance
+            - final_state.get_account_by_id(acc1()).balance,
+        101
+    );
+    assert_eq!(
+        final_state.get_account_by_id(acc2()).balance
+            - prev_breakpoint.get_account_by_id(acc2()).balance,
+        101
+    );
+    assert_eq!(
+        breakpoint.get_account_by_id(acc1()).balance
+            - final_state.get_account_by_id(acc1()).balance,
+        1
+    );
+    assert_eq!(
+        final_state.get_account_by_id(acc2()).balance
+            - breakpoint.get_account_by_id(acc2()).balance,
+        1
+    );
+}
+
+#[test]
+fn simple_maps() {
+    let temp_dir = tempdir().unwrap();
+    let temdir_path = temp_dir.path();
+
+    let dbio = RocksDBIO::open_or_create(temdir_path, &initial_state()).unwrap();
+
+    let from = acc1();
+    let to = acc2();
+    let sign_key = acc1_sign_key();
+
+    let transfer_tx =
+        common::test_utils::create_transaction_native_token_transfer(from, 0, to, 1, &sign_key);
+    let block = produce_dummy_block(1, None, vec![transfer_tx]);
+
+    let control_hash1 = block.header.hash;
+
+    dbio.put_block(&block, [1; 32]).unwrap();
+
+    let last_id = dbio.get_meta_last_block_id_in_db().unwrap().unwrap();
+    let last_block = dbio.get_block(last_id).unwrap().unwrap();
+
+    let prev_hash = last_block.header.hash;
+    let transfer_tx =
+        common::test_utils::create_transaction_native_token_transfer(from, 1, to, 1, &sign_key);
+    let block = produce_dummy_block(2, Some(prev_hash), vec![transfer_tx]);
+
+    let control_hash2 = block.header.hash;
+
+    dbio.put_block(&block, [2; 32]).unwrap();
+
+    let last_id = dbio.get_meta_last_block_id_in_db().unwrap().unwrap();
+    let last_block = dbio.get_block(last_id).unwrap().unwrap();
+
+    let prev_hash = last_block.header.hash;
+    let transfer_tx =
+        common::test_utils::create_transaction_native_token_transfer(from, 2, to, 1, &sign_key);
+
+    let control_tx_hash1 = transfer_tx.hash();
+
+    let block = produce_dummy_block(3, Some(prev_hash), vec![transfer_tx]);
+    dbio.put_block(&block, [3; 32]).unwrap();
+
+    let last_id = dbio.get_meta_last_block_id_in_db().unwrap().unwrap();
+    let last_block = dbio.get_block(last_id).unwrap().unwrap();
+
+    let prev_hash = last_block.header.hash;
+    let transfer_tx =
+        common::test_utils::create_transaction_native_token_transfer(from, 3, to, 1, &sign_key);
+
+    let control_tx_hash2 = transfer_tx.hash();
+
+    let block = produce_dummy_block(4, Some(prev_hash), vec![transfer_tx]);
+    dbio.put_block(&block, [4; 32]).unwrap();
+
+    let control_block_id1 = dbio.get_block_id_by_hash(control_hash1.0).unwrap().unwrap();
+    let control_block_id2 = dbio.get_block_id_by_hash(control_hash2.0).unwrap().unwrap();
+    let control_block_id3 = dbio
+        .get_block_id_by_tx_hash(control_tx_hash1.0)
+        .unwrap()
+        .unwrap();
+    let control_block_id4 = dbio
+        .get_block_id_by_tx_hash(control_tx_hash2.0)
+        .unwrap()
+        .unwrap();
+
+    assert_eq!(control_block_id1, 1);
+    assert_eq!(control_block_id2, 2);
+    assert_eq!(control_block_id3, 3);
+    assert_eq!(control_block_id4, 4);
+}
+
+#[test]
+fn block_batch() {
+    let temp_dir = tempdir().unwrap();
+    let temdir_path = temp_dir.path();
+
+    let mut block_res = vec![];
+
+    let dbio = RocksDBIO::open_or_create(temdir_path, &initial_state()).unwrap();
+
+    let from = acc1();
+    let to = acc2();
+    let sign_key = acc1_sign_key();
+
+    let transfer_tx =
+        common::test_utils::create_transaction_native_token_transfer(from, 0, to, 1, &sign_key);
+    let block = produce_dummy_block(1, None, vec![transfer_tx]);
+
+    block_res.push(block.clone());
+    dbio.put_block(&block, [1; 32]).unwrap();
+
+    let last_id = dbio.get_meta_last_block_id_in_db().unwrap().unwrap();
+    let last_block = dbio.get_block(last_id).unwrap().unwrap();
+
+    let prev_hash = last_block.header.hash;
+    let transfer_tx =
+        common::test_utils::create_transaction_native_token_transfer(from, 1, to, 1, &sign_key);
+    let block = produce_dummy_block(2, Some(prev_hash), vec![transfer_tx]);
+
+    block_res.push(block.clone());
+    dbio.put_block(&block, [2; 32]).unwrap();
+
+    let last_id = dbio.get_meta_last_block_id_in_db().unwrap().unwrap();
+    let last_block = dbio.get_block(last_id).unwrap().unwrap();
+
+    let prev_hash = last_block.header.hash;
+    let transfer_tx =
+        common::test_utils::create_transaction_native_token_transfer(from, 2, to, 1, &sign_key);
+
+    let block = produce_dummy_block(3, Some(prev_hash), vec![transfer_tx]);
+    block_res.push(block.clone());
+    dbio.put_block(&block, [3; 32]).unwrap();
+
+    let last_id = dbio.get_meta_last_block_id_in_db().unwrap().unwrap();
+    let last_block = dbio.get_block(last_id).unwrap().unwrap();
+
+    let prev_hash = last_block.header.hash;
+    let transfer_tx =
+        common::test_utils::create_transaction_native_token_transfer(from, 3, to, 1, &sign_key);
+
+    let block = produce_dummy_block(4, Some(prev_hash), vec![transfer_tx]);
+    block_res.push(block.clone());
+    dbio.put_block(&block, [4; 32]).unwrap();
+
+    let block_hashes_mem: Vec<[u8; 32]> =
+        block_res.into_iter().map(|bl| bl.header.hash.0).collect();
+
+    // Get blocks before ID 5 (i.e., starting from 4 going backwards), limit 4
+    // This should return blocks 4, 3, 2, 1 in descending order
+    let mut batch_res = dbio.get_block_batch(Some(5), 4).unwrap();
+    batch_res.reverse(); // Reverse to match ascending order for comparison
+
+    let block_hashes_db: Vec<[u8; 32]> = batch_res.into_iter().map(|bl| bl.header.hash.0).collect();
+
+    assert_eq!(block_hashes_mem, block_hashes_db);
+
+    let block_hashes_mem_limited = &block_hashes_mem[1..];
+
+    // Get blocks before ID 5, limit 3
+    // This should return blocks 4, 3, 2 in descending order
+    let mut batch_res_limited = dbio.get_block_batch(Some(5), 3).unwrap();
+    batch_res_limited.reverse(); // Reverse to match ascending order for comparison
+
+    let block_hashes_db_limited: Vec<[u8; 32]> = batch_res_limited
+        .into_iter()
+        .map(|bl| bl.header.hash.0)
+        .collect();
+
+    assert_eq!(block_hashes_mem_limited, block_hashes_db_limited.as_slice());
+
+    let block_batch_seq = dbio.get_block_batch_seq(1..=5).unwrap();
+    let block_batch_ids = block_batch_seq
+        .into_iter()
+        .map(|block| block.header.block_id)
+        .collect::<Vec<_>>();
+
+    assert_eq!(block_batch_ids, vec![1, 2, 3, 4]);
+}
+
+#[test]
+fn account_map() {
+    let temp_dir = tempdir().unwrap();
+    let temdir_path = temp_dir.path();
+
+    let dbio = RocksDBIO::open_or_create(temdir_path, &initial_state()).unwrap();
+
+    let from = acc1();
+    let to = acc2();
+    let sign_key = acc1_sign_key();
+
+    let mut tx_hash_res = vec![];
+
+    let transfer_tx1 =
+        common::test_utils::create_transaction_native_token_transfer(from, 0, to, 1, &sign_key);
+    let transfer_tx2 =
+        common::test_utils::create_transaction_native_token_transfer(from, 1, to, 1, &sign_key);
+    tx_hash_res.push(transfer_tx1.hash().0);
+    tx_hash_res.push(transfer_tx2.hash().0);
+
+    let block = produce_dummy_block(1, None, vec![transfer_tx1, transfer_tx2]);
+
+    dbio.put_block(&block, [1; 32]).unwrap();
+
+    let last_id = dbio.get_meta_last_block_id_in_db().unwrap().unwrap();
+    let last_block = dbio.get_block(last_id).unwrap().unwrap();
+
+    let prev_hash = last_block.header.hash;
+    let transfer_tx1 =
+        common::test_utils::create_transaction_native_token_transfer(from, 2, to, 1, &sign_key);
+    let transfer_tx2 =
+        common::test_utils::create_transaction_native_token_transfer(from, 3, to, 1, &sign_key);
+    tx_hash_res.push(transfer_tx1.hash().0);
+    tx_hash_res.push(transfer_tx2.hash().0);
+
+    let block = produce_dummy_block(2, Some(prev_hash), vec![transfer_tx1, transfer_tx2]);
+
+    dbio.put_block(&block, [2; 32]).unwrap();
+
+    let last_id = dbio.get_meta_last_block_id_in_db().unwrap().unwrap();
+    let last_block = dbio.get_block(last_id).unwrap().unwrap();
+
+    let prev_hash = last_block.header.hash;
+    let transfer_tx1 =
+        common::test_utils::create_transaction_native_token_transfer(from, 4, to, 1, &sign_key);
+    let transfer_tx2 =
+        common::test_utils::create_transaction_native_token_transfer(from, 5, to, 1, &sign_key);
+    tx_hash_res.push(transfer_tx1.hash().0);
+    tx_hash_res.push(transfer_tx2.hash().0);
+
+    let block = produce_dummy_block(3, Some(prev_hash), vec![transfer_tx1, transfer_tx2]);
+
+    dbio.put_block(&block, [3; 32]).unwrap();
+
+    let last_id = dbio.get_meta_last_block_id_in_db().unwrap().unwrap();
+    let last_block = dbio.get_block(last_id).unwrap().unwrap();
+
+    let prev_hash = last_block.header.hash;
+    let transfer_tx =
+        common::test_utils::create_transaction_native_token_transfer(from, 6, to, 1, &sign_key);
+    tx_hash_res.push(transfer_tx.hash().0);
+
+    let block = produce_dummy_block(4, Some(prev_hash), vec![transfer_tx]);
+
+    dbio.put_block(&block, [4; 32]).unwrap();
+
+    let acc1_tx = dbio.get_acc_transactions(*acc1().value(), 0, 7).unwrap();
+    let acc1_tx_hashes: Vec<[u8; 32]> = acc1_tx.into_iter().map(|tx| tx.hash().0).collect();
+
+    assert_eq!(acc1_tx_hashes, tx_hash_res);
+
+    let acc1_tx_limited = dbio.get_acc_transactions(*acc1().value(), 1, 4).unwrap();
+    let acc1_tx_limited_hashes: Vec<[u8; 32]> =
+        acc1_tx_limited.into_iter().map(|tx| tx.hash().0).collect();
+
+    assert_eq!(acc1_tx_limited_hashes.as_slice(), &tx_hash_res[1..5]);
+}

lez/storage/src/indexer/write_atomic.rs

@@ -52,44 +52,8 @@ impl RocksDBIO {
         acc_id: [u8; 32],
         tx_hashes: &[[u8; 32]],
     ) -> DbResult<()> {
-        let acc_num_tx = self.get_acc_meta_num_tx(acc_id)?.unwrap_or(0);
-        let cf_att = self.account_id_to_tx_hash_column();
         let mut write_batch = WriteBatch::new();
-
-        for (tx_id, tx_hash) in tx_hashes.iter().enumerate() {
-            let put_id = acc_num_tx
-                .checked_add(tx_id.try_into().expect("Must fit into u64"))
-                .expect("Tx count should be lesser that u64::MAX");
-
-            let mut prefix = borsh::to_vec(&acc_id).map_err(|berr| {
-                DbError::borsh_cast_message(berr, Some("Failed to serialize account id".to_owned()))
-            })?;
-            let suffix = borsh::to_vec(&put_id).map_err(|berr| {
-                DbError::borsh_cast_message(berr, Some("Failed to serialize tx id".to_owned()))
-            })?;
-
-            prefix.extend_from_slice(&suffix);
-
-            write_batch.put_cf(
-                &cf_att,
-                prefix,
-                borsh::to_vec(tx_hash).map_err(|berr| {
-                    DbError::borsh_cast_message(
-                        berr,
-                        Some("Failed to serialize tx hash".to_owned()),
-                    )
-                })?,
-            );
-        }
-
-        self.update_acc_meta_batch(
-            acc_id,
-            acc_num_tx
-                .checked_add(tx_hashes.len().try_into().expect("Must fit into u64"))
-                .expect("Tx count should be lesser that u64::MAX"),
-            &mut write_batch,
-        )?;
-
+        self.put_account_transactions_dependant(acc_id, tx_hashes, &mut write_batch)?;
         self.db.write(write_batch).map_err(|rerr| {
             DbError::rocksdb_cast_message(rerr, Some("Failed to write batch".to_owned()))
         })

lez/wallet-ffi/src/label.rs

@@ -0,0 +1,316 @@
+use std::{
+    ffi::{c_char, CString},
+    str::FromStr as _,
+};
+
+use crate::{
+    c_str_to_string,
+    error::{print_error, WalletFfiError},
+    wallet::get_wallet,
+    FfiAccountIdWithPrivacy, WalletHandle,
+};
+
+#[repr(C)]
+pub struct LabelAvailability {
+    pub is_available: bool,
+    pub error: WalletFfiError,
+}
+
+impl LabelAvailability {
+    #[must_use]
+    pub const fn availability(is_available: bool) -> Self {
+        Self {
+            is_available,
+            error: WalletFfiError::Success,
+        }
+    }
+
+    #[must_use]
+    pub const fn error(error: WalletFfiError) -> Self {
+        Self {
+            is_available: false,
+            error,
+        }
+    }
+}
+
+#[repr(C)]
+#[derive(Debug, Clone, Copy)]
+pub struct AccountIdResolvedFromLabel {
+    pub account_id: FfiAccountIdWithPrivacy,
+    pub error: WalletFfiError,
+}
+
+impl AccountIdResolvedFromLabel {
+    #[must_use]
+    pub const fn account_id(account_id: FfiAccountIdWithPrivacy) -> Self {
+        Self {
+            account_id,
+            error: WalletFfiError::Success,
+        }
+    }
+
+    #[must_use]
+    pub fn error(error: WalletFfiError) -> Self {
+        Self {
+            account_id: FfiAccountIdWithPrivacy::default(),
+            error,
+        }
+    }
+}
+
+#[repr(C)]
+pub struct LabelList {
+    pub labels_data: *mut *const c_char,
+    pub labels_size: usize,
+    pub error: WalletFfiError,
+}
+
+impl LabelList {
+    #[must_use]
+    pub fn from_labels(labels: Vec<*const c_char>) -> Self {
+        let labels_size = labels.len();
+        let boxed_slice = labels.into_boxed_slice();
+        let labels_data = Box::into_raw(boxed_slice).cast::<*const c_char>();
+
+        Self {
+            labels_data,
+            labels_size,
+            error: WalletFfiError::Success,
+        }
+    }
+
+    #[must_use]
+    pub const fn error(error: WalletFfiError) -> Self {
+        Self {
+            labels_data: std::ptr::null_mut(),
+            labels_size: 0,
+            error,
+        }
+    }
+}
+
+/// Check if label is available.
+///
+/// # Parameters
+/// - `handle`: Valid wallet handle
+/// - `label`: Input null terminated C string for a label
+///
+/// # Returns
+/// - `LabelAvailability` struct
+///
+/// # Safety
+/// - `handle` must be a valid wallet handle from `wallet_ffi_create_new` or `wallet_ffi_open`
+/// - `label` must be a valid pointer to a null-terminated C string
+#[no_mangle]
+pub unsafe extern "C" fn wallet_ffi_check_label_available(
+    handle: *mut WalletHandle,
+    label: *const c_char,
+) -> LabelAvailability {
+    let wrapper = match get_wallet(handle) {
+        Ok(w) => w,
+        Err(e) => return LabelAvailability::error(e),
+    };
+
+    let label = match c_str_to_string(label, "label") {
+        Ok(value) => value,
+        Err(e) => return LabelAvailability::error(e),
+    };
+
+    let wallet = match wrapper.core.lock() {
+        Ok(w) => w,
+        Err(e) => {
+            print_error(format!("Failed to lock wallet: {e}"));
+            return LabelAvailability::error(WalletFfiError::InternalError);
+        }
+    };
+
+    let is_available = wallet
+        .storage()
+        .check_label_availability(&label.into())
+        .is_ok();
+
+    LabelAvailability::availability(is_available)
+}
+
+/// Add new label.
+///
+/// # Parameters
+/// - `handle`: Valid wallet handle
+/// - `label`: Input null terminated C string for a label
+/// - `account_id_with_privacy`: The account ID (32 bytes) and its privacy.
+///
+/// # Returns
+/// - `Success` on successful query
+/// - Error code on failure
+///
+/// # Safety
+/// - `handle` must be a valid wallet handle from `wallet_ffi_create_new` or `wallet_ffi_open`
+/// - `label` must be a valid pointer to a null-terminated C string
+#[no_mangle]
+pub unsafe extern "C" fn wallet_ffi_add_label(
+    handle: *mut WalletHandle,
+    label: *const c_char,
+    account_id_with_privacy: FfiAccountIdWithPrivacy,
+) -> WalletFfiError {
+    let wrapper = match get_wallet(handle) {
+        Ok(w) => w,
+        Err(e) => return e,
+    };
+
+    let label = match c_str_to_string(label, "label") {
+        Ok(value) => value,
+        Err(e) => return e,
+    };
+
+    let mut wallet = match wrapper.core.lock() {
+        Ok(w) => w,
+        Err(e) => {
+            print_error(format!("Failed to lock wallet: {e}"));
+            return WalletFfiError::InternalError;
+        }
+    };
+
+    match wallet
+        .storage_mut()
+        .add_label(label.into(), account_id_with_privacy.into())
+    {
+        Ok(()) => WalletFfiError::Success,
+        Err(err) => {
+            print_error(format!("Failed to add label : {err}"));
+            WalletFfiError::InternalError
+        }
+    }
+}
+
+/// Resolve a label.
+///
+/// # Parameters
+/// - `handle`: Valid wallet handle
+/// - `label`: Input null terminated C string for a label
+///
+/// # Returns
+/// - `AccountIdResolvedFromLabel` struct
+///
+/// # Safety
+/// - `handle` must be a valid wallet handle from `wallet_ffi_create_new` or `wallet_ffi_open`
+/// - `label` must be a valid pointer to a null-terminated C string
+#[no_mangle]
+pub unsafe extern "C" fn wallet_ffi_resolve_label(
+    handle: *mut WalletHandle,
+    label: *const c_char,
+) -> AccountIdResolvedFromLabel {
+    let wrapper = match get_wallet(handle) {
+        Ok(w) => w,
+        Err(e) => return AccountIdResolvedFromLabel::error(e),
+    };
+
+    let label = match c_str_to_string(label, "label") {
+        Ok(value) => value,
+        Err(e) => return AccountIdResolvedFromLabel::error(e),
+    };
+
+    let mut wallet = match wrapper.core.lock() {
+        Ok(w) => w,
+        Err(e) => {
+            print_error(format!("Failed to lock wallet: {e}"));
+            return AccountIdResolvedFromLabel::error(WalletFfiError::InternalError);
+        }
+    };
+
+    wallet
+        .storage_mut()
+        .resolve_label(&label.into())
+        .map_or_else(
+            || {
+                print_error("Failed to resolve label");
+                AccountIdResolvedFromLabel::error(WalletFfiError::InternalError)
+            },
+            |acc_id| AccountIdResolvedFromLabel::account_id(acc_id.into()),
+        )
+}
+
+/// Get all labels for account.
+///
+/// # Parameters
+/// - `handle`: Valid wallet handle
+/// - `account_id_with_privacy`: The account ID (32 bytes) and its privacy.
+///
+/// # Returns
+/// - `LabelList` struct
+///
+/// # Safety
+/// - `handle` must be a valid wallet handle from `wallet_ffi_create_new` or `wallet_ffi_open`
+#[no_mangle]
+pub unsafe extern "C" fn wallet_ffi_get_all_labels_for_account(
+    handle: *mut WalletHandle,
+    account_id_with_privacy: FfiAccountIdWithPrivacy,
+) -> LabelList {
+    let wrapper = match get_wallet(handle) {
+        Ok(w) => w,
+        Err(e) => return LabelList::error(e),
+    };
+
+    let wallet = match wrapper.core.lock() {
+        Ok(w) => w,
+        Err(e) => {
+            print_error(format!("Failed to lock wallet: {e}"));
+            return LabelList::error(WalletFfiError::InternalError);
+        }
+    };
+
+    let mut labels = vec![];
+
+    for label in wallet
+        .storage()
+        .labels_for_account(account_id_with_privacy.into())
+    {
+        let Ok(label_c) = CString::from_str(label.as_ref()) else {
+            print_error(format!("Failed to cast label into C string: {label}"));
+            return LabelList::error(WalletFfiError::InternalError);
+        };
+
+        let label_raw = label_c.into_raw().cast_const();
+
+        labels.push(label_raw);
+    }
+
+    LabelList::from_labels(labels)
+}
+
+/// Free label list.
+///
+/// # Parameters
+/// - `label_list`: Input list of labels
+///
+/// # Returns
+/// - `Success` on successful query
+/// - Error code on failure
+///
+/// # Safety
+/// - `label_list` must be a valid pointer to `LabelList`, received from
+///   `wallet_ffi_get_all_labels_for_account`
+#[no_mangle]
+pub unsafe extern "C" fn wallet_ffi_free_label_list(label_list: *mut LabelList) -> WalletFfiError {
+    if label_list.is_null() {
+        return WalletFfiError::NullPointer;
+    }
+
+    let labels_raw = unsafe { &*label_list };
+
+    if !labels_raw.labels_data.is_null() && labels_raw.labels_size > 0 {
+        let labels_slice =
+            std::slice::from_raw_parts_mut(labels_raw.labels_data, labels_raw.labels_size);
+
+        for label_ptr in labels_slice.iter() {
+            if !(*label_ptr).is_null() {
+                drop(CString::from_raw((*label_ptr).cast_mut()));
+            }
+        }
+
+        let boxed_slice = Box::from_raw(std::ptr::from_mut::<[*const c_char]>(labels_slice));
+        drop(boxed_slice);
+    }
+
+    WalletFfiError::Success
+}

lez/wallet-ffi/src/lib.rs

@@ -46,6 +46,7 @@ pub mod bridge;
 pub mod error;
 pub mod generic_transaction;
 pub mod keys;
+pub mod label;
 pub mod pinata;
 pub mod program_deployment;
 pub mod sync;