add instruction to the program output

nssa/core/src/program.rs

@@ -26,23 +26,32 @@ pub struct ChainedCall {
 #[derive(Serialize, Deserialize, Clone)]
 #[cfg_attr(any(feature = "host", test), derive(Debug, PartialEq, Eq))]
 pub struct ProgramOutput {
+    pub instruction_data: InstructionData,
     pub pre_states: Vec<AccountWithMetadata>,
     pub post_states: Vec<Account>,
     pub chained_call: Option<ChainedCall>,
 }
 
-pub fn read_nssa_inputs<T: DeserializeOwned>() -> ProgramInput<T> {
+pub fn read_nssa_inputs<T: DeserializeOwned>() -> (ProgramInput<T>, InstructionData) {
     let pre_states: Vec<AccountWithMetadata> = env::read();
     let instruction_words: InstructionData = env::read();
     let instruction = T::deserialize(&mut Deserializer::new(instruction_words.as_ref())).unwrap();
-    ProgramInput {
+    (
-        pre_states,
+        ProgramInput {
-        instruction,
+            pre_states,
-    }
+            instruction,
+        },
+        instruction_words,
+    )
 }
 
-pub fn write_nssa_outputs(pre_states: Vec<AccountWithMetadata>, post_states: Vec<Account>) {
+pub fn write_nssa_outputs(
+    instruction_data: InstructionData,
+    pre_states: Vec<AccountWithMetadata>,
+    post_states: Vec<Account>,
+) {
     let output = ProgramOutput {
+        instruction_data,
         pre_states,
         post_states,
         chained_call: None,
@@ -51,11 +60,13 @@ pub fn write_nssa_outputs(pre_states: Vec<AccountWithMetadata>, post_states: Vec
 }
 
 pub fn write_nssa_outputs_with_chained_call(
+    instruction_data: InstructionData,
     pre_states: Vec<AccountWithMetadata>,
     post_states: Vec<Account>,
     chained_call: Option<ChainedCall>,
 ) {
     let output = ProgramOutput {
+        instruction_data,
         pre_states,
         post_states,
         chained_call,

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

@@ -2,37 +2,42 @@ use nssa_core::{
     account::{Account, AccountWithMetadata},
     program::{ProgramInput, read_nssa_inputs, write_nssa_outputs},
 };
+use risc0_zkvm::serde::to_vec;
 
 /// Initializes a default account under the ownership of this program.
 /// This is achieved by a noop.
-fn initialize_account(pre_state: AccountWithMetadata) {
+fn initialize_account(pre_state: AccountWithMetadata) -> (AccountWithMetadata, Account) {
     let account_to_claim = 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::default() {
-        return;
+        panic!("Invalid input");
     }
 
     // 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]);
+    (pre_state, 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<AccountWithMetadata>, Vec<Account>) {
     // 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
@@ -40,23 +45,30 @@ fn transfer(sender: AccountWithMetadata, recipient: AccountWithMetadata, balance
     let mut recipient_post = recipient.account.clone();
     sender_post.balance -= balance_to_move;
     recipient_post.balance += balance_to_move;
-
+    (vec![sender, recipient], vec![sender_post, recipient_post])
-    write_nssa_outputs(vec![sender, recipient], 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 {
+    let (
-        pre_states,
+        ProgramInput {
-        instruction: balance_to_move,
+            pre_states,
-    } = read_nssa_inputs();
+            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()),
+        ([account_to_claim], 0) => {
+            let (pre, post) = initialize_account(account_to_claim.clone());
+            write_nssa_outputs(instruction_words, vec![pre], vec![post]);
+        }
         ([sender, recipient], balance_to_move) => {
-            transfer(sender.clone(), recipient.clone(), balance_to_move)
+            let (pre_states, post_states) =
+                transfer(sender.clone(), recipient.clone(), balance_to_move);
+            write_nssa_outputs(instruction_words, pre_states, post_states);
         }
         _ => panic!("invalid params"),
     }

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

@@ -1,5 +1,8 @@
 use nssa_core::program::{ProgramInput, read_nssa_inputs, write_nssa_outputs};
-use risc0_zkvm::sha::{Impl, Sha256};
+use risc0_zkvm::{
+    serde::to_vec,
+    sha::{Impl, Sha256},
+};
 
 const PRIZE: u128 = 150;
 
@@ -44,10 +47,13 @@ impl Challenge {
 fn main() {
     // Read input accounts.
     // It is expected to receive only two accounts: [pinata_account, winner_account]
-    let ProgramInput {
+    let (
-        pre_states,
+        ProgramInput {
-        instruction: solution,
+            pre_states,
-    } = read_nssa_inputs::<Instruction>();
+            instruction: solution,
+        },
+        instruction_data,
+    ) = read_nssa_inputs::<Instruction>();
 
     let [pinata, winner] = match pre_states.try_into() {
         Ok(array) => array,
@@ -66,5 +72,9 @@ fn main() {
     pinata_post.data = data.next_data().to_vec();
     winner_post.balance += PRIZE;
 
-    write_nssa_outputs(vec![pinata, winner], vec![pinata_post, winner_post]);
+    write_nssa_outputs(
+        to_vec(&solution).unwrap(),
+        vec![pinata, winner],
+        vec![pinata_post, winner_post],
+    );
 }

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

@@ -8,7 +8,7 @@ use nssa_core::{
     account::{Account, AccountId, AccountWithMetadata},
     compute_digest_for_path,
     encryption::Ciphertext,
-    program::{DEFAULT_PROGRAM_ID, MAX_NUMBER_CHAINED_CALLS, ProgramOutput, validate_execution},
+    program::{DEFAULT_PROGRAM_ID, MAX_NUMBER_CHAINED_CALLS, validate_execution},
 };
 
 fn main() {
@@ -24,12 +24,30 @@ fn main() {
     let mut pre_states: Vec<AccountWithMetadata> = Vec::new();
     let mut state_diff: HashMap<AccountId, Account> = HashMap::new();
 
-    let mut program_output = program_outputs[0].clone();
+    let num_calls = program_outputs.len();
+    if num_calls > MAX_NUMBER_CHAINED_CALLS {
+        panic!("Max deapth is exceeded");
+    }
+
+    if program_outputs[num_calls - 1].chained_call.is_some() {
+        panic!("Call stack is incomplete");
+    }
+
+    for i in 0..(program_outputs.len() - 1) {
+        let Some(chained_call) = program_outputs[i].chained_call.clone() else {
+            panic!("Expected chained call");
+        };
+
+        // Check that instruction data in caller is the instruction data in callee
+        if chained_call.instruction_data != program_outputs[i + 1].instruction_data {
+            panic!("Invalid instruction data");
+        }
+    }
+
+    for program_output in program_outputs {
+        let mut program_output = program_output.clone();
 
-    for _i in 0..MAX_NUMBER_CHAINED_CALLS {
         // Check that `program_output` is consistent with the execution of the corresponding program.
-        // TODO: Program output should contain the instruction data to verify the chain of call si
-        // performed correctly.
         env::verify(program_id, &to_vec(&program_output).unwrap()).unwrap();
 
         // Check that the program is well behaved.
@@ -54,7 +72,11 @@ fn main() {
             .iter()
             .zip(&program_output.post_states)
         {
-            if !state_diff.contains_key(&pre.account_id) {
+            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.clone());
@@ -62,29 +84,6 @@ fn main() {
 
         if let Some(next_chained_call) = &program_output.chained_call {
             program_id = next_chained_call.program_id;
-
-            // // Build post states with metadata for next call
-            // let mut post_states_with_metadata = Vec::new();
-            // for (pre, post) in program_output
-            //     .pre_states
-            //     .iter()
-            //     .zip(program_output.post_states)
-            // {
-            //     let mut post_with_metadata = pre.clone();
-            //     post_with_metadata.account = post.clone();
-            //     post_states_with_metadata.push(post_with_metadata);
-            // }
-
-            // input_pre_states = next_chained_call
-            //     .account_indices
-            //     .iter()
-            //     .map(|&i| {
-            //         post_states_with_metadata
-            //             .get(i)
-            //             .ok_or_else(|| NssaError::InvalidInput("Invalid account indices".into()))
-            //             .cloned()
-            //     })
-            //     .collect::<Result<Vec<_>, NssaError>>()?;
         } else {
             break;
         };

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

@@ -170,10 +170,13 @@ fn new_definition(
 type Instruction = [u8; 23];
 
 fn main() {
-    let ProgramInput {
+    let (
-        pre_states,
+        ProgramInput {
-        instruction,
+            pre_states,
-    } = read_nssa_inputs::<Instruction>();
+            instruction,
+        },
+        instruction_words,
+    ) = read_nssa_inputs::<Instruction>();
 
     match instruction[0] {
         0 => {
@@ -184,7 +187,7 @@ fn main() {
 
             // Execute
             let post_states = new_definition(&pre_states, name, total_supply);
-            write_nssa_outputs(pre_states, post_states);
+            write_nssa_outputs(instruction_words, pre_states, post_states);
         }
         1 => {
             // Parse instruction
@@ -194,7 +197,7 @@ fn main() {
 
             // Execute
             let post_states = transfer(&pre_states, balance_to_move);
-            write_nssa_outputs(pre_states, post_states);
+            write_nssa_outputs(instruction_words, pre_states, post_states);
         }
         _ => panic!("Invalid instruction"),
     };
@@ -204,7 +207,7 @@ fn main() {
 mod tests {
     use nssa_core::account::{Account, AccountId, AccountWithMetadata};
 
-    use crate::{new_definition, transfer, TOKEN_HOLDING_DATA_SIZE, TOKEN_HOLDING_TYPE};
+    use crate::{TOKEN_HOLDING_DATA_SIZE, TOKEN_HOLDING_TYPE, new_definition, transfer};
 
     #[should_panic(expected = "Invalid number of input accounts")]
     #[test]

nssa/src/privacy_preserving_transaction/circuit.rs

@@ -27,6 +27,7 @@ pub fn execute_and_prove(
 ) -> Result<(PrivacyPreservingCircuitOutput, Proof), NssaError> {
     let mut env_builder = ExecutorEnv::builder();
     let mut program_outputs = Vec::new();
+
     for _i in 0..MAX_NUMBER_CHAINED_CALLS {
         let inner_receipt = execute_and_prove_program(program, pre_states, instruction_data)?;
 

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

@@ -1,12 +1,15 @@
-use nssa_core::program::{read_nssa_inputs, write_nssa_outputs, ProgramInput};
+use nssa_core::program::{ProgramInput, read_nssa_inputs, write_nssa_outputs};
 
 type Instruction = u128;
 
 fn main() {
-    let ProgramInput {
+    let (
-        pre_states,
+        ProgramInput {
-        instruction: balance_to_burn,
+            pre_states,
-    } = read_nssa_inputs::<Instruction>();
+            instruction: balance_to_burn,
+        },
+        instruction_words,
+    ) = read_nssa_inputs::<Instruction>();
 
     let [pre] = match pre_states.try_into() {
         Ok(array) => array,
@@ -17,5 +20,5 @@ fn main() {
     let mut account_post = account_pre.clone();
     account_post.balance -= balance_to_burn;
 
-    write_nssa_outputs(vec![pre], vec![account_post]);
+    write_nssa_outputs(instruction_words, vec![pre], vec![account_post]);
 }

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

@@ -8,10 +8,13 @@ type Instruction = (u128, ProgramId);
 /// A program that calls another program.
 /// It permutes the order of the input accounts on the subsequent call
 fn main() {
-    let ProgramInput {
+    let (
-        pre_states,
+        ProgramInput {
-        instruction: (balance, program_id),
+            pre_states,
-    } = read_nssa_inputs::<Instruction>();
+            instruction: (balance, program_id),
+        },
+        instruction_words,
+    ) = read_nssa_inputs::<Instruction>();
 
     let [sender_pre, receiver_pre] = match pre_states.try_into() {
         Ok(array) => array,
@@ -27,6 +30,7 @@ fn main() {
     });
 
     write_nssa_outputs_with_chained_call(
+        instruction_words,
         vec![sender_pre.clone(), receiver_pre.clone()],
         vec![sender_pre.account, receiver_pre.account],
         chained_call,

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

@@ -1,9 +1,9 @@
-use nssa_core::program::{read_nssa_inputs, write_nssa_outputs, ProgramInput};
+use nssa_core::program::{ProgramInput, read_nssa_inputs, write_nssa_outputs};
 
 type Instruction = ();
 
 fn main() {
-    let ProgramInput { pre_states, .. } = read_nssa_inputs::<Instruction>();
+    let (ProgramInput { pre_states, .. }, instruction_words) = read_nssa_inputs::<Instruction>();
 
     let [pre] = match pre_states.try_into() {
         Ok(array) => array,
@@ -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]);
+    write_nssa_outputs(instruction_words, vec![pre], vec![account_post]);
 }

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

@@ -1,12 +1,12 @@
 use nssa_core::{
     account::Account,
-    program::{read_nssa_inputs, write_nssa_outputs, ProgramInput},
+    program::{ProgramInput, read_nssa_inputs, write_nssa_outputs},
 };
 
 type Instruction = ();
 
 fn main() {
-    let ProgramInput { pre_states, .. } = read_nssa_inputs::<Instruction>();
+    let (ProgramInput { pre_states, .. }, instruction_words) = read_nssa_inputs::<Instruction>();
 
     let [pre] = match pre_states.try_into() {
         Ok(array) => array,
@@ -15,5 +15,9 @@ fn main() {
 
     let account_pre = pre.account.clone();
 
-    write_nssa_outputs(vec![pre], vec![account_pre, Account::default()]);
+    write_nssa_outputs(
+        instruction_words,
+        vec![pre],
+        vec![account_pre, Account::default()],
+    );
 }

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

@@ -1,9 +1,9 @@
-use nssa_core::program::{read_nssa_inputs, write_nssa_outputs, ProgramInput};
+use nssa_core::program::{ProgramInput, read_nssa_inputs, write_nssa_outputs};
 
 type Instruction = ();
 
 fn main() {
-    let ProgramInput { pre_states, .. } = read_nssa_inputs::<Instruction>();
+    let (ProgramInput { pre_states, .. }, instruction_words) = read_nssa_inputs::<Instruction>();
 
     let [pre] = match pre_states.try_into() {
         Ok(array) => array,
@@ -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]);
+    write_nssa_outputs(instruction_words, vec![pre], vec![account_post]);
 }

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

@@ -1,9 +1,9 @@
-use nssa_core::program::{read_nssa_inputs, write_nssa_outputs, ProgramInput};
+use nssa_core::program::{ProgramInput, read_nssa_inputs, write_nssa_outputs};
 
 type Instruction = ();
 
 fn main() {
-    let ProgramInput { pre_states, .. } = read_nssa_inputs::<Instruction>();
+    let (ProgramInput { pre_states, .. }, instruction_words) = read_nssa_inputs::<Instruction>();
 
     let [pre1, pre2] = match pre_states.try_into() {
         Ok(array) => array,
@@ -12,5 +12,5 @@ fn main() {
 
     let account_pre1 = pre1.account.clone();
 
-    write_nssa_outputs(vec![pre1, pre2], vec![account_pre1]);
+    write_nssa_outputs(instruction_words, vec![pre1, pre2], vec![account_pre1]);
 }

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

@@ -3,7 +3,7 @@ use nssa_core::program::{read_nssa_inputs, write_nssa_outputs, ProgramInput};
 type Instruction = ();
 
 fn main() {
-    let ProgramInput { pre_states, .. } = read_nssa_inputs::<Instruction>();
+    let (ProgramInput { pre_states, .. } , instruction_words) = read_nssa_inputs::<Instruction>();
 
     let [pre] = match pre_states.try_into() {
         Ok(array) => array,
@@ -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]);
+    write_nssa_outputs(instruction_words, vec![pre], vec![account_post]);
 }

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

@@ -1,9 +1,9 @@
-use nssa_core::program::{read_nssa_inputs, write_nssa_outputs, ProgramInput};
+use nssa_core::program::{ProgramInput, read_nssa_inputs, write_nssa_outputs};
 
 type Instruction = ();
 
 fn main() {
-    let ProgramInput { pre_states, .. } = read_nssa_inputs::<Instruction>();
+    let (ProgramInput { pre_states, .. }, instruction_words) = read_nssa_inputs::<Instruction>();
 
     let [pre] = match pre_states.try_into() {
         Ok(array) => array,
@@ -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]);
+    write_nssa_outputs(instruction_words, vec![pre], vec![account_post]);
 }

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

@@ -1,12 +1,15 @@
-use nssa_core::program::{read_nssa_inputs, write_nssa_outputs, ProgramInput};
+use nssa_core::program::{ProgramInput, read_nssa_inputs, write_nssa_outputs};
 
 type Instruction = u128;
 
 fn main() {
-    let ProgramInput {
+    let (
-        pre_states,
+        ProgramInput {
-        instruction: balance,
+            pre_states,
-    } = read_nssa_inputs::<Instruction>();
+            instruction: balance,
+        },
+        instruction_words,
+    ) = read_nssa_inputs::<Instruction>();
 
     let [sender_pre, receiver_pre] = match pre_states.try_into() {
         Ok(array) => array,
@@ -19,6 +22,7 @@ fn main() {
     receiver_post.balance += balance;
 
     write_nssa_outputs(
+        instruction_words,
         vec![sender_pre, receiver_pre],
         vec![sender_post, receiver_post],
     );