improve struct interface

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<Account> 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());
+    }
+
 }

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.

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),
+        ],
+    );
 }

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<Ac
     let sender_post = {
         let mut this = sender.account.clone();
         this.data = sender_holding.into_data();
-        this
+        AccountPostState::new(this)
     };
+
     let recipient_post = {
         let mut this = recipient.account.clone();
         this.data = recipient_holding.into_data();
-        this
+
+        // Claim the recipient account if it has default program owner
+        if this.program_owner == DEFAULT_PROGRAM_ID {
+            AccountPostState::new_claimed(this)
+        } else {
+            AccountPostState::new(this)
+        }
     };
 
-    vec![sender_post.into(), recipient_post.into()]
+    vec![sender_post, recipient_post]
 }
 
 fn new_definition(
@@ -196,7 +205,10 @@ fn new_definition(
     let mut holding_target_account_post = holding_target_account.account.clone();
     holding_target_account_post.data = token_holding.into_data();
 
-    vec![definition_target_account_post.into(), holding_target_account_post.into()]
+    vec![
+        AccountPostState::new_claimed(definition_target_account_post),
+        AccountPostState::new_claimed(holding_target_account_post),
+    ]
 }
 
 fn initialize_account(pre_states: &[AccountWithMetadata]) -> Vec<AccountPostState> {
@@ -220,10 +232,13 @@ fn initialize_account(pre_states: &[AccountWithMetadata]) -> Vec<AccountPostStat
     let holding_values = TokenHolding::new(&definition.account_id);
 
     let definition_post = definition.account.clone();
-    let mut account_to_initialize_post = account_to_initialize.account.clone();
-    account_to_initialize_post.data = holding_values.into_data();
+    let mut account_to_initialize = account_to_initialize.account.clone();
+    account_to_initialize.data = holding_values.into_data();
 
-    vec![definition_post.into(), account_to_initialize_post.into()]
+    vec![
+        AccountPostState::new(definition_post),
+        AccountPostState::new_claimed(account_to_initialize),
+    ]
 }
 
 type Instruction = [u8; 23];

nssa/src/public_transaction/transaction.rs

@@ -153,10 +153,16 @@ impl PublicTransaction {
                 return Err(NssaError::InvalidProgramBehavior);
             }
 
-            // The invoked program claims the accounts with default program id.
             for post in program_output.post_states.iter_mut() {
+                if !post.requires_claim() {
+                    continue;
+                }
+
+                // The invoked program can only claim accounts with default program id.
                 if post.account.program_owner == DEFAULT_PROGRAM_ID {
                     post.account.program_owner = chained_call.program_id;
+                } else {
+                    return Err(NssaError::InvalidProgramBehavior);
                 }
             }
 

nssa/test_program_methods/guest/src/bin/burner.rs

@@ -1,4 +1,4 @@
-use nssa_core::program::{read_nssa_inputs, write_nssa_outputs, ProgramInput};
+use nssa_core::program::{AccountPostState, ProgramInput, read_nssa_inputs, write_nssa_outputs};
 
 type Instruction = u128;
 
@@ -17,5 +17,5 @@ fn main() {
     let mut account_post = account_pre.clone();
     account_post.balance -= balance_to_burn;
 
-    write_nssa_outputs(vec![pre], vec![account_post.into()]);
+    write_nssa_outputs(vec![pre], vec![AccountPostState::new(account_post)]);
 }

nssa/test_program_methods/guest/src/bin/chain_caller.rs

@@ -1,5 +1,6 @@
 use nssa_core::program::{
-    ChainedCall, ProgramId, ProgramInput, read_nssa_inputs, write_nssa_outputs_with_chained_call,
+    AccountPostState, ChainedCall, ProgramId, ProgramInput, read_nssa_inputs,
+    write_nssa_outputs_with_chained_call,
 };
 use risc0_zkvm::serde::to_vec;
 
@@ -37,7 +38,10 @@ fn main() {
 
     write_nssa_outputs_with_chained_call(
         vec![sender_pre.clone(), receiver_pre.clone()],
-        vec![sender_pre.account.into(), receiver_pre.account.into()],
+        vec![
+            AccountPostState::new(sender_pre.account),
+            AccountPostState::new(receiver_pre.account),
+        ],
         chained_call,
     );
 }

nssa/test_program_methods/guest/src/bin/data_changer.rs

@@ -1,4 +1,4 @@
-use nssa_core::program::{read_nssa_inputs, write_nssa_outputs, ProgramInput};
+use nssa_core::program::{AccountPostState, ProgramInput, read_nssa_inputs, write_nssa_outputs};
 
 type Instruction = ();
 
@@ -14,5 +14,5 @@ fn main() {
     let mut account_post = account_pre.clone();
     account_post.data.push(0);
 
-    write_nssa_outputs(vec![pre], vec![account_post.into()]);
+    write_nssa_outputs(vec![pre], vec![AccountPostState::new(account_post)]);
 }

nssa/test_program_methods/guest/src/bin/extra_output.rs

@@ -1,6 +1,6 @@
 use nssa_core::{
     account::Account,
-    program::{read_nssa_inputs, write_nssa_outputs, ProgramInput},
+    program::{AccountPostState, ProgramInput, read_nssa_inputs, write_nssa_outputs},
 };
 
 type Instruction = ();
@@ -15,5 +15,11 @@ fn main() {
 
     let account_pre = pre.account.clone();
 
-    write_nssa_outputs(vec![pre], vec![account_pre.into(), Account::default().into()]);
+    write_nssa_outputs(
+        vec![pre],
+        vec![
+            AccountPostState::new(account_pre),
+            AccountPostState::new(Account::default()),
+        ],
+    );
 }

nssa/test_program_methods/guest/src/bin/minter.rs

@@ -1,4 +1,4 @@
-use nssa_core::program::{read_nssa_inputs, write_nssa_outputs, ProgramInput};
+use nssa_core::program::{read_nssa_inputs, write_nssa_outputs, AccountPostState, ProgramInput};
 
 type Instruction = ();
 
@@ -14,5 +14,5 @@ fn main() {
     let mut account_post = account_pre.clone();
     account_post.balance += 1;
 
-    write_nssa_outputs(vec![pre], vec![account_post.into()]);
+    write_nssa_outputs(vec![pre], vec![AccountPostState::new(account_post)]);
 }

nssa/test_program_methods/guest/src/bin/missing_output.rs

@@ -1,4 +1,4 @@
-use nssa_core::program::{read_nssa_inputs, write_nssa_outputs, ProgramInput};
+use nssa_core::program::{AccountPostState, ProgramInput, read_nssa_inputs, write_nssa_outputs};
 
 type Instruction = ();
 
@@ -12,5 +12,5 @@ fn main() {
 
     let account_pre1 = pre1.account.clone();
 
-    write_nssa_outputs(vec![pre1, pre2], vec![account_pre1.into()]);
+    write_nssa_outputs(vec![pre1, pre2], vec![AccountPostState::new(account_pre1)]);
 }

nssa/test_program_methods/guest/src/bin/nonce_changer.rs

@@ -1,4 +1,4 @@
-use nssa_core::program::{read_nssa_inputs, write_nssa_outputs, ProgramInput};
+use nssa_core::program::{read_nssa_inputs, write_nssa_outputs, AccountPostState, ProgramInput};
 
 type Instruction = ();
 
@@ -14,5 +14,5 @@ fn main() {
     let mut account_post = account_pre.clone();
     account_post.nonce += 1;
 
-    write_nssa_outputs(vec![pre], vec![account_post.into()]);
+    write_nssa_outputs(vec![pre], vec![AccountPostState::new(account_post)]);
 }

nssa/test_program_methods/guest/src/bin/program_owner_changer.rs

@@ -1,4 +1,4 @@
-use nssa_core::program::{read_nssa_inputs, write_nssa_outputs, ProgramInput};
+use nssa_core::program::{read_nssa_inputs, write_nssa_outputs, AccountPostState, ProgramInput};
 
 type Instruction = ();
 
@@ -14,5 +14,5 @@ fn main() {
     let mut account_post = account_pre.clone();
     account_post.program_owner = [0, 1, 2, 3, 4, 5, 6, 7];
 
-    write_nssa_outputs(vec![pre], vec![account_post.into()]);
+    write_nssa_outputs(vec![pre], vec![AccountPostState::new(account_post)]);
 }

nssa/test_program_methods/guest/src/bin/simple_balance_transfer.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};
 
 type Instruction = u128;
 
@@ -20,6 +20,9 @@ fn main() {
 
     write_nssa_outputs(
         vec![sender_pre, receiver_pre],
-        vec![sender_post.into(), receiver_post.into()],
+        vec![
+            AccountPostState::new(sender_post),
+            AccountPostState::new(receiver_post),
+        ],
     );
 }