From 960bdfbb33e3cb4cbd557a48182b1dc1b9f16c2b Mon Sep 17 00:00:00 2001 From: Marvin Jones Date: Thu, 13 Aug 2026 14:26:34 -0400 Subject: [PATCH] feat(lee): store deployed programs as Account-shaped state, keyed by AccountId Program-as-Account migration, first slice: V03State.programs becomes HashMap instead of HashMap, with the elf held directly in Account.data. The map key is derived from ProgramId via a new 1:1 From for AccountId conversion (both types are exactly 32 bytes) rather than a hash, since ProgramId is already content-derived from the elf. Account.program_owner stays ProgramId-typed everywhere - this only changes how deployed programs are stored and looked up host-side, not the dispatch/authorization model any guest program logic depends on. Dispatch resolves a ChainedCall's program_id by converting to AccountId, fetching the Account, and reconstructing a Program via new_unchecked for execution. DATA_MAX_LENGTH is raised from 100 KiB to 700 KiB to fit real program elfs (observed 375 KB-631 KB) directly in Account.data; noted in its docstring as a rough placeholder pending real transaction/block-size budget analysis. --- lee/state_machine/core/src/account/data.rs | 10 +++- lee/state_machine/core/src/program/mod.rs | 17 ++++++ lee/state_machine/src/state/mod.rs | 45 +++++++++----- lee/state_machine/src/state/tests/genesis.rs | 7 ++- lee/state_machine/src/state/tests/mod.rs | 60 +++++++++---------- .../src/validated_state_diff/mod.rs | 21 ++++++- 6 files changed, 108 insertions(+), 52 deletions(-) diff --git a/lee/state_machine/core/src/account/data.rs b/lee/state_machine/core/src/account/data.rs index 272f01915..867c82746 100644 --- a/lee/state_machine/core/src/account/data.rs +++ b/lee/state_machine/core/src/account/data.rs @@ -4,7 +4,15 @@ use borsh::{BorshDeserialize, BorshSerialize}; use bytesize::ByteSize; use serde::{Deserialize, Serialize}; -pub const DATA_MAX_LENGTH: ByteSize = ByteSize::kib(100); +/// Raised from the original 100 KiB to accommodate program elfs stored directly in +/// `Account.data` under the Program-as-Account migration. +/// +/// Observed elfs currently run 375 KB-520 KB, plus 631 KB for the fixed +/// privacy-preserving circuit itself. This value is a rough placeholder, not a considered +/// protocol constant yet — it still needs to be refined against real transaction/block-size +/// budgets (e.g. `SequencerConfig::max_block_size`, currently 1 MiB) before this is something +/// production traffic should rely on. +pub const DATA_MAX_LENGTH: ByteSize = ByteSize::kib(700); #[derive(Debug, Default, Clone, PartialEq, Eq, Serialize, BorshSerialize)] pub struct Data(Vec); diff --git a/lee/state_machine/core/src/program/mod.rs b/lee/state_machine/core/src/program/mod.rs index 770bcf2db..52bb0c795 100644 --- a/lee/state_machine/core/src/program/mod.rs +++ b/lee/state_machine/core/src/program/mod.rs @@ -14,6 +14,23 @@ pub const DEFAULT_PROGRAM_ID: ProgramId = [0; 8]; pub const MAX_NUMBER_CHAINED_CALLS: usize = 10; pub type ProgramId = [u32; 8]; + +/// Derives the `AccountId` under which a program's data is stored, directly from its +/// `ProgramId`, by reinterpreting the 8 little-endian `u32` words as 32 raw bytes. +/// +/// A 1:1, information-preserving mapping (both types are exactly 32 bytes) rather than a +/// hash — `ProgramId` is already content-derived (RISC0's `image_id`), so no extra domain +/// separation is needed just to use it as a `HashMap` key. +impl From for AccountId { + fn from(program_id: ProgramId) -> Self { + let bytes: Vec = program_id + .iter() + .flat_map(|word| word.to_le_bytes()) + .collect(); + Self::new(bytes.try_into().expect("8 u32 words are exactly 32 bytes")) + } +} + pub type InstructionData = Vec; pub struct ProgramInput { pub self_program_id: ProgramId, diff --git a/lee/state_machine/src/state/mod.rs b/lee/state_machine/src/state/mod.rs index 8b91b6983..cdce2de61 100644 --- a/lee/state_machine/src/state/mod.rs +++ b/lee/state_machine/src/state/mod.rs @@ -4,8 +4,7 @@ use borsh::{BorshDeserialize, BorshSerialize}; use lee_core::{ BlockId, Commitment, CommitmentSetDigest, DUMMY_COMMITMENT, MembershipProof, Nullifier, Timestamp, - account::{Account, AccountId}, - program::ProgramId, + account::{Account, AccountId, Data}, }; use crate::{ @@ -114,7 +113,12 @@ impl BorshDeserialize for NullifierSet { pub struct V03State { public_state: HashMap, private_state: (CommitmentSet, NullifierSet), - programs: HashMap, + /// Deployed programs, stored as `Account`s keyed by `AccountId::from(program_id)` (see that + /// impl's doc comment) rather than by `ProgramId` directly, with the elf held in + /// `Account.data`. Kept as its own map rather than folded into `public_state`: nothing in + /// dispatch/execution reads or writes it, so it isn't part of the account-mutation surface + /// `program_owner`-based authorization governs — this is host-side bookkeeping only. + programs: HashMap, } impl Default for V03State { @@ -190,13 +194,19 @@ impl V03State { #[must_use] pub fn with_programs(mut self, programs: impl IntoIterator) -> Self { for program in programs { - self.insert_program(program); + self.insert_program(&program); } self } - pub(crate) fn insert_program(&mut self, program: Program) { - self.programs.insert(program.id(), program); + pub(crate) fn insert_program(&mut self, program: &Program) { + let account_id = AccountId::from(program.id()); + let account = Account { + data: Data::try_from(program.elf().to_vec()) + .expect("elf must fit under DATA_MAX_LENGTH"), + ..Account::default() + }; + self.programs.insert(account_id, account); } pub fn apply_state_diff(&mut self, diff: ValidatedStateDiff) { @@ -222,7 +232,7 @@ impl V03State { self.private_state.0.extend(&new_commitments); self.private_state.1.extend(&new_nullifiers); if let Some(program) = program { - self.insert_program(program); + self.insert_program(&program); } } @@ -281,7 +291,7 @@ impl V03State { self.private_state.0.get_proof_for(commitment) } - pub(crate) const fn programs(&self) -> &HashMap { + pub(crate) const fn programs(&self) -> &HashMap { &self.programs } @@ -314,11 +324,14 @@ impl V03State { let mut accounts: Vec<(&AccountId, &Account)> = public_state.iter().collect(); accounts.sort_by(|a, b| a.0.as_ref().cmp(b.0.as_ref())); - let mut program_ids: Vec = programs.keys().copied().collect(); - program_ids.sort_unstable(); + // `programs` is `Account`-shaped now, same as `public_state` — reuse the identical + // sort-then-hash-id-plus-encoded-account pattern rather than a bespoke `ProgramId` loop. + let mut program_accounts: Vec<(&AccountId, &Account)> = programs.iter().collect(); + program_accounts.sort_by(|a, b| a.0.as_ref().cmp(b.0.as_ref())); let account_count = u64::try_from(accounts.len()).expect("account count fits in u64"); - let program_count = u64::try_from(program_ids.len()).expect("program count fits in u64"); + let program_count = + u64::try_from(program_accounts.len()).expect("program count fits in u64"); let mut hasher = Sha256::new(); hasher.update(account_count.to_le_bytes()); @@ -330,10 +343,12 @@ impl V03State { hasher.update(&bytes); } hasher.update(program_count.to_le_bytes()); - for id in program_ids { - for word in id { - hasher.update(word.to_le_bytes()); - } + for (id, account) in program_accounts { + hasher.update(id.as_ref()); + let bytes = borsh::to_vec(account).expect("Account is BorshSerialize"); + let len = u64::try_from(bytes.len()).expect("program account encoding fits in u64"); + hasher.update(len.to_le_bytes()); + hasher.update(&bytes); } hasher.update(private_state.0.digest()); diff --git a/lee/state_machine/src/state/tests/genesis.rs b/lee/state_machine/src/state/tests/genesis.rs index f67628eee..4a07b5086 100644 --- a/lee/state_machine/src/state/tests/genesis.rs +++ b/lee/state_machine/src/state/tests/genesis.rs @@ -67,11 +67,12 @@ 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)); + let account_id = lee_core::account::AccountId::from(program_id); + assert!(!state.programs.contains_key(&account_id)); - state.insert_program(program_to_insert); + state.insert_program(&program_to_insert); - assert!(state.programs.contains_key(&program_id)); + assert!(state.programs.contains_key(&account_id)); } #[test] diff --git a/lee/state_machine/src/state/tests/mod.rs b/lee/state_machine/src/state/tests/mod.rs index b1c67b9b5..d6004853d 100644 --- a/lee/state_machine/src/state/tests/mod.rs +++ b/lee/state_machine/src/state/tests/mod.rs @@ -45,36 +45,36 @@ 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::dropped_account()); - 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.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::dropped_account()); + 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 } diff --git a/lee/state_machine/src/validated_state_diff/mod.rs b/lee/state_machine/src/validated_state_diff/mod.rs index 738cc5153..276c8bcec 100644 --- a/lee/state_machine/src/validated_state_diff/mod.rs +++ b/lee/state_machine/src/validated_state_diff/mod.rs @@ -1,4 +1,5 @@ use std::{ + borrow::Cow, collections::{HashMap, HashSet, VecDeque}, hash::Hash, }; @@ -111,10 +112,21 @@ impl ValidatedStateDiff { LeeError::MaxChainedCallsDepthExceeded ); - // Check that the `program_id` corresponds to a deployed program - let Some(program) = state.programs().get(&chained_call.program_id) else { + // Check that the `program_id` corresponds to a deployed program. `programs` is + // keyed by `AccountId::from(program_id)`, not `program_id` itself (see that impl's + // doc comment), and holds the elf as a plain `Account`; reconstruct a `Program` from + // it via `new_unchecked` to execute, skipping a redundant image-id recomputation + // since the id/elf pairing was already validated once, at deployment time. + let Some(program_account) = state + .programs() + .get(&AccountId::from(chained_call.program_id)) + else { return Err(LeeError::InvalidInput("Unknown program".into())); }; + let program = Program::new_unchecked( + chained_call.program_id, + Cow::Owned(program_account.data.to_vec()), + ); debug!( "Program {:?} pre_states: {:?}, instruction_data: {:?}", @@ -444,7 +456,10 @@ impl ValidatedStateDiff { ) -> Result { // TODO: remove clone let program = Program::new(tx.message.bytecode.clone().into())?; - if state.programs().contains_key(&program.id()) { + if state + .programs() + .contains_key(&AccountId::from(program.id())) + { return Err(LeeError::ProgramAlreadyExists); } Ok(Self(StateDiff {