From a694e705ea5159d554053b37e490d10aa03e358b Mon Sep 17 00:00:00 2001
From: Sergio Chouhy <sergio.chouhy@gmail.com>
Date: Thu, 14 Aug 2025 14:48:20 -0300
Subject: [PATCH] move claiming accounts logic to state transition

---
 nssa/src/program.rs | 18 +-----------------
 nssa/src/state.rs   | 10 +++++++++-
 2 files changed, 10 insertions(+), 18 deletions(-)

diff --git a/nssa/src/program.rs b/nssa/src/program.rs
index 7ab6d74..b104eca 100644
--- a/nssa/src/program.rs
+++ b/nssa/src/program.rs
@@ -45,28 +45,15 @@ impl Program {
 
         // Get outputs
         let ProgramOutput {
-            post_states: mut post_states,
-            ..
+            mut post_states, ..
         } = session_info
             .journal
             .decode()
             .map_err(|e| NssaError::ProgramExecutionFailed(e.to_string()))?;
 
-        // TODO: Move this logic to `V01State::transition_from_public_transaction`.
-        self.claim_accounts_with_default_program_owner(&mut post_states);
-
         Ok(post_states)
     }
 
-    fn claim_accounts_with_default_program_owner(&self, post_states: &mut [Account]) {
-        // Claim any output account with default program owner field
-        for account in post_states.iter_mut() {
-            if account.program_owner == DEFAULT_PROGRAM_ID {
-                account.program_owner = self.id;
-            }
-        }
-    }
-
     /// Executes and proves the program `P`.
     /// Returns the proof
     fn execute_and_prove(
@@ -218,13 +205,10 @@ mod tests {
 
         let expected_sender_post = Account {
             balance: 77665544332211 - balance_to_move,
-            program_owner: program.id(),
             ..Account::default()
         };
         let expected_recipient_post = Account {
             balance: balance_to_move,
-            // Program claims the account since the pre_state has default prorgam owner
-            program_owner: program.id(),
             ..Account::default()
         };
         let [sender_post, recipient_post] = program
diff --git a/nssa/src/state.rs b/nssa/src/state.rs
index f6d932e..3060185 100644
--- a/nssa/src/state.rs
+++ b/nssa/src/state.rs
@@ -1,7 +1,10 @@
 use crate::{
     address::Address, error::NssaError, program::Program, public_transaction::PublicTransaction,
 };
-use nssa_core::{account::Account, program::ProgramId};
+use nssa_core::{
+    account::Account,
+    program::{DEFAULT_PROGRAM_ID, ProgramId},
+};
 use std::collections::HashMap;
 
 pub struct V01State {
@@ -48,7 +51,12 @@ impl V01State {
 
         for (address, post) in state_diff.into_iter() {
             let current_account = self.get_account_by_address_mut(address);
+
             *current_account = post;
+            // The invoked program claims the accounts with default program id.
+            if current_account.program_owner == DEFAULT_PROGRAM_ID {
+                current_account.program_owner = tx.message().program_id;
+            }
         }
 
         for address in tx.signer_addresses() {