diff --git a/nssa/core/src/program.rs b/nssa/core/src/program.rs index c79a841..8874773 100644 --- a/nssa/core/src/program.rs +++ b/nssa/core/src/program.rs @@ -24,16 +24,26 @@ pub struct ChainedCall { #[cfg_attr(any(feature = "host", test), derive(Debug, PartialEq, Eq))] pub struct AccountPostState { pub account: Account, - pub claim: bool, + claim: bool, } -impl From for AccountPostState { - fn from(account: Account) -> Self { - AccountPostState { +impl AccountPostState { + pub fn new(account: Account) -> Self { + Self { account, claim: false, } } + pub fn new_claimed(account: Account) -> Self { + Self { + account, + claim: true, + } + } + + pub fn requires_claim(&self) -> bool { + self.claim + } } impl AccountPostState { @@ -153,7 +163,7 @@ mod tests { use super::*; #[test] - fn test_account_post_state_from_account_constructor() { + fn test_post_state_new_without_claim_constructor() { let account = Account { program_owner: [1, 2, 3, 4, 5, 6, 7, 8], balance: 1337, @@ -161,9 +171,25 @@ mod tests { nonce: 10, }; - let account_post_state: AccountPostState = account.clone().into(); + let account_post_state = AccountPostState::new_claimed(account.clone()); assert_eq!(account, account_post_state.account); - assert!(!account_post_state.claim); + assert!(account_post_state.requires_claim()); } + + #[test] + fn test_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], + nonce: 10, + }; + + let account_post_state = AccountPostState::new(account.clone()); + + assert_eq!(account, account_post_state.account); + assert!(!account_post_state.requires_claim()); + } + } diff --git a/nssa/program_methods/guest/src/bin/authenticated_transfer.rs b/nssa/program_methods/guest/src/bin/authenticated_transfer.rs index 14aded4..f711f20 100644 --- a/nssa/program_methods/guest/src/bin/authenticated_transfer.rs +++ b/nssa/program_methods/guest/src/bin/authenticated_transfer.rs @@ -1,12 +1,14 @@ use nssa_core::{ account::{Account, AccountWithMetadata}, - program::{AccountPostState, ProgramInput, read_nssa_inputs, write_nssa_outputs}, + program::{ + AccountPostState, DEFAULT_PROGRAM_ID, ProgramInput, read_nssa_inputs, write_nssa_outputs, + }, }; /// Initializes a default account under the ownership of this program. /// This is achieved by a noop. fn initialize_account(pre_state: AccountWithMetadata) { - let account_to_claim: AccountPostState = pre_state.account.clone().into(); + let account_to_claim = AccountPostState::new_claimed(pre_state.account.clone()); let is_authorized = pre_state.is_authorized; // Continue only if the account to claim has default values @@ -36,15 +38,27 @@ fn transfer(sender: AccountWithMetadata, recipient: AccountWithMetadata, balance } // Create accounts post states, with updated balances - let mut sender_post = sender.account.clone(); - let mut recipient_post = recipient.account.clone(); - sender_post.balance -= balance_to_move; - recipient_post.balance += balance_to_move; + let sender_post: AccountPostState = { + // Modify sender's balance + let mut sender_post_account = sender.account.clone(); + sender_post_account.balance -= balance_to_move; + AccountPostState::new(sender_post_account) + }; - write_nssa_outputs( - vec![sender, recipient], - vec![sender_post.into(), recipient_post.into()], - ); + let recipient_post = { + // Modify recipient's balance + let mut recipient_post_account = recipient.account.clone(); + recipient_post_account.balance += balance_to_move; + + // Claim recipient account if it has default program owner + if recipient_post_account.program_owner == DEFAULT_PROGRAM_ID { + AccountPostState::new_claimed(recipient_post_account) + } else { + AccountPostState::new(recipient_post_account) + } + }; + + write_nssa_outputs(vec![sender, recipient], vec![sender_post, recipient_post]); } /// A transfer of balance program. diff --git a/nssa/program_methods/guest/src/bin/pinata.rs b/nssa/program_methods/guest/src/bin/pinata.rs index 9337ab7..50aac7b 100644 --- a/nssa/program_methods/guest/src/bin/pinata.rs +++ b/nssa/program_methods/guest/src/bin/pinata.rs @@ -1,4 +1,4 @@ -use nssa_core::program::{ProgramInput, read_nssa_inputs, write_nssa_outputs}; +use nssa_core::program::{AccountPostState, ProgramInput, read_nssa_inputs, write_nssa_outputs}; use risc0_zkvm::sha::{Impl, Sha256}; const PRIZE: u128 = 150; @@ -66,5 +66,11 @@ fn main() { pinata_post.data = data.next_data().to_vec(); winner_post.balance += PRIZE; - write_nssa_outputs(vec![pinata, winner], vec![pinata_post.into(), winner_post.into()]); + write_nssa_outputs( + vec![pinata, winner], + vec![ + AccountPostState::new(pinata_post), + AccountPostState::new(winner_post), + ], + ); } diff --git a/nssa/program_methods/guest/src/bin/token.rs b/nssa/program_methods/guest/src/bin/token.rs index 1e9cc80..ce4558a 100644 --- a/nssa/program_methods/guest/src/bin/token.rs +++ b/nssa/program_methods/guest/src/bin/token.rs @@ -1,6 +1,8 @@ use nssa_core::{ account::{Account, AccountId, AccountWithMetadata, Data}, - program::{read_nssa_inputs, write_nssa_outputs, AccountPostState, ProgramInput}, + program::{ + AccountPostState, DEFAULT_PROGRAM_ID, ProgramInput, read_nssa_inputs, write_nssa_outputs, + }, }; // The token program has three functions: @@ -148,15 +150,22 @@ fn transfer(pre_states: &[AccountWithMetadata], balance_to_move: u128) -> Vec Vec { @@ -220,10 +232,13 @@ fn initialize_account(pre_states: &[AccountWithMetadata]) -> Vec