Merge pull request #336 from 0xPolygonMiden/bobbin-acct-code

`AccountCode` cleanup

objects/src/accounts/code.rs

@@ -1,5 +1,10 @@
 use super::{AccountError, Assembler, AssemblyContext, Digest, ModuleAst, Vec};
-use crate::crypto::merkle::SimpleSmt;
+use crate::{
+    crypto::merkle::SimpleSmt,
+    utils::serde::{ByteReader, ByteWriter, Deserializable, DeserializationError, Serializable},
+};
+use assembly::ast::AstSerdeOptions;
+use core::cell::OnceCell;
 
 // CONSTANTS
 // ================================================================================================
@@ -10,6 +15,9 @@ const ACCOUNT_CODE_TREE_DEPTH: u8 = 8;
 /// The maximum number of account interface procedures.
 const MAX_ACCOUNT_PROCEDURES: usize = 2_usize.pow(ACCOUNT_CODE_TREE_DEPTH as u32);
 
+/// Default serialization options for account code AST.
+const MODULE_SERDE_OPTIONS: AstSerdeOptions = AstSerdeOptions::new(true);
+
 // ACCOUNT CODE
 // ================================================================================================
 
@@ -18,61 +26,59 @@ const MAX_ACCOUNT_PROCEDURES: usize = 2_usize.pow(ACCOUNT_CODE_TREE_DEPTH as u32
 /// Account's public interface consists of a set of account procedures, each procedure being a Miden
 /// VM program. Thus, MAST root of each procedure commits to the underlying program. We commit to
 /// the entire account interface by building a simple Merkle tree out of all procedure MAST roots.
-#[derive(Debug, Clone, PartialEq, Eq)]
-#[cfg_attr(feature = "serde", derive(serde::Deserialize, serde::Serialize))]
+#[derive(Debug, Clone)]
 pub struct AccountCode {
-    #[cfg_attr(feature = "serde", serde(with = "serialization"))]
     module: ModuleAst,
     procedures: Vec<Digest>,
-    procedure_tree: SimpleSmt,
+    procedure_tree: OnceCell<SimpleSmt>,
 }
 
 impl AccountCode {
-    // CONSTANTS
-    // --------------------------------------------------------------------------------------------
-    pub const ACCOUNT_CODE_NAMESPACE_BASE: &'static str = "context::account";
-
     // CONSTRUCTOR
     // --------------------------------------------------------------------------------------------
     /// Returns a new definition of an account's interface compiled from the specified source code.
     ///
     /// # Errors
     /// Returns an error if:
     /// - Compilation of the provided module fails.
-    /// - The number of procedures exported from the provided module is greater than 256.
-    pub fn new(account_module: ModuleAst, assembler: &Assembler) -> Result<Self, AccountError> {
+    /// - The number of procedures exported from the provided module is smaller than 1 or greater
+    ///   than 256.
+    pub fn new(module: ModuleAst, assembler: &Assembler) -> Result<Self, AccountError> {
         // compile the module and make sure the number of exported procedures is within the limit
-        let procedure_digests = assembler
-            .compile_module(&account_module, None, &mut AssemblyContext::for_module(false))
+        let procedures = assembler
+            .compile_module(&module, None, &mut AssemblyContext::for_module(false))
             .map_err(AccountError::AccountCodeAssemblerError)?;
 
-        if procedure_digests.len() > MAX_ACCOUNT_PROCEDURES {
+        // make sure the number of procedures is between 1 and 256 (both inclusive)
+        if procedures.is_empty() {
+            return Err(AccountError::AccountCodeNoProcedures);
+        } else if procedures.len() > MAX_ACCOUNT_PROCEDURES {
             return Err(AccountError::AccountCodeTooManyProcedures {
                 max: MAX_ACCOUNT_PROCEDURES,
-                actual: procedure_digests.len(),
+                actual: procedures.len(),
             });
         }
 
         Ok(Self {
-            procedure_tree: build_procedure_tree(&procedure_digests),
-            module: account_module,
-            procedures: procedure_digests,
+            procedure_tree: OnceCell::new(),
+            module,
+            procedures,
         })
     }
 
     /// Returns a new definition of an account's interface instantiated from the provided
     /// module and list of procedure digests.
     ///
-    /// # Safety
-    /// This function assumes that the list of provided procedure digests resulted from the
-    /// compilation of the provided module, but this is not checked.
+    /// **Note**: This function assumes that the list of provided procedure digests resulted from
+    /// the compilation of the provided module, but this is not checked.
     ///
     /// # Panics
-    /// Panics if the number of procedures is greater than 256.
-    pub unsafe fn from_parts(module: ModuleAst, procedures: Vec<Digest>) -> Self {
+    /// Panics if the number of procedures is smaller than 1 or greater than 256.
+    pub fn from_parts(module: ModuleAst, procedures: Vec<Digest>) -> Self {
+        assert!(!procedures.is_empty(), "no account procedures");
         assert!(procedures.len() <= MAX_ACCOUNT_PROCEDURES, "too many account procedures");
         Self {
-            procedure_tree: build_procedure_tree(&procedures),
+            procedure_tree: OnceCell::new(),
             module,
             procedures,
         }
@@ -98,7 +104,8 @@ impl AccountCode {
 
     /// Returns a reference to the procedure tree.
     pub fn procedure_tree(&self) -> &SimpleSmt {
-        &self.procedure_tree
+        // build procedure tree only when requested
+        self.procedure_tree.get_or_init(|| build_procedure_tree(&self.procedures))
     }
 
     /// Returns the number of public interface procedures defined for this account.
@@ -108,8 +115,7 @@ impl AccountCode {
 
     /// Returns true if a procedure with the specified root is defined for this account.
     pub fn has_procedure(&self, root: Digest) -> bool {
-        let root_bytes = root.as_bytes();
-        self.procedures.binary_search_by(|r| r.as_bytes().cmp(&root_bytes)).is_ok()
+        self.procedures.contains(&root)
     }
 
     /// Returns a procedure digest for the procedure with the specified index.
@@ -123,36 +129,44 @@ impl AccountCode {
     /// Returns the procedure index for the procedure with the specified root or None if such
     /// procedure is not defined for this account.
     pub fn get_procedure_index_by_root(&self, root: Digest) -> Option<usize> {
-        let root_bytes = root.as_bytes();
-        self.procedures.binary_search_by(|x| x.as_bytes().cmp(&root_bytes)).ok()
+        self.procedures.iter().position(|r| r == &root)
     }
 }
 
-// SERIALIZATION
+// EQUALITY
 // ================================================================================================
 
-#[cfg(feature = "serde")]
-mod serialization {
-    use assembly::ast::{AstSerdeOptions, ModuleAst};
+impl PartialEq for AccountCode {
+    fn eq(&self, other: &Self) -> bool {
+        // TODO: consider checking equality based only on the set of procedures
+        self.module == other.module && self.procedures == other.procedures
+    }
+}
 
-    pub fn serialize<S>(module: &super::ModuleAst, serializer: S) -> Result<S::Ok, S::Error>
-    where
-        S: serde::Serializer,
-    {
-        let bytes = module.to_bytes(AstSerdeOptions {
-            serialize_imports: true,
-        });
+impl Eq for AccountCode {}
 
-        serializer.serialize_bytes(&bytes)
+// SERIALIZATION
+// ================================================================================================
+
+impl Serializable for AccountCode {
+    fn write_into<W: ByteWriter>(&self, target: &mut W) {
+        self.module.write_into(target, MODULE_SERDE_OPTIONS);
+        self.module.write_source_locations(target);
+        // since the number of procedures is guaranteed to be between 1 and 256, we can store the
+        // number as a single byte - but we do have to subtract 1 to store 256 as 255.
+        target.write_u8((self.procedures.len() - 1) as u8);
+        self.procedures.write_into(target);
     }
+}
 
-    pub fn deserialize<'de, D>(deserializer: D) -> Result<super::ModuleAst, D::Error>
-    where
-        D: serde::Deserializer<'de>,
-    {
-        let bytes: Vec<u8> = <Vec<u8> as serde::Deserialize>::deserialize(deserializer)?;
+impl Deserializable for AccountCode {
+    fn read_from<R: ByteReader>(source: &mut R) -> Result<Self, DeserializationError> {
+        let mut module = ModuleAst::read_from(source, MODULE_SERDE_OPTIONS)?;
+        module.load_source_locations(source)?;
+        let num_procedures = (source.read_u8()? as usize) + 1;
+        let procedures = Digest::read_batch_from(source, num_procedures)?;
 
-        ModuleAst::from_bytes(&bytes).map_err(serde::de::Error::custom)
+        Ok(Self::from_parts(module, procedures))
     }
 }
 
@@ -177,3 +191,34 @@ fn build_procedure_tree(procedures: &[Digest]) -> SimpleSmt {
     )
     .expect("failed to build procedure tree")
 }
+
+// TESTS
+// ================================================================================================
+
+#[cfg(test)]
+mod tests {
+    use super::{AccountCode, Assembler, Deserializable, ModuleAst, Serializable};
+
+    #[test]
+    fn serialize_code() {
+        let source = "
+            export.foo
+                push.1 push.2 mul
+            end
+
+            export.bar
+                push.1 push.2 add
+            end
+        ";
+
+        // build account code from source
+        let module = ModuleAst::parse(source).unwrap();
+        let assembler = Assembler::default();
+        let code1 = AccountCode::new(module, &assembler).unwrap();
+
+        // serialize and deserialize the code; make sure deserialized version matches the original
+        let bytes = code1.to_bytes();
+        let code2 = AccountCode::read_from_bytes(&bytes).unwrap();
+        assert_eq!(code1, code2)
+    }
+}

objects/src/accounts/mod.rs

@@ -73,6 +73,7 @@ pub struct Account {
     id: AccountId,
     vault: AccountVault,
     storage: AccountStorage,
+    #[cfg_attr(feature = "serde", serde(with = "serialization"))]
     code: AccountCode,
     nonce: Felt,
 }
@@ -210,8 +211,35 @@ impl ToAdviceInputs for Account {
     }
 }
 
+// SERIALIZATION
+// ================================================================================================
+
+#[cfg(feature = "serde")]
+mod serialization {
+    use super::AccountCode;
+    use crate::utils::serde::{Deserializable, Serializable};
+
+    pub fn serialize<S>(code: &AccountCode, serializer: S) -> Result<S::Ok, S::Error>
+    where
+        S: serde::Serializer,
+    {
+        let bytes = code.to_bytes();
+        serializer.serialize_bytes(&bytes)
+    }
+
+    pub fn deserialize<'de, D>(deserializer: D) -> Result<AccountCode, D::Error>
+    where
+        D: serde::Deserializer<'de>,
+    {
+        let bytes: Vec<u8> = <Vec<u8> as serde::Deserialize>::deserialize(deserializer)?;
+
+        AccountCode::read_from_bytes(&bytes).map_err(serde::de::Error::custom)
+    }
+}
+
 // HELPERS
 // ================================================================================================
+
 /// Returns hash of an account with the specified ID, nonce, vault root, storage root, and code root.
 ///
 /// Hash of an account is computed as hash(id, nonce, vault_root, storage_root, code_root).
@@ -234,6 +262,7 @@ pub fn hash_account(
 
 // DIFF IMPLEMENTATION
 // ================================================================================================
+
 impl TryApplyDiff<Digest, StoreNode> for Account {
     type DiffType = AccountDelta;
     type Error = AccountError;

objects/src/errors.rs

@@ -14,6 +14,7 @@ use core::fmt;
 #[derive(Debug, Clone, Eq, PartialEq)]
 pub enum AccountError {
     AccountCodeAssemblerError(AssemblyError),
+    AccountCodeNoProcedures,
     AccountCodeTooManyProcedures {
         max: usize,
         actual: usize,

objects/src/lib.rs

@@ -52,4 +52,10 @@ pub mod crypto {
 pub mod utils {
     pub use miden_crypto::utils::{format, vec};
     pub use vm_core::utils::{collections, string};
+
+    pub mod serde {
+        pub use miden_crypto::utils::{
+            ByteReader, ByteWriter, Deserializable, DeserializationError, Serializable,
+        };
+    }
 }