add test for balance invariance after program execution

nssa/src/program.rs

@@ -135,4 +135,24 @@ impl Program {
             elf: DATA_CHANGER_ELF,
         }
     }
+
+    /// A program that mints balance
+    pub fn minter() -> Self {
+        use test_program_methods::{MINTER_ELF, MINTER_ID};
+
+        Program {
+            id: MINTER_ID,
+            elf: MINTER_ELF,
+        }
+    }
+
+    /// A program that mints balance
+    pub fn burner() -> Self {
+        use test_program_methods::{BURNER_ELF, BURNER_ID};
+
+        Program {
+            id: BURNER_ID,
+            elf: BURNER_ELF,
+        }
+    }
 }

nssa/src/state.rs

@@ -152,6 +152,8 @@ impl V01State {
         self.insert_program(Program::program_owner_changer());
         self.insert_program(Program::simple_balance_transfer());
         self.insert_program(Program::data_changer());
+        self.insert_program(Program::minter());
+        self.insert_program(Program::burner());
         self
     }
 
@@ -182,4 +184,14 @@ impl V01State {
         );
         self
     }
+
+    pub fn with_account_owned_by_burner_program(mut self) -> Self {
+        let account = Account {
+            program_owner: Program::burner().id(),
+            balance: 100,
+            ..Default::default()
+        };
+        self.public_state.insert(Address::new([252; 32]), account);
+        self
+    }
 }

nssa/src/tests/state_tests.rs

@@ -4,7 +4,6 @@ use crate::{
 };
 use nssa_core::account::Account;
 
-
 fn transfer_transaction(
     from: Address,
     from_key: PrivateKey,

nssa/src/tests/valid_execution_tests.rs

@@ -101,7 +101,7 @@ fn test_program_should_fail_if_modifies_program_owner_with_only_non_default_nonc
         .with_non_default_accounts_but_default_program_owners();
     let address = Address::new([254; 32]);
     let account = state.get_account_by_address(&address);
-    // Assert the target account only differs from the default account in balance field
+    // Assert the target account only differs from the default account in nonce field
     assert_eq!(account.program_owner, Account::default().program_owner);
     assert_eq!(account.balance, Account::default().balance);
     assert_ne!(account.nonce, Account::default().nonce);
@@ -124,7 +124,7 @@ fn test_program_should_fail_if_modifies_program_owner_with_only_non_default_data
         .with_non_default_accounts_but_default_program_owners();
     let address = Address::new([253; 32]);
     let account = state.get_account_by_address(&address);
-    // Assert the target account only differs from the default account in balance field
+    // Assert the target account only differs from the default account in data field
     assert_eq!(account.program_owner, Account::default().program_owner);
     assert_eq!(account.balance, Account::default().balance);
     assert_eq!(account.nonce, Account::default().nonce);
@@ -186,3 +186,43 @@ fn test_program_should_fail_if_modifies_data_of_non_owned_account() {
 
     assert!(matches!(result, Err(NssaError::InvalidProgramBehavior)));
 }
+
+#[test]
+fn test_program_should_fail_if_does_not_preserve_total_balance_by_minting() {
+    let initial_data = [];
+    let mut state = V01State::new_with_genesis_accounts(&initial_data).with_test_programs();
+    let address = Address::new([1; 32]);
+    let program_id = Program::minter().id();
+
+    let message = public_transaction::Message::new(program_id, vec![address], vec![], 0);
+    let witness_set = public_transaction::WitnessSet::for_message(&message, &[]);
+    let tx = PublicTransaction::new(message, witness_set);
+
+    let result = state.transition_from_public_transaction(&tx);
+
+    assert!(matches!(result, Err(NssaError::InvalidProgramBehavior)));
+}
+
+#[test]
+fn test_program_should_fail_if_does_not_preserve_total_balance_by_burning() {
+    let initial_data = [];
+    let mut state = V01State::new_with_genesis_accounts(&initial_data)
+        .with_test_programs()
+        .with_account_owned_by_burner_program();
+    let program_id = Program::burner().id();
+    let address = Address::new([252; 32]);
+    assert_eq!(
+        state.get_account_by_address(&address).program_owner,
+        program_id
+    );
+    let balance_to_burn = 1;
+    assert!(state.get_account_by_address(&address).balance > balance_to_burn);
+
+    let message =
+        public_transaction::Message::new(program_id, vec![address], vec![], balance_to_burn);
+    let witness_set = public_transaction::WitnessSet::for_message(&message, &[]);
+    let tx = PublicTransaction::new(message, witness_set);
+    let result = state.transition_from_public_transaction(&tx);
+
+    assert!(matches!(result, Err(NssaError::InvalidProgramBehavior)));
+}

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

@@ -0,0 +1,20 @@
+use nssa_core::account::AccountWithMetadata;
+use risc0_zkvm::guest::env;
+
+fn main() {
+    let input_accounts: Vec<AccountWithMetadata> = env::read();
+    let balance_to_burn: u128 = env::read();
+
+    let [pre] = match input_accounts.try_into() {
+        Ok(array) => array,
+        Err(_) => return,
+    };
+
+    let account_pre = pre.account;
+    let mut account_post = account_pre.clone();
+    account_post.balance -= balance_to_burn;
+
+    env::commit(&vec![account_post]);
+}
+
+

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

@@ -0,0 +1,19 @@
+use nssa_core::account::AccountWithMetadata;
+use risc0_zkvm::guest::env;
+
+fn main() {
+    let input_accounts: Vec<AccountWithMetadata> = env::read();
+    let _instruction_data: u128 = env::read();
+
+    let [pre] = match input_accounts.try_into() {
+        Ok(array) => array,
+        Err(_) => return,
+    };
+
+    let account_pre = pre.account;
+    let mut account_post = account_pre.clone();
+    account_post.balance += 1;
+
+    env::commit(&vec![account_post]);
+}
+