refactor logic

nssa/src/error.rs

@@ -5,9 +5,6 @@ pub enum NssaError {
     #[error("Invalid input: {0}")]
     InvalidInput(String),
 
-    #[error("Operation failed")]
-    OperationFailed,
-
     #[error("Risc0 error: {0}")]
     ProgramExecutionFailed(String),
 

nssa/src/program.rs

@@ -146,7 +146,7 @@ impl Program {
         }
     }
 
-    /// A program that mints balance
+    /// A program that burns balance
     pub fn burner() -> Self {
         use test_program_methods::{BURNER_ELF, BURNER_ID};
 

nssa/src/public_transaction/mod.rs

@@ -1,11 +1,13 @@
-use nssa_core::{account::Nonce, program::ProgramId};
+use std::collections::{HashMap, HashSet};
+
+use nssa_core::{
+    account::{Account, AccountWithMetadata},
+    program::validate_execution,
+};
 use serde::{Deserialize, Serialize};
 use sha2::{Digest, digest::FixedOutput};
 
-use crate::{
+use crate::{V01State, address::Address, error::NssaError};
-    address::Address,
-    signature::{PrivateKey, PublicKey, Signature},
-};
 
 mod message;
 mod witness_set;
@@ -49,4 +51,71 @@ impl PublicTransaction {
         hasher.update(&bytes);
         hasher.finalize_fixed().into()
     }
+
+    pub(crate) fn validate_and_compute_post_states(
+        &self,
+        state: &V01State,
+    ) -> Result<HashMap<Address, Account>, NssaError> {
+        let message = self.message();
+        let witness_set = self.witness_set();
+
+        // All addresses must be different
+        if message.addresses.iter().collect::<HashSet<_>>().len() != message.addresses.len() {
+            return Err(NssaError::InvalidInput(
+                "Duplicate addresses found in message".into(),
+            ));
+        }
+
+        if message.nonces.len() != witness_set.signatures_and_public_keys.len() {
+            return Err(NssaError::InvalidInput(
+                "Mismatch between number of nonces and signatures/public keys".into(),
+            ));
+        }
+
+        let mut authorized_addresses = Vec::new();
+        for ((signature, public_key), nonce) in witness_set.iter_signatures().zip(&message.nonces) {
+            // Check the signature is valid
+            if !signature.is_valid_for(message, public_key) {
+                return Err(NssaError::InvalidInput(
+                    "Invalid signature for given message and public key".into(),
+                ));
+            }
+
+            // Check the nonce corresponds to the current nonce on the public state.
+            let address = Address::from_public_key(public_key);
+            let current_nonce = state.get_account_by_address(&address).nonce;
+            if current_nonce != *nonce {
+                return Err(NssaError::InvalidInput("Nonce mismatch".into()));
+            }
+
+            authorized_addresses.push(address);
+        }
+
+        // Build pre_states for execution
+        let pre_states: Vec<_> = message
+            .addresses
+            .iter()
+            .map(|address| AccountWithMetadata {
+                account: state.get_account_by_address(address),
+                is_authorized: authorized_addresses.contains(address),
+            })
+            .collect();
+
+        // Check the `program_id` corresponds to a built-in program
+        // Only allowed program so far is the authenticated transfer program
+        let Some(program) = state.builtin_programs().get(&message.program_id) else {
+            return Err(NssaError::InvalidInput("Unknown program".into()));
+        };
+
+        // // Execute program
+        let post_states = program.execute(&pre_states, message.instruction_data)?;
+
+        // Verify execution corresponds to a well-behaved program.
+        // See the # Programs section for the definition of the `validate_execution` method.
+        if !validate_execution(&pre_states, &post_states, message.program_id) {
+            return Err(NssaError::InvalidProgramBehavior);
+        }
+
+        Ok(message.addresses.iter().cloned().zip(post_states).collect())
+    }
 }

nssa/src/public_transaction/witness_set.rs

@@ -7,14 +7,14 @@ pub struct WitnessSet {
     pub(crate) signatures_and_public_keys: Vec<(Signature, PublicKey)>,
 }
 
