Merge branch 'main' into marvin/refactor-amm-program

nssa/Cargo.toml

@@ -2,6 +2,7 @@
 name = "nssa"
 version = "0.1.0"
 edition = "2024"
+license = { workspace = true }
 
 [dependencies]
 nssa_core = { workspace = true, features = ["host"] }

nssa/core/Cargo.toml

@@ -2,6 +2,7 @@
 name = "nssa_core"
 version = "0.1.0"
 edition = "2024"
+license = { workspace = true }
 
 [dependencies]
 risc0-zkvm.workspace = true

nssa/core/src/commitment.rs

@@ -5,7 +5,10 @@ use serde::{Deserialize, Serialize};
 use crate::{NullifierPublicKey, account::Account};
 
 #[derive(Serialize, Deserialize, BorshSerialize, BorshDeserialize)]
-#[cfg_attr(any(feature = "host", test), derive(Debug, Clone, PartialEq, Eq, Hash))]
+#[cfg_attr(
+    any(feature = "host", test),
+    derive(Debug, Clone, PartialEq, Eq, Hash, PartialOrd, Ord)
+)]
 pub struct Commitment(pub(super) [u8; 32]);
 
 /// A commitment to all zero data.

nssa/core/src/nullifier.rs

@@ -42,7 +42,10 @@ impl From<&NullifierSecretKey> for NullifierPublicKey {
 pub type NullifierSecretKey = [u8; 32];
 
 #[derive(Serialize, Deserialize, BorshSerialize, BorshDeserialize)]
-#[cfg_attr(any(feature = "host", test), derive(Debug, Clone, PartialEq, Eq, Hash))]
+#[cfg_attr(
+    any(feature = "host", test),
+    derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)
+)]
 pub struct Nullifier(pub(super) [u8; 32]);
 
 impl Nullifier {

nssa/src/lib.rs

@@ -14,7 +14,7 @@ mod state;
 
 pub use nssa_core::{
     SharedSecretKey,
-    account::{Account, AccountId},
+    account::{Account, AccountId, Data},
     encryption::EphemeralPublicKey,
     program::ProgramId,
 };

nssa/src/merkle_tree/mod.rs

@@ -1,3 +1,4 @@
+use borsh::{BorshDeserialize, BorshSerialize};
 use sha2::{Digest, Sha256};
 
 mod default_values;
@@ -20,6 +21,7 @@ fn hash_value(value: &Value) -> Node {
 }
 
 #[cfg_attr(test, derive(Debug, PartialEq, Eq))]
+#[derive(BorshSerialize, BorshDeserialize)]
 pub struct MerkleTree {
     nodes: Vec<Node>,
     capacity: usize,

nssa/src/privacy_preserving_transaction/transaction.rs

@@ -5,6 +5,7 @@ use nssa_core::{
     Commitment, CommitmentSetDigest, Nullifier, PrivacyPreservingCircuitOutput,
     account::{Account, AccountWithMetadata},
 };
+use sha2::{Digest as _, digest::FixedOutput as _};
 
 use super::{message::Message, witness_set::WitnessSet};
 use crate::{
@@ -131,6 +132,13 @@ impl PrivacyPreservingTransaction {
         &self.witness_set
     }
 
+    pub fn hash(&self) -> [u8; 32] {
+        let bytes = self.to_bytes();
+        let mut hasher = sha2::Sha256::new();
+        hasher.update(&bytes);
+        hasher.finalize_fixed().into()
+    }
+
     pub(crate) fn signer_account_ids(&self) -> Vec<AccountId> {
         self.witness_set
             .signatures_and_public_keys()

nssa/src/program.rs

@@ -1,3 +1,4 @@
+use borsh::{BorshDeserialize, BorshSerialize};
 use nssa_core::{
     account::AccountWithMetadata,
     program::{InstructionData, ProgramId, ProgramOutput},
@@ -14,7 +15,7 @@ use crate::{
 /// TODO: Make this variable when fees are implemented
 const MAX_NUM_CYCLES_PUBLIC_EXECUTION: u64 = 1024 * 1024 * 32; // 32M cycles
 
-#[derive(Clone, Debug, PartialEq, Eq)]
+#[derive(Clone, Debug, PartialEq, Eq, BorshSerialize, BorshDeserialize)]
 pub struct Program {
     id: ProgramId,
     elf: Vec<u8>,

nssa/src/program_deployment_transaction/transaction.rs

@@ -1,4 +1,5 @@
 use borsh::{BorshDeserialize, BorshSerialize};
+use sha2::{Digest as _, digest::FixedOutput as _};
 
 use crate::{
     V02State, error::NssaError, program::Program, program_deployment_transaction::message::Message,
@@ -6,7 +7,7 @@ use crate::{
 
 #[derive(Debug, Clone, PartialEq, Eq, BorshSerialize, BorshDeserialize)]
 pub struct ProgramDeploymentTransaction {
-    pub(crate) message: Message,
+    pub message: Message,
 }
 
 impl ProgramDeploymentTransaction {
@@ -30,4 +31,11 @@ impl ProgramDeploymentTransaction {
             Ok(program)
         }
     }
+
+    pub fn hash(&self) -> [u8; 32] {
+        let bytes = self.to_bytes();
+        let mut hasher = sha2::Sha256::new();
+        hasher.update(&bytes);
+        hasher.finalize_fixed().into()
+    }
 }

nssa/src/public_transaction/transaction.rs

@@ -17,8 +17,8 @@ use crate::{
 
 #[derive(Debug, Clone, PartialEq, Eq, BorshSerialize, BorshDeserialize)]
 pub struct PublicTransaction {
-    message: Message,
-    witness_set: WitnessSet,
+    pub message: Message,
+    pub witness_set: WitnessSet,
 }
 
 impl PublicTransaction {

nssa/src/state.rs

@@ -1,5 +1,6 @@
-use std::collections::{HashMap, HashSet};
+use std::collections::{BTreeSet, HashMap, HashSet};
 
+use borsh::{BorshDeserialize, BorshSerialize};
 use nssa_core::{
     Commitment, CommitmentSetDigest, DUMMY_COMMITMENT, MembershipProof, Nullifier,
     account::{Account, AccountId},
@@ -15,6 +16,8 @@ use crate::{
 
 pub const MAX_NUMBER_CHAINED_CALLS: usize = 10;
 
+#[derive(BorshSerialize, BorshDeserialize)]
+#[cfg_attr(test, derive(Debug, PartialEq, Eq))]
 pub(crate) struct CommitmentSet {
     merkle_tree: MerkleTree,
     commitments: HashMap<Commitment, usize>,
@@ -60,8 +63,49 @@ impl CommitmentSet {
     }
 }
 
-type NullifierSet = HashSet<Nullifier>;
+#[cfg_attr(test, derive(Debug, PartialEq, Eq))]
+struct NullifierSet(BTreeSet<Nullifier>);
 
+impl NullifierSet {
+    fn new() -> Self {
+        Self(BTreeSet::new())
+    }
+
+    fn extend(&mut self, new_nullifiers: Vec<Nullifier>) {
+        self.0.extend(new_nullifiers);
+    }
+
+    fn contains(&self, nullifier: &Nullifier) -> bool {
+        self.0.contains(nullifier)
+    }
+}
+
+impl BorshSerialize for NullifierSet {
+    fn serialize<W: std::io::Write>(&self, writer: &mut W) -> std::io::Result<()> {
+        self.0.iter().collect::<Vec<_>>().serialize(writer)
+    }
+}
+
+impl BorshDeserialize for NullifierSet {
+    fn deserialize_reader<R: std::io::Read>(reader: &mut R) -> std::io::Result<Self> {
+        let vec = Vec::<Nullifier>::deserialize_reader(reader)?;
+
+        let mut set = BTreeSet::new();
+        for n in vec {
+            if !set.insert(n) {
+                return Err(std::io::Error::new(
+                    std::io::ErrorKind::InvalidData,
+                    "duplicate nullifier in NullifierSet",
+                ));
+            }
+        }
+
+        Ok(Self(set))
+    }
+}
+
+#[derive(BorshSerialize, BorshDeserialize)]
+#[cfg_attr(test, derive(Debug, PartialEq, Eq))]
 pub struct V02State {
     public_state: HashMap<AccountId, Account>,
     private_state: (CommitmentSet, NullifierSet),
@@ -4309,4 +4353,15 @@ pub mod tests {
         // Assert - should fail because the malicious program tries to manipulate is_authorized
         assert!(matches!(result, Err(NssaError::CircuitProvingError(_))));
     }
+
+    #[test]
+    fn test_state_serialization_roundtrip() {
+        let account_id_1 = AccountId::new([1; 32]);
+        let account_id_2 = AccountId::new([2; 32]);
+        let initial_data = [(account_id_1, 100u128), (account_id_2, 151u128)];
+        let state = V02State::new_with_genesis_accounts(&initial_data, &[]).with_test_programs();
+        let bytes = borsh::to_vec(&state).unwrap();
+        let state_from_bytes: V02State = borsh::from_slice(&bytes).unwrap();
+        assert_eq!(state, state_from_bytes);
+    }
 }