Merge branch 'main' into Pravdyvy/amm-wallet-integration

nssa/program_methods/guest/Cargo.lock

@@ -1579,6 +1579,7 @@ dependencies = [
  "chacha20",
  "risc0-zkvm",
  "serde",
+ "thiserror",
 ]
 
 [[package]]

nssa/program_methods/guest/src/bin/amm.rs

@@ -209,17 +209,17 @@ impl TokenHolding {
         bytes[0] = self.account_type;
         bytes[1..33].copy_from_slice(&self.definition_id.to_bytes());
         bytes[33..].copy_from_slice(&self.balance.to_le_bytes());
-        bytes.into()
+        bytes.to_vec().try_into().expect("Data too big")
     }
 }
 
 
 type Instruction = Vec<u8>;
 fn main() {
-    let ProgramInput {
+    let ( ProgramInput {
         pre_states,
         instruction,
-    } = read_nssa_inputs::<Instruction>();
+    }, instruction_data) = read_nssa_inputs::<Instruction>();
 
     let (post_states, chained_calls) = match instruction[0] {
         0 => {
@@ -266,7 +266,7 @@ fn main() {
         _ => panic!("Invalid instruction"),
     };
 
-    write_nssa_outputs_with_chained_call(pre_states, post_states, chained_calls);
+    write_nssa_outputs_with_chained_call(instruction_data, pre_states, post_states, chained_calls);
 }
 
 
@@ -434,7 +434,7 @@ fn new_definition (
             active: true, 
     };
 
-    pool_post.data = pool_post_definition.into_data();
+    pool_post.data = pool_post_definition.into_data().try_into().expect("Data too big");
     let pool_post: AccountPostState = 
         if pool.account == Account::default() { AccountPostState::new_claimed(pool_post.clone()) }
         else { AccountPostState::new(pool_post.clone()) };
@@ -590,7 +590,7 @@ fn swap(
             active: true, 
     };
 
-    pool_post.data = pool_post_definition.into_data();
+    pool_post.data = pool_post_definition.into_data().try_into().expect("Data too big");
     
     let post_states = vec![
         AccountPostState::new(pool_post.clone()),
@@ -769,7 +769,7 @@ fn add_liquidity(pre_states: &[AccountWithMetadata],
             active: true,  
     };
     
-    pool_post.data = pool_post_definition.into_data();
+    pool_post.data = pool_post_definition.into_data().try_into().expect("Data too big");
     let mut chained_call = Vec::new();
 
     // Chain call for Token A (UserHoldingA -> Vault_A)
@@ -925,7 +925,7 @@ fn remove_liquidity(pre_states: &[AccountWithMetadata],
             active,  
     };
 
-    pool_post.data = pool_post_definition.into_data();
+    pool_post.data = pool_post_definition.into_data().try_into().expect("Data too big");
 
     let mut chained_calls = Vec::new();
 

nssa/program_methods/guest/src/bin/authenticated_transfer.rs

@@ -6,34 +6,37 @@ use nssa_core::{
 };
 
 /// Initializes a default account under the ownership of this program.
-fn initialize_account(pre_state: AccountWithMetadata) {
+fn initialize_account(pre_state: AccountWithMetadata) -> AccountPostState {
     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
     if account_to_claim.account() != &Account::default() {
-        return;
+        panic!("Account must be uninitialized");
     }
 
     // Continue only if the owner authorized this operation
     if !is_authorized {
-        return;
+        panic!("Invalid input");
     }
 
-    // Noop will result in account being claimed for this program
-    write_nssa_outputs(vec![pre_state], vec![account_to_claim]);
+    account_to_claim
 }
 
 /// Transfers `balance_to_move` native balance from `sender` to `recipient`.
-fn transfer(sender: AccountWithMetadata, recipient: AccountWithMetadata, balance_to_move: u128) {
+fn transfer(
+    sender: AccountWithMetadata,
+    recipient: AccountWithMetadata,
+    balance_to_move: u128,
+) -> Vec<AccountPostState> {
     // Continue only if the sender has authorized this operation
     if !sender.is_authorized {
-        return;
+        panic!("Invalid input");
     }
 
     // Continue only if the sender has enough balance
     if sender.account.balance < balance_to_move {
-        return;
+        panic!("Invalid input");
     }
 
     // Create accounts post states, with updated balances
@@ -57,23 +60,31 @@ fn transfer(sender: AccountWithMetadata, recipient: AccountWithMetadata, balance
         }
     };
 
-    write_nssa_outputs(vec![sender, recipient], vec![sender_post, recipient_post]);
+    vec![sender_post, recipient_post]
 }
 
 /// A transfer of balance program.
 /// To be used both in public and private contexts.
 fn main() {
     // Read input accounts.
-    let ProgramInput {
-        pre_states,
-        instruction: balance_to_move,
-    } = read_nssa_inputs();
+    let (
+        ProgramInput {
+            pre_states,
+            instruction: balance_to_move,
+        },
+        instruction_words,
+    ) = read_nssa_inputs();
 
-    match (pre_states.as_slice(), balance_to_move) {
-        ([account_to_claim], 0) => initialize_account(account_to_claim.clone()),
+    let post_states = match (pre_states.as_slice(), balance_to_move) {
+        ([account_to_claim], 0) => {
+            let post = initialize_account(account_to_claim.clone());
+            vec![post]
+        }
         ([sender, recipient], balance_to_move) => {
             transfer(sender.clone(), recipient.clone(), balance_to_move)
         }
         _ => panic!("invalid params"),
-    }
+    };
+
+    write_nssa_outputs(instruction_words, pre_states, post_states);
 }

nssa/program_methods/guest/src/bin/pinata.rs

@@ -44,10 +44,13 @@ impl Challenge {
 fn main() {
     // Read input accounts.
     // It is expected to receive only two accounts: [pinata_account, winner_account]
-    let ProgramInput {
-        pre_states,
-        instruction: solution,
-    } = read_nssa_inputs::<Instruction>();
+    let (
+        ProgramInput {
+            pre_states,
+            instruction: solution,
+        },
+        instruction_words,
+    ) = read_nssa_inputs::<Instruction>();
 
     let [pinata, winner] = match pre_states.try_into() {
         Ok(array) => array,
@@ -63,10 +66,15 @@ fn main() {
     let mut pinata_post = pinata.account.clone();
     let mut winner_post = winner.account.clone();
     pinata_post.balance -= PRIZE;
-    pinata_post.data = data.next_data().to_vec();
+    pinata_post.data = data
+        .next_data()
+        .to_vec()
+        .try_into()
+        .expect("33 bytes should fit into Data");
     winner_post.balance += PRIZE;
 
     write_nssa_outputs(
+        instruction_words,
         vec![pinata, winner],
         vec![
             AccountPostState::new(pinata_post),

nssa/program_methods/guest/src/bin/pinata_token.rs

@@ -1,8 +1,14 @@
-use nssa_core::program::{
-    read_nssa_inputs, write_nssa_outputs_with_chained_call, AccountPostState, ChainedCall, PdaSeed, ProgramInput
+use nssa_core::{
+    account::Data,
+    program::{
+        AccountPostState, ChainedCall, PdaSeed, ProgramInput, read_nssa_inputs,
+        write_nssa_outputs_with_chained_call,
+    },
+};
+use risc0_zkvm::{
+    serde::to_vec,
+    sha::{Impl, Sha256},
 };
-use risc0_zkvm::serde::to_vec;
-use risc0_zkvm::sha::{Impl, Sha256};
 
 const PRIZE: u128 = 150;
 
@@ -35,22 +41,26 @@ impl Challenge {
         digest[..difficulty].iter().all(|&b| b == 0)
     }
 
-    fn next_data(self) -> [u8; 33] {
+    fn next_data(self) -> Data {
         let mut result = [0; 33];
         result[0] = self.difficulty;
         result[1..].copy_from_slice(Impl::hash_bytes(&self.seed).as_bytes());
-        result
+        result.to_vec().try_into().expect("should fit")
     }
 }
 
 /// A pinata program
 fn main() {
     // Read input accounts.
-    // It is expected to receive three accounts: [pinata_definition, pinata_token_holding, winner_token_holding]
-    let ProgramInput {
-        pre_states,
-        instruction: solution,
-    } = read_nssa_inputs::<Instruction>();
+    // It is expected to receive three accounts: [pinata_definition, pinata_token_holding,
+    // winner_token_holding]
+    let (
+        ProgramInput {
+            pre_states,
+            instruction: solution,
+        },
+        instruction_words,
+    ) = read_nssa_inputs::<Instruction>();
 
     let [
         pinata_definition,
@@ -70,7 +80,7 @@ fn main() {
     let mut pinata_definition_post = pinata_definition.account.clone();
     let pinata_token_holding_post = pinata_token_holding.account.clone();
     let winner_token_holding_post = winner_token_holding.account.clone();
-    pinata_definition_post.data = data.next_data().to_vec();
+    pinata_definition_post.data = data.next_data();
 
     let mut instruction_data: [u8; 23] = [0; 23];
     instruction_data[0] = 1;
@@ -83,11 +93,15 @@ fn main() {
     let chained_calls = vec![ChainedCall {
         program_id: pinata_token_holding_post.program_owner,
         instruction_data: to_vec(&instruction_data).unwrap(),
-        pre_states: vec![pinata_token_holding_for_chain_call, winner_token_holding.clone()],
+        pre_states: vec![
+            pinata_token_holding_for_chain_call,
+            winner_token_holding.clone(),
+        ],
         pda_seeds: vec![PdaSeed::new([0; 32])],
     }];
 
     write_nssa_outputs_with_chained_call(
+        instruction_words,
         vec![
             pinata_definition,
             pinata_token_holding,

nssa/program_methods/guest/src/bin/privacy_preserving_circuit.rs

@@ -1,49 +1,119 @@
-use std::collections::HashSet;
-
-use risc0_zkvm::{guest::env, serde::to_vec};
+use std::collections::HashMap;
 
 use nssa_core::{
-    Commitment, CommitmentSetDigest, DUMMY_COMMITMENT_HASH, EncryptionScheme,
-    Nullifier, NullifierPublicKey, PrivacyPreservingCircuitInput, PrivacyPreservingCircuitOutput,
+    Commitment, CommitmentSetDigest, DUMMY_COMMITMENT_HASH, EncryptionScheme, Nullifier,
+    NullifierPublicKey, PrivacyPreservingCircuitInput, PrivacyPreservingCircuitOutput,
     account::{Account, AccountId, AccountWithMetadata},
     compute_digest_for_path,
     encryption::Ciphertext,
-    program::{DEFAULT_PROGRAM_ID, ProgramOutput, validate_execution},
+    program::{DEFAULT_PROGRAM_ID, MAX_NUMBER_CHAINED_CALLS, validate_execution},
 };
+use risc0_zkvm::{guest::env, serde::to_vec};
 
 fn main() {
     let PrivacyPreservingCircuitInput {
-        program_output,
+        program_outputs,
         visibility_mask,
         private_account_nonces,
         private_account_keys,
         private_account_auth,
-        program_id,
+        mut program_id,
     } = env::read();
 
-    // Check that `program_output` is consistent with the execution of the corresponding program.
-    env::verify(program_id, &to_vec(&program_output).unwrap()).unwrap();
+    let mut pre_states: Vec<AccountWithMetadata> = Vec::new();
+    let mut state_diff: HashMap<AccountId, Account> = HashMap::new();
 
-    let ProgramOutput {
-        pre_states,
-        post_states,
-        chained_calls,
-    } = program_output;
-
-    // TODO: implement chained calls for privacy preserving transactions
-    if !chained_calls.is_empty() {
-        panic!("Privacy preserving transactions do not support yet chained calls.")
+    let num_calls = program_outputs.len();
+    if num_calls > MAX_NUMBER_CHAINED_CALLS {
+        panic!("Max chained calls depth is exceeded");
     }
 
-    // Check that there are no repeated account ids
-    if !validate_uniqueness_of_account_ids(&pre_states) {
-        panic!("Repeated account ids found")
+    let Some(last_program_call) = program_outputs.last() else {
+        panic!("Program outputs is empty")
+    };
+
+    if !last_program_call.chained_calls.is_empty() {
+        panic!("Call stack is incomplete");
     }
 
-    // Check that the program is well behaved.
-    // See the # Programs section for the definition of the `validate_execution` method.
-    if !validate_execution(&pre_states, &post_states, program_id) {
-        panic!("Bad behaved program");
+    for window in program_outputs.windows(2) {
+        let caller = &window[0];
+        let callee = &window[1];
+
+        if caller.chained_calls.len() > 1 {
+            panic!("Privacy Multi-chained calls are not supported yet");
+        }
+
+        // TODO: Modify when multi-chain calls are supported in the circuit
+        let Some(caller_chained_call) = &caller.chained_calls.first() else {
+            panic!("Expected chained call");
+        };
+
+        // Check that instruction data in caller is the instruction data in callee
+        if caller_chained_call.instruction_data != callee.instruction_data {
+            panic!("Invalid instruction data");
+        }
+
+        // Check that account pre_states in caller are the ones in calle
+        if caller_chained_call.pre_states != callee.pre_states {
+            panic!("Invalid pre states");
+        }
+    }
+
+    for (i, program_output) in program_outputs.iter().enumerate() {
+        let mut program_output = program_output.clone();
+
+        // Check that `program_output` is consistent with the execution of the corresponding program.
+        let program_output_words =
+            &to_vec(&program_output).expect("program_output must be serializable");
+        env::verify(program_id, program_output_words)
+            .expect("program output must match the program's execution");
+
+        // Check that the program is well behaved.
+        // See the # Programs section for the definition of the `validate_execution` method.
+        if !validate_execution(
+            &program_output.pre_states,
+            &program_output.post_states,
+            program_id,
+        ) {
+            panic!("Bad behaved program");
+        }
+
+        // The invoked program claims the accounts with default program id.
+        for post in program_output
+            .post_states
+            .iter_mut()
+            .filter(|post| post.requires_claim())
+        {
+            // The invoked program can only claim accounts with default program id.
+            if post.account().program_owner == DEFAULT_PROGRAM_ID {
+                post.account_mut().program_owner = program_id;
+            } else {
+                panic!("Cannot claim an initialized account")
+            }
+        }
+
+        for (pre, post) in program_output
+            .pre_states
+            .iter()
+            .zip(&program_output.post_states)
+        {
+            if let Some(account_pre) = state_diff.get(&pre.account_id) {
+                if account_pre != &pre.account {
+                    panic!("Invalid input");
+                }
+            } else {
+                pre_states.push(pre.clone());
+            }
+            state_diff.insert(pre.account_id.clone(), post.account().clone());
+        }
+
+        // TODO: Modify when multi-chain calls are supported in the circuit
+        if let Some(next_chained_call) = &program_output.chained_calls.first() {
+            program_id = next_chained_call.program_id;
+        } else if i != program_outputs.len() - 1 {
+            panic!("Inner call without a chained call found")
+        };
     }
 
     let n_accounts = pre_states.len();
@@ -70,10 +140,8 @@ fn main() {
                 // Public account
                 public_pre_states.push(pre_states[i].clone());
 
-                let mut post = post_states[i].account().clone();
-                if pre_states[i].is_authorized {
-                    post.nonce += 1;
-                }
+                let mut post = state_diff.get(&pre_states[i].account_id).unwrap().clone();
+
                 if post.program_owner == DEFAULT_PROGRAM_ID {
                     // Claim account
                     post.program_owner = program_id;
@@ -126,7 +194,8 @@ fn main() {
                 }
 
                 // Update post-state with new nonce
-                let mut post_with_updated_values = post_states[i].account().clone();
+                let mut post_with_updated_values =
+                    state_diff.get(&pre_states[i].account_id).unwrap().clone();
                 post_with_updated_values.nonce = *new_nonce;
 
                 if post_with_updated_values.program_owner == DEFAULT_PROGRAM_ID {
@@ -175,14 +244,3 @@ fn main() {
 
     env::commit(&output);
 }
-
-fn validate_uniqueness_of_account_ids(pre_states: &[AccountWithMetadata]) -> bool {
-    let number_of_accounts = pre_states.len();
-    let number_of_account_ids = pre_states
-        .iter()
-        .map(|account| account.account_id.clone())
-        .collect::<HashSet<_>>()
-        .len();
-
-    number_of_accounts == number_of_account_ids
-}

nssa/program_methods/guest/src/bin/token.rs

@@ -25,14 +25,16 @@ use nssa_core::{
 //      * Two accounts: [definition_account, account_to_initialize].
 //      * An dummy byte string of length 23, with the following layout
 //        [0x02 || 0x00 || 0x00 || 0x00 || ... || 0x00 || 0x00].
-// 4. Burn tokens from a Toking Holding account (thus lowering total supply)
+// 4. Burn tokens from a Token Holding account (thus lowering total supply)
 //    Arguments to this function are:
 //      * Two accounts: [definition_account, holding_account].
+//      * Authorization required: holding_account
 //      * An instruction data byte string of length 23, indicating the balance to burn with the folloiwng layout
 //       [0x03 || amount (little-endian 16 bytes) || 0x00 || 0x00 || 0x00 || 0x00 || 0x00 || 0x00].
-// 5. Mint additional supply of tokens tokens to a Toking Holding account (thus increasing total supply)
+// 5. Mint additional supply of tokens tokens to a Token Holding account (thus increasing total supply)
 //    Arguments to this function are:
 //      * Two accounts: [definition_account, holding_account].
+//      * Authorization required: definition_account
 //      * An instruction data byte string of length 23, indicating the balance to mint with the folloiwng layout
 //       [0x04 || amount (little-endian 16 bytes) || 0x00 || 0x00 || 0x00 || 0x00 || 0x00 || 0x00].
 
@@ -55,12 +57,15 @@ struct TokenHolding {
 }
 
 impl TokenDefinition {
-    fn into_data(self) -> Vec<u8> {
+    fn into_data(self) -> Data {
         let mut bytes = [0; TOKEN_DEFINITION_DATA_SIZE];
         bytes[0] = self.account_type;
         bytes[1..7].copy_from_slice(&self.name);
         bytes[7..].copy_from_slice(&self.total_supply.to_le_bytes());
-        bytes.into()
+        bytes
+            .to_vec()
+            .try_into()
+            .expect("23 bytes should fit into Data")
     }
 
     fn parse(data: &[u8]) -> Option<Self> {
@@ -120,7 +125,10 @@ impl TokenHolding {
         bytes[0] = self.account_type;
         bytes[1..33].copy_from_slice(&self.definition_id.to_bytes());
         bytes[33..].copy_from_slice(&self.balance.to_le_bytes());
-        bytes.into()
+        bytes
+            .to_vec()
+            .try_into()
+            .expect("33 bytes should fit into Data")
     }
 }
 
@@ -155,7 +163,7 @@ fn transfer(pre_states: &[AccountWithMetadata], balance_to_move: u128) -> Vec<Ac
     recipient_holding.balance = recipient_holding
         .balance
         .checked_add(balance_to_move)
-        .expect("Recipient balance overflow.");
+        .expect("Recipient balance overflow");
 
     let sender_post = {
         let mut this = sender.account.clone();
@@ -282,13 +290,19 @@ fn burn(pre_states: &[AccountWithMetadata], balance_to_burn: u128) -> Vec<Accoun
     post_user_holding.data = TokenHolding::into_data(TokenHolding {
         account_type: user_values.account_type,
         definition_id: user_values.definition_id,
-        balance: user_values.balance - balance_to_burn,
+        balance: user_values
+            .balance
+            .checked_sub(balance_to_burn)
+            .expect("Checked above"),
     });
 
     post_definition.data = TokenDefinition::into_data(TokenDefinition {
         account_type: definition_values.account_type,
         name: definition_values.name,
-        total_supply: definition_values.total_supply - balance_to_burn,
+        total_supply: definition_values
+            .total_supply
+            .checked_sub(balance_to_burn)
+            .expect("Total supply underflow"),
     });
 
     vec![
@@ -315,9 +329,6 @@ fn mint_additional_supply(
     let definition_values =
         TokenDefinition::parse(&definition.account.data).expect("Definition account must be valid");
 
-    //TODO: add overflow protection
-    // TokenDefinition.supply_limit + amount_to_mint
-
     let token_holding_values: TokenHolding = if token_holding.account == Account::default() {
         TokenHolding::new(&definition.account_id)
     } else {
@@ -331,13 +342,21 @@ fn mint_additional_supply(
     let token_holding_post_data = TokenHolding {
         account_type: token_holding_values.account_type,
         definition_id: token_holding_values.definition_id,
-        balance: token_holding_values.balance + amount_to_mint,
+        balance: token_holding_values
+            .balance
+            .checked_add(amount_to_mint)
+            .expect("New balance overflow"),
     };
 
+    let post_total_supply = definition_values
+        .total_supply
+        .checked_add(amount_to_mint)
+        .expect("Total supply overflow");
+
     let post_definition_data = TokenDefinition {
         account_type: definition_values.account_type,
         name: definition_values.name,
-        total_supply: definition_values.total_supply + amount_to_mint,
+        total_supply: post_total_supply,
     };
 
     let post_definition = {
@@ -363,10 +382,13 @@ fn mint_additional_supply(
 type Instruction = [u8; 23];
 
 fn main() {
-    let ProgramInput {
-        pre_states,
-        instruction,
-    } = read_nssa_inputs::<Instruction>();
+    let (
+        ProgramInput {
+            pre_states,
+            instruction,
+        },
+        instruction_words,
+    ) = read_nssa_inputs::<Instruction>();
 
     let post_states = match instruction[0] {
         0 => {
@@ -437,7 +459,7 @@ fn main() {
         _ => panic!("Invalid instruction"),
     };
 
-    write_nssa_outputs(pre_states, post_states);
+    write_nssa_outputs(instruction_words, pre_states, post_states);
 }
 
 #[cfg(test)]
@@ -549,15 +571,15 @@ mod tests {
         let post_states = new_definition(&pre_states, [0xca, 0xfe, 0xca, 0xfe, 0xca, 0xfe], 10);
         let [definition_account, holding_account] = post_states.try_into().ok().unwrap();
         assert_eq!(
-            definition_account.account().data,
-            vec![
+            definition_account.account().data.as_ref(),
+            &[
                 0, 0xca, 0xfe, 0xca, 0xfe, 0xca, 0xfe, 10, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
                 0, 0
             ]
         );
         assert_eq!(
-            holding_account.account().data,
-            vec![
+            holding_account.account().data.as_ref(),
+            &[
                 1, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22,
                 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 10, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
                 0, 0
@@ -607,7 +629,9 @@ mod tests {
             AccountWithMetadata {
                 account: Account {
                     // First byte should be `TOKEN_HOLDING_TYPE` for token holding accounts
-                    data: vec![invalid_type; TOKEN_HOLDING_DATA_SIZE],
+                    data: vec![invalid_type; TOKEN_HOLDING_DATA_SIZE]
+                        .try_into()
+                        .unwrap(),
                     ..Account::default()
                 },
                 is_authorized: true,
@@ -629,7 +653,7 @@ mod tests {
             AccountWithMetadata {
                 account: Account {
                     // Data must be of exact length `TOKEN_HOLDING_DATA_SIZE`
-                    data: vec![1; TOKEN_HOLDING_DATA_SIZE - 1],
+                    data: vec![1; TOKEN_HOLDING_DATA_SIZE - 1].try_into().unwrap(),
                     ..Account::default()
                 },
                 is_authorized: true,
@@ -651,7 +675,7 @@ mod tests {
             AccountWithMetadata {
                 account: Account {
                     // Data must be of exact length `TOKEN_HOLDING_DATA_SIZE`
-                    data: vec![1; TOKEN_HOLDING_DATA_SIZE + 1],
+                    data: vec![1; TOKEN_HOLDING_DATA_SIZE + 1].try_into().unwrap(),
                     ..Account::default()
                 },
                 is_authorized: true,
@@ -672,7 +696,7 @@ mod tests {
         let pre_states = vec![
             AccountWithMetadata {
                 account: Account {
-                    data: vec![1; TOKEN_HOLDING_DATA_SIZE],
+                    data: vec![1; TOKEN_HOLDING_DATA_SIZE].try_into().unwrap(),
                     ..Account::default()
                 },
                 is_authorized: true,
@@ -680,10 +704,12 @@ mod tests {
             },
             AccountWithMetadata {
                 account: Account {
-                    data: vec![1]
+                    data: [1]
                         .into_iter()
                         .chain(vec![2; TOKEN_HOLDING_DATA_SIZE - 1])
-                        .collect(),
+                        .collect::<Vec<_>>()
+                        .try_into()
+                        .unwrap(),
                     ..Account::default()
                 },
                 is_authorized: true,
@@ -700,10 +726,12 @@ mod tests {
             AccountWithMetadata {
                 account: Account {
                     // Account with balance 37
-                    data: vec![1; TOKEN_HOLDING_DATA_SIZE - 16]
+                    data: [1; TOKEN_HOLDING_DATA_SIZE - 16]
                         .into_iter()
                         .chain(u128::to_le_bytes(37))
-                        .collect(),
+                        .collect::<Vec<_>>()
+                        .try_into()
+                        .unwrap(),
                     ..Account::default()
                 },
                 is_authorized: true,
@@ -711,7 +739,7 @@ mod tests {
             },
             AccountWithMetadata {
                 account: Account {
-                    data: vec![1; TOKEN_HOLDING_DATA_SIZE],
+                    data: vec![1; TOKEN_HOLDING_DATA_SIZE].try_into().unwrap(),
                     ..Account::default()
                 },
                 is_authorized: true,
@@ -729,10 +757,12 @@ mod tests {
             AccountWithMetadata {
                 account: Account {
                     // Account with balance 37
-                    data: vec![1; TOKEN_HOLDING_DATA_SIZE - 16]
+                    data: [1; TOKEN_HOLDING_DATA_SIZE - 16]
                         .into_iter()
                         .chain(u128::to_le_bytes(37))
-                        .collect(),
+                        .collect::<Vec<_>>()
+                        .try_into()
+                        .unwrap(),
                     ..Account::default()
                 },
                 is_authorized: false,
@@ -740,7 +770,7 @@ mod tests {
             },
             AccountWithMetadata {
                 account: Account {
-                    data: vec![1; TOKEN_HOLDING_DATA_SIZE],
+                    data: vec![1; TOKEN_HOLDING_DATA_SIZE].try_into().unwrap(),
                     ..Account::default()
                 },
                 is_authorized: true,
@@ -756,10 +786,12 @@ mod tests {
             AccountWithMetadata {
                 account: Account {
                     // Account with balance 37
-                    data: vec![1; TOKEN_HOLDING_DATA_SIZE - 16]
+                    data: [1; TOKEN_HOLDING_DATA_SIZE - 16]
                         .into_iter()
                         .chain(u128::to_le_bytes(37))
-                        .collect(),
+                        .collect::<Vec<_>>()
+                        .try_into()
+                        .unwrap(),
                     ..Account::default()
                 },
                 is_authorized: true,
@@ -768,10 +800,12 @@ mod tests {
             AccountWithMetadata {
                 account: Account {
                     // Account with balance 255
-                    data: vec![1; TOKEN_HOLDING_DATA_SIZE - 16]
+                    data: [1; TOKEN_HOLDING_DATA_SIZE - 16]
                         .into_iter()
                         .chain(u128::to_le_bytes(255))
-                        .collect(),
+                        .collect::<Vec<_>>()
+                        .try_into()
+                        .unwrap(),
                     ..Account::default()
                 },
                 is_authorized: true,
@@ -781,15 +815,15 @@ mod tests {
         let post_states = transfer(&pre_states, 11);
         let [sender_post, recipient_post] = post_states.try_into().ok().unwrap();
         assert_eq!(
-            sender_post.account().data,
-            vec![
+            sender_post.account().data.as_ref(),
+            [
                 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
                 1, 1, 1, 1, 1, 26, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0
             ]
         );
         assert_eq!(
-            recipient_post.account().data,
-            vec![
+            recipient_post.account().data.as_ref(),
+            [
                 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
                 1, 1, 1, 1, 1, 10, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0
             ]
@@ -805,7 +839,9 @@ mod tests {
                     data: [0; TOKEN_DEFINITION_DATA_SIZE - 16]
                         .into_iter()
                         .chain(u128::to_le_bytes(1000))
-                        .collect(),
+                        .collect::<Vec<_>>()
+                        .try_into()
+                        .unwrap(),
                     ..Account::default()
                 },
                 is_authorized: false,
@@ -819,10 +855,13 @@ mod tests {
         ];
         let post_states = initialize_account(&pre_states);
         let [definition, holding] = post_states.try_into().ok().unwrap();
-        assert_eq!(definition.account().data, pre_states[0].account.data);
         assert_eq!(
-            holding.account().data,
-            vec![
+            definition.account().data.as_ref(),
+            pre_states[0].account.data.as_ref()
+        );
+        assert_eq!(
+            holding.account().data.as_ref(),
+            [
                 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
                 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0
             ]
@@ -839,6 +878,7 @@ mod tests {
         MintSuccess,
         InitSupplyMint,
         HoldingBalanceMint,
+        MintOverflow,
     }
 
     enum AccountsEnum {
@@ -847,12 +887,14 @@ mod tests {
         HoldingDiffDef,
         HoldingSameDefAuth,
         HoldingSameDefNotAuth,
+        HoldingSameDefNotAuthOverflow,
         DefinitionAccountPostBurn,
         HoldingAccountPostBurn,
         Uninit,
         InitMint,
         DefinitionAccountMint,
         HoldingSameDefMint,
+        HoldingSameDefAuthLargeBalance,
     }
 
     enum IdEnum {
@@ -933,6 +975,20 @@ mod tests {
                 is_authorized: false,
                 account_id: helper_id_constructor(IdEnum::HoldingId),
             },
+            AccountsEnum::HoldingSameDefNotAuthOverflow => AccountWithMetadata {
+                account: Account {
+                    program_owner: [5u32; 8],
+                    balance: 0u128,
+                    data: TokenHolding::into_data(TokenHolding {
+                        account_type: TOKEN_HOLDING_TYPE,
+                        definition_id: helper_id_constructor(IdEnum::PoolDefinitionId),
+                        balance: helper_balance_constructor(BalanceEnum::InitSupply),
+                    }),
+                    nonce: 0,
+                },
+                is_authorized: false,
+                account_id: helper_id_constructor(IdEnum::HoldingId),
+            },
             AccountsEnum::DefinitionAccountPostBurn => AccountWithMetadata {
                 account: Account {
                     program_owner: [5u32; 8],
@@ -1008,6 +1064,20 @@ mod tests {
                 is_authorized: true,
                 account_id: helper_id_constructor(IdEnum::PoolDefinitionId),
             },
+            AccountsEnum::HoldingSameDefAuthLargeBalance => AccountWithMetadata {
+                account: Account {
+                    program_owner: [5u32; 8],
+                    balance: 0u128,
+                    data: TokenHolding::into_data(TokenHolding {
+                        account_type: TOKEN_HOLDING_TYPE,
+                        definition_id: helper_id_constructor(IdEnum::PoolDefinitionId),
+                        balance: helper_balance_constructor(BalanceEnum::MintOverflow),
+                    }),
+                    nonce: 0,
+                },
+                is_authorized: true,
+                account_id: helper_id_constructor(IdEnum::PoolDefinitionId),
+            },
             _ => panic!("Invalid selection"),
         }
     }
@@ -1023,6 +1093,7 @@ mod tests {
             BalanceEnum::MintSuccess => 50_000,
             BalanceEnum::InitSupplyMint => 150_000,
             BalanceEnum::HoldingBalanceMint => 51_000,
+            BalanceEnum::MintOverflow => (2 as u128).pow(128) - 40_000,
             _ => panic!("Invalid selection"),
         }
     }
@@ -1086,6 +1157,19 @@ mod tests {
         );
     }
 
+    #[test]
+    #[should_panic(expected = "Total supply underflow")]
+    fn test_burn_total_supply_underflow() {
+        let pre_states = vec![
+            helper_account_constructor(AccountsEnum::DefinitionAccountAuth),
+            helper_account_constructor(AccountsEnum::HoldingSameDefAuthLargeBalance),
+        ];
+        let _post_states = burn(
+            &pre_states,
+            helper_balance_constructor(BalanceEnum::MintOverflow),
+        );
+    }
+
     #[test]
     fn test_burn_success() {
         let pre_states = vec![
@@ -1208,4 +1292,30 @@ mod tests {
         );
         assert!(holding_post.requires_claim() == true);
     }
+
+    #[test]
+    #[should_panic(expected = "Total supply overflow")]
+    fn test_mint_total_supply_overflow() {
+        let pre_states = vec![
+            helper_account_constructor(AccountsEnum::DefinitionAccountAuth),
+            helper_account_constructor(AccountsEnum::HoldingSameDefNotAuth),
+        ];
+        let _post_states = mint_additional_supply(
+            &pre_states,
+            helper_balance_constructor(BalanceEnum::MintOverflow),
+        );
+    }
+
+    #[test]
+    #[should_panic(expected = "New balance overflow")]
+    fn test_mint_holding_account_overflow() {
+        let pre_states = vec![
+            helper_account_constructor(AccountsEnum::DefinitionAccountAuth),
+            helper_account_constructor(AccountsEnum::HoldingSameDefNotAuthOverflow),
+        ];
+        let _post_states = mint_additional_supply(
+            &pre_states,
+            helper_balance_constructor(BalanceEnum::MintOverflow),
+        );
+    }
 }