-fn serialize_message_to_bytes(message: &Message) -> Vec<u8> {
+fn message_to_bytes(_message: &Message) -> Vec<u8> {
     //TODO: implement
     vec![0, 0]
 }
 
 impl WitnessSet {
     pub fn for_message(message: &Message, private_keys: &[&PrivateKey]) -> Self {
-        let message_bytes = serialize_message_to_bytes(&message);
+        let message_bytes = message_to_bytes(message);
         let signatures_and_public_keys = private_keys
             .iter()
             .map(|&key| (Signature::new(key, &message_bytes), PublicKey::new(key)))

nssa/src/state.rs

@@ -1,11 +1,8 @@
 use crate::{
     address::Address, error::NssaError, program::Program, public_transaction::PublicTransaction,
 };
-use nssa_core::{
+use nssa_core::{account::Account, program::ProgramId};
-    account::{Account, AccountWithMetadata},
+use std::collections::HashMap;
-    program::{ProgramId, validate_execution},
-};
-use std::collections::{HashMap, HashSet};
 
 pub struct V01State {
     public_state: HashMap<Address, Account>,
@@ -47,7 +44,7 @@ impl V01State {
         &mut self,
         tx: &PublicTransaction,
     ) -> Result<(), NssaError> {
-        let state_diff = self.execute_and_verify_public_transaction(tx)?;
+        let state_diff = tx.validate_and_compute_post_states(self)?;
 
         for (address, post) in state_diff.into_iter() {
             let current_account = self.get_account_by_address_mut(address);
@@ -73,71 +70,8 @@ impl V01State {
             .unwrap_or(Account::default())
     }
 
-    fn execute_and_verify_public_transaction(
+    pub(crate) fn builtin_programs(&self) -> &HashMap<ProgramId, Program> {
-        &mut self,
+        &self.builtin_programs
-        tx: &PublicTransaction,
-    ) -> Result<HashMap<Address, Account>, NssaError> {
-        let message = tx.message();
-        let witness_set = tx.witness_set();
-
-        // All addresses must be different
-        if message.addresses.iter().collect::<HashSet<_>>().len() != message.addresses.len() {
-            return Err(NssaError::InvalidInput(
-                "Duplicate addresses found in message".into(),
-            ));
-        }
-
-        if message.nonces.len() != witness_set.signatures_and_public_keys.len() {
-            return Err(NssaError::InvalidInput(
-                "Mismatch between number of nonces and signatures/public keys".into(),
-            ));
-        }
-
-        let mut authorized_addresses = Vec::new();
-        for ((signature, public_key), nonce) in witness_set.iter_signatures().zip(&message.nonces) {
-            // Check the signature is valid
-            if !signature.is_valid_for(message, public_key) {
-                return Err(NssaError::InvalidInput(
-                    "Invalid signature for given message and public key".into(),
-                ));
-            }
-
-            // Check the nonce corresponds to the current nonce on the public state.
-            let address = Address::from_public_key(public_key);
-            let current_nonce = self.get_account_by_address(&address).nonce;
-            if current_nonce != *nonce {
-                return Err(NssaError::InvalidInput("Nonce mismatch".into()));
-            }
-
-            authorized_addresses.push(address);
-        }
-
-        // Build pre_states for execution
-        let pre_states: Vec<_> = message
-            .addresses
-            .iter()
-            .map(|address| AccountWithMetadata {
-                account: self.get_account_by_address(address),
-                is_authorized: authorized_addresses.contains(address),
-            })
-            .collect();
-
-        // Check the `program_id` corresponds to a built-in program
-        // Only allowed program so far is the authenticated transfer program
-        let Some(program) = self.builtin_programs.get(&message.program_id) else {
-            return Err(NssaError::InvalidInput("Unknown program".into()));
-        };
-
-        // // Execute program
-        let post_states = program.execute(&pre_states, message.instruction_data)?;
-
-        // Verify execution corresponds to a well-behaved program.
-        // See the # Programs section for the definition of the `validate_execution` method.
-        if !validate_execution(&pre_states, &post_states, message.program_id) {
-            return Err(NssaError::InvalidProgramBehavior);
-        }
-
-        Ok(message.addresses.iter().cloned().zip(post_states).collect())
     }
 }