Merge pull request #1167 from 0xPolygonMiden/hacka-kernel-serde

Kernel serde

assembly/src/assembler/context.rs

@@ -140,7 +140,7 @@ impl AssemblyContext {
     ///
     /// This pops the module off the module stack and return all local procedures of the module
     /// (both exported and internal) together with the combined callset of module's procedures.
-    pub fn complete_module(&mut self) -> (Vec<NamedProcedure>, CallSet) {
+    pub fn complete_module(&mut self) -> Result<(Vec<NamedProcedure>, CallSet), AssemblyError> {
         let module_ctx = self.module_stack.pop().expect("no modules");
         if self.is_kernel && self.module_stack.is_empty() {
             // if we are compiling a kernel and this is the last module on the module stack, then
@@ -152,11 +152,11 @@ impl AssemblyContext {
                 .filter(|proc| proc.is_export())
                 .map(|proc| proc.mast_root())
                 .collect::<Vec<_>>();
-            self.kernel = Some(Kernel::new(&proc_roots));
+            self.kernel = Some(Kernel::new(&proc_roots).map_err(AssemblyError::KernelError)?);
         }
 
         // return compiled procedures and callset from the module
-        (module_ctx.compiled_procs, module_ctx.callset)
+        Ok((module_ctx.compiled_procs, module_ctx.callset))
     }
 
     // PROCEDURE PROCESSORS

assembly/src/assembler/mod.rs

@@ -240,7 +240,7 @@ impl Assembler {
         for proc_ast in module.procs().iter() {
             self.compile_procedure(proc_ast, context)?;
         }
-        let (module_procs, module_callset) = context.complete_module();
+        let (module_procs, module_callset) = context.complete_module()?;
 
         // add the compiled procedures to the assembler's cache. the procedures are added to the
         // cache only if:

assembly/src/errors.rs

@@ -1,6 +1,6 @@
 use super::{
-    ast::ProcReExport, crypto::hash::RpoDigest, tokens::SourceLocation, LibraryNamespace,
-    ProcedureId, ProcedureName, String, ToString, Token, Vec,
+    ast::ProcReExport, crypto::hash::RpoDigest, tokens::SourceLocation, KernelError,
+    LibraryNamespace, ProcedureId, ProcedureName, String, ToString, Token, Vec,
 };
 use core::fmt;
 
@@ -11,28 +11,29 @@ use core::fmt;
 #[derive(Debug, Clone, Eq, PartialEq)]
 pub enum AssemblyError {
     CallInKernel(String),
-    CallerOutOKernel,
     CallSetProcedureNotFound(RpoDigest),
+    CallerOutOKernel,
     CircularModuleDependency(Vec<String>),
     ConflictingNumLocals(String),
     DivisionByZero,
-    DuplicateProcName(String, String),
     DuplicateProcId(ProcedureId),
+    DuplicateProcName(String, String),
     ExportedProcInProgram(String),
     ImportedProcModuleNotFound(ProcedureId, String),
-    ReExportedProcModuleNotFound(ProcReExport),
     ImportedProcNotFoundInModule(ProcedureId, String),
-    InvalidProgramAssemblyContext,
     InvalidCacheLock,
+    InvalidProgramAssemblyContext,
+    Io(String),
+    KernelError(KernelError),
     KernelProcNotFound(ProcedureId),
+    LibraryError(String),
     LocalProcNotFound(u16, String),
-    ParsingError(String),
     ParamOutOfBounds(u64, u64, u64),
+    ParsingError(String),
     PhantomCallsNotAllowed(RpoDigest),
     ProcedureNameError(String),
+    ReExportedProcModuleNotFound(ProcReExport),
     SysCallInKernel(String),
-    LibraryError(String),
-    Io(String),
 }
 
 impl AssemblyError {
@@ -134,25 +135,26 @@ impl fmt::Display for AssemblyError {
         use AssemblyError::*;
         match self {
             CallInKernel(proc_name) => write!(f, "call instruction used kernel procedure '{proc_name}'"),
-            CallerOutOKernel => write!(f, "caller instruction used outside of kernel"),
             CallSetProcedureNotFound(mast_root) => write!(f, "callset procedure not found in assembler cache for procedure with MAST root {mast_root}"),
+            CallerOutOKernel => write!(f, "caller instruction used outside of kernel"),
             CircularModuleDependency(dep_chain) => write!(f, "circular module dependency in the following chain: {dep_chain:?}"),
             ConflictingNumLocals(proc_name) => write!(f, "procedure `{proc_name}` has the same MAST as another procedure but different number of locals"),
             DivisionByZero => write!(f, "division by zero"),
-            DuplicateProcName(proc_name, module_path) => write!(f, "duplicate proc name '{proc_name}' in module {module_path}"),
             DuplicateProcId(proc_id) => write!(f, "duplicate proc id {proc_id}"),
+            DuplicateProcName(proc_name, module_path) => write!(f, "duplicate proc name '{proc_name}' in module {module_path}"),
             ExportedProcInProgram(proc_name) => write!(f, "exported procedure '{proc_name}' in executable program"),
             ImportedProcModuleNotFound(proc_id, proc_name) => write!(f, "module for imported procedure `{proc_name}` with ID {proc_id} not found"),
-            ReExportedProcModuleNotFound(reexport) => write!(f, "re-exported proc {} with id {} not found", reexport.name(), reexport.proc_id()),
             ImportedProcNotFoundInModule(proc_id, module_path) => write!(f, "imported procedure {proc_id} not found in module {module_path}"),
-            InvalidProgramAssemblyContext => write!(f, "assembly context improperly initialized for program compilation"),
             InvalidCacheLock => write!(f, "an attempt was made to lock a borrowed procedures cache"),
+            InvalidProgramAssemblyContext => write!(f, "assembly context improperly initialized for program compilation"),
             Io(description) => write!(f, "I/O error: {description}"),
+            KernelError(error) => write!(f, "{}", error),
             KernelProcNotFound(proc_id) => write!(f, "procedure {proc_id} not found in kernel"),
             LibraryError(err) | ParsingError(err) | ProcedureNameError(err) => write!(f, "{err}"),
             LocalProcNotFound(proc_idx, module_path) => write!(f, "procedure at index {proc_idx} not found in module {module_path}"),
             ParamOutOfBounds(value, min, max) => write!(f, "parameter value must be greater than or equal to {min} and less than or equal to {max}, but was {value}"),
             PhantomCallsNotAllowed(mast_root) => write!(f, "cannot call phantom procedure with MAST root {mast_root}: phantom calls not allowed"),
+            ReExportedProcModuleNotFound(reexport) => write!(f, "re-exported proc {} with id {} not found", reexport.name(), reexport.proc_id()),
             SysCallInKernel(proc_name) => write!(f, "syscall instruction used in kernel procedure '{proc_name}'"),
         }
     }

assembly/src/lib.rs

@@ -7,6 +7,7 @@ extern crate alloc;
 use vm_core::{
     code_blocks::CodeBlock,
     crypto,
+    errors::KernelError,
     utils::{
         collections::{btree_map, BTreeMap, BTreeSet, Vec},
         string::{String, ToString},

core/src/errors.rs

@@ -28,6 +28,7 @@ impl std::error::Error for InputError {}
 
 // OUTPUT ERROR
 // ================================================================================================
+
 #[derive(Clone, Debug)]
 pub enum OutputError {
     InvalidOverflowAddress(u64),
@@ -56,5 +57,30 @@ impl fmt::Display for OutputError {
     }
 }
 
+// KERNEL ERROR
+// ================================================================================================
+
 #[cfg(feature = "std")]
 impl std::error::Error for OutputError {}
+
+#[derive(Debug, Copy, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
+pub enum KernelError {
+    DuplicatedProcedures,
+    TooManyProcedures(usize, usize),
+}
+
+impl fmt::Display for KernelError {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        match self {
+            KernelError::DuplicatedProcedures => {
+                write!(f, "Kernel can not have duplicated procedures",)
+            }
+            KernelError::TooManyProcedures(max, count) => {
+                write!(f, "Kernel can have at most {} procedures, received {}", max, count)
+            }
+        }
+    }
+}
+
+#[cfg(feature = "std")]
+impl std::error::Error for KernelError {}

core/src/program/info.rs

@@ -70,7 +70,7 @@ impl From<Program> for ProgramInfo {
 impl Serializable for ProgramInfo {
     fn write_into<W: ByteWriter>(&self, target: &mut W) {
         self.program_hash.write_into(target);
-        <Kernel as Serializable>::write_into(&self.kernel, target);
+        self.kernel.write_into(target);
     }
 }
 

core/src/program/mod.rs

@@ -1,5 +1,6 @@
 use super::{
     chiplets::hasher::{self, Digest},
+    errors,
     utils::{
         collections::{BTreeMap, Vec},
         Box,
@@ -126,15 +127,25 @@ impl CodeBlockTable {
 #[derive(Debug, Clone, Default, PartialEq, Eq)]
 pub struct Kernel(Vec<Digest>);
 
+pub const MAX_KERNEL_PROCEDURES: usize = u8::MAX as usize;
+
 impl Kernel {
     /// Returns a new [Kernel] instantiated with the specified procedure hashes.
-    pub fn new(proc_hashes: &[Digest]) -> Self {
-        // make sure procedure roots are ordered consistently
-        let mut hash_map: BTreeMap<[u8; 32], Digest> = BTreeMap::new();
-        proc_hashes.iter().cloned().for_each(|r| {
-            hash_map.insert(r.into(), r);
-        });
-        Self(hash_map.values().copied().collect())
+    pub fn new(proc_hashes: &[Digest]) -> Result<Self, errors::KernelError> {
+        if proc_hashes.len() > MAX_KERNEL_PROCEDURES {
+            Err(errors::KernelError::TooManyProcedures(MAX_KERNEL_PROCEDURES, proc_hashes.len()))
+        } else {
+            let mut hashes = proc_hashes.to_vec();
+            hashes.sort_by_key(|v| v.as_bytes()); // ensure consistent order
+
+            let duplicated = hashes.windows(2).any(|data| data[0] == data[1]);
+
+            if duplicated {
+                Err(errors::KernelError::DuplicatedProcedures)
+            } else {
+                Ok(Self(hashes))
+            }
+        }
     }
 
     /// Returns true if this kernel does not contain any procedures.
@@ -158,11 +169,7 @@ impl Kernel {
 // this is required by AIR as public inputs will be serialized with the proof
 impl Serializable for Kernel {
     fn write_into<W: ByteWriter>(&self, target: &mut W) {
-        // TODO the serialization of MAST will not support values greater than `u16::MAX`, so we
-        // reflect the same restriction here. however, this should be tweaked in the future. This
-        // value will likely be capped to `u8::MAX`.
-
-        debug_assert!(self.0.len() <= u16::MAX as usize);
+        debug_assert!(self.0.len() <= MAX_KERNEL_PROCEDURES);
         target.write_u16(self.0.len() as u16);
         Digest::write_batch_into(&self.0, target)
     }
@@ -171,7 +178,7 @@ impl Serializable for Kernel {
 impl Deserializable for Kernel {
     fn read_from<R: ByteReader>(source: &mut R) -> Result<Self, DeserializationError> {
         let len = source.read_u16()?;
-        let kernel = (0..len).map(|_| source.read::<Digest>()).collect::<Result<_, _>>()?;
+        let kernel = Digest::read_batch_from(source, len.into())?;
         Ok(Self(kernel))
     }
 }

core/src/program/tests.rs

@@ -23,7 +23,7 @@ proptest! {
                 Some(digest_from_seed(*seed))
             })
             .collect();
-        let kernel = Kernel::new(&kernel);
+        let kernel = Kernel::new(&kernel).unwrap();
         let program_info = ProgramInfo::new(program_hash, kernel);
         let bytes = program_info.to_bytes();
         let deser = ProgramInfo::read_from_bytes(&bytes).unwrap();

processor/src/chiplets/kernel_rom/tests.rs

@@ -134,7 +134,7 @@ fn kernel_rom_with_access() {
 
 /// Creates a kernel with two dummy procedures
 fn build_kernel() -> Kernel {
-    Kernel::new(&[PROC1_HASH.into(), PROC2_HASH.into()])
+    Kernel::new(&[PROC1_HASH.into(), PROC2_HASH.into()]).unwrap()
 }
 
 /// Builds a trace of the specified length and fills it with data from the provided KernelRom

processor/src/chiplets/tests.rs

@@ -100,7 +100,7 @@ fn stacked_chiplet_trace() {
 fn build_kernel() -> Kernel {
     let proc_hash1: Digest = [ONE, ZERO, ONE, ZERO].into();
     let proc_hash2: Digest = [ONE, ONE, ONE, ONE].into();
-    Kernel::new(&[proc_hash1, proc_hash2])
+    Kernel::new(&[proc_hash1, proc_hash2]).unwrap()
 }
 
 /// Builds a sample trace by executing a span block containing the specified operations. This

processor/src/decoder/tests.rs

@@ -1713,7 +1713,7 @@ fn build_call_trace(
     kernel_proc: Option<CodeBlock>,
 ) -> (SystemTrace, DecoderTrace, AuxTraceHints, usize) {
     let kernel = match kernel_proc {
-        Some(ref proc) => Kernel::new(&[proc.hash()]),
+        Some(ref proc) => Kernel::new(&[proc.hash()]).unwrap(),
         None => Kernel::default(),
     };
     let host = DefaultHost::default();