migrate traces to NVM

Cargo.toml

@@ -38,6 +38,7 @@ publish = false
 clap = { version = "4.3", features = ["derive"] }
 serde = { version = "1.0", features = ["derive"] }
 serde_json = "1.0"
+elf = { version = "0.7", default-features = false, features = ["std"] }
 
 ark-crypto-primitives = { version = "0.4.0", features = ["r1cs", "sponge", "crh", "merkle_tree"] }
 ark-std = "0.4.0"

nvm/Cargo.toml

@@ -6,9 +6,8 @@ publish = false
 
 [dependencies]
 thiserror = "1.0"
-
-# temporary dependency needed for riscv module
-elf = { version = "0.7", default-features = false, features = ["std"] }
+elf.workspace = true
+serde.workspace = true
 
 nexus-riscv = { path = "../riscv" }
 

nvm/src/ark_serde.rs

@@ -0,0 +1,57 @@
+use std::fmt;
+use serde::{Serializer, Deserializer, de::Visitor};
+use ark_serialize::{CanonicalSerialize, CanonicalDeserialize};
+
+pub fn serialize<T, S>(t: &T, s: S) -> Result<S::Ok, S::Error>
+where
+    S: Serializer,
+    T: CanonicalSerialize,
+{
+    let mut v = Vec::new();
+    t.serialize_uncompressed(&mut v)
+        .map_err(|_| serde::ser::Error::custom("ark error"))?;
+    s.serialize_bytes(&v)
+}
+
+pub fn deserialize<'a, D, T>(d: D) -> Result<T, D::Error>
+where
+    D: Deserializer<'a>,
+    T: CanonicalDeserialize,
+{
+    let v = d.deserialize_bytes(BV)?;
+    let t = T::deserialize_uncompressed(v.as_slice())
+        .map_err(|_| serde::de::Error::custom("ark Error"))?;
+    Ok(t)
+}
+
+struct BV;
+
+impl<'a> Visitor<'a> for BV {
+    type Value = Vec<u8>;
+
+    fn expecting(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
+        formatter.write_str("a byte sequence")
+    }
+
+    fn visit_bytes<E: serde::de::Error>(self, v: &[u8]) -> Result<Self::Value, E> {
+        Ok(v.to_vec())
+    }
+
+    fn visit_byte_buf<E: serde::de::Error>(self, v: Vec<u8>) -> Result<Self::Value, E> {
+        Ok(v)
+    }
+
+    fn visit_seq<A>(self, mut seq: A) -> Result<Self::Value, A::Error>
+    where
+        A: serde::de::SeqAccess<'a>,
+    {
+        let mut v = Vec::new();
+        loop {
+            match seq.next_element() {
+                Ok(Some(x)) => v.push(x),
+                Ok(None) => return Ok(v),
+                Err(e) => return Err(e),
+            }
+        }
+    }
+}

nvm/src/error.rs

@@ -5,7 +5,7 @@ use thiserror::Error;
 pub enum NVMError {
     /// Invalid instruction format, could not parse
     #[error("invalid instruction {1} at {0}")]
-    InvalidInstruction(u32, u32),
+    InvalidInstruction(u64, u32),
 
     /// Unknown ECALL number
     #[error("unknown ecall {1} at {0}")]

nvm/src/eval.rs

@@ -1,14 +1,38 @@
-use crate::error::Result;
+//! Evaluation for Nexus VM programs.
+
+use num_traits::FromPrimitive;
+
+use crate::error::{Result, NVMError::InvalidInstruction};
 use crate::instructions::{Inst, Opcode, Opcode::*, Width};
 use crate::memory::{Memory, path::Path};
 
+/// State of a running Nexus VM program.
+#[derive(Default)]
 pub struct NVM {
+    /// Current program counter.
     pub pc: u32,
+    /// Register file.
     pub regs: [u32; 32],
+    /// Most recent instruction.
+    pub inst: Inst,
+    /// Result of most recent instruction.
+    pub Z: u32,
+    /// Machine memory.
     pub memory: Memory,
+    /// Merkle proof for current instruction at pc
     pub pc_path: Path,
-    pub read_path: Path,
-    pub write_path: Path,
+    /// Merkle proof for load/store instructions.
+    pub read_path: Option<Path>,
+    /// Merkle proof for store instructions.
+    pub write_path: Option<Path>,
+}
+
+/// Generate a trivial VM with a single HALT instruction.
+pub fn halt_vm() -> NVM {
+    let mut vm = NVM { pc: 0x1000, ..NVM::default() };
+    let inst = Inst { opcode: HALT, ..Inst::default() };
+    vm.memory.write_inst(vm.pc, inst.into()).unwrap();
+    vm
 }
 
 #[inline]
@@ -21,6 +45,7 @@ fn sub32(x: u32, y: u32) -> u32 {
     x.overflowing_sub(y).0
 }
 
+// Evaluator for branch conditions.
 fn brcc(opcode: Opcode, x: u32, y: u32) -> bool {
     match opcode {
         BEQ => x == y,
@@ -33,8 +58,16 @@ fn brcc(opcode: Opcode, x: u32, y: u32) -> bool {
     }
 }
 
-pub fn step(vm: &mut NVM) -> Result<()> {
-    let inst = Inst::default();
+/// Execute one step of a running Nexus VM.
+/// This function will load the next instruction at the address
+/// located at the program counter, execute the instruction,
+/// and update the register file, program counter, and merkle
+/// proofs.
+pub fn eval_step(vm: &mut NVM) -> Result<()> {
+    let (dword, path) = vm.memory.read_inst(vm.pc)?;
+    let Some(inst) = Inst::from_u64(dword) else {
+        return Err(InvalidInstruction(dword, vm.pc));
+    };
 
     let I = inst.imm;
     let X = vm.regs[inst.rs1 as usize];
@@ -43,9 +76,14 @@ pub fn step(vm: &mut NVM) -> Result<()> {
     let YI = add32(Y, I);
     let shamt = YI & 0x1f;
 
-    let mut Z = 0u32;
     let mut PC = 0u32;
 
+    vm.inst = inst;
+    vm.Z = 0;
+    vm.pc_path = path;
+    vm.read_path = None;
+    vm.write_path = None;
+
     match inst.opcode {
         NOP => {}
         HALT => {
@@ -59,7 +97,7 @@ pub fn step(vm: &mut NVM) -> Result<()> {
         }
 
         JAL => {
-            Z = add32(vm.pc, 8);
+            vm.Z = add32(vm.pc, 8);
             PC = add32(X, I);
         }
         BEQ | BNE | BLT | BGE | BLTU | BGEU => {
@@ -73,39 +111,55 @@ pub fn step(vm: &mut NVM) -> Result<()> {
             let width = Width::try_from(inst.opcode).unwrap();
             let addr = add32(X, I);
             let (val, path) = vm.memory.load(width, addr)?;
-            vm.read_path = path;
-            Z = val;
+            vm.read_path = Some(path);
+            vm.Z = val;
         }
         SB | SH | SW => {
             // Note: unwrap cannot fail
             let width = Width::try_from(inst.opcode).unwrap();
             let addr = add32(X, I);
             let (_, path) = vm.memory.load(width, addr)?;
-            vm.read_path = path;
-            vm.write_path = vm.memory.store(width, addr, Y)?;
+            vm.read_path = Some(path);
+            vm.write_path = Some(vm.memory.store(width, addr, Y)?);
         }
 
-        ADD => Z = add32(X, YI),
-        SUB => Z = sub32(X, YI),
-        SLT => Z = ((X as i32) < (YI as i32)) as u32,
-        SLTU => Z = (X < Y) as u32,
-        SLL => Z = X << shamt,
-        SRL => Z = X >> shamt,
-        SRA => Z = ((X as i32) >> shamt) as u32,
-        AND => Z = X & YI,
-        OR => Z = X | YI,
-        XOR => Z = X ^ YI,
+        ADD => vm.Z = add32(X, YI),
+        SUB => vm.Z = sub32(X, YI),
+        SLT => vm.Z = ((X as i32) < (YI as i32)) as u32,
+        SLTU => vm.Z = (X < YI) as u32,
+        SLL => vm.Z = X << shamt,
+        SRL => vm.Z = X >> shamt,
+        SRA => vm.Z = ((X as i32) >> shamt) as u32,
+        AND => vm.Z = X & YI,
+        OR => vm.Z = X | YI,
+        XOR => vm.Z = X ^ YI,
     }
 
     if inst.rd > 0 {
-        vm.regs[inst.rd as usize] = Z;
+        vm.regs[inst.rd as usize] = vm.Z;
     }
 
     if PC == 0 {
-        vm.pc = add32(PC, 8);
+        vm.pc = add32(vm.pc, 8);
     } else {
         vm.pc = PC;
     }
 
     Ok(())
 }
+
+/// Run a VM to completion. The VM will stop when it encounters
+/// a HALT instruction.
+pub fn eval(vm: &mut NVM, verbose: bool) -> Result<()> {
+    loop {
+        let pc = vm.pc;
+        eval_step(vm)?;
+        if verbose {
+            println!("{:x} {:?}", pc, vm.inst);
+        }
+        if vm.inst.opcode == HALT {
+            break;
+        }
+    }
+    Ok(())
+}

nvm/src/lib.rs

@@ -5,5 +5,7 @@ pub mod error;
 pub mod instructions;
 mod memory;
 pub mod eval;
+pub mod trace;
 
+mod ark_serde;
 pub mod riscv;

nvm/src/memory.rs

@@ -30,6 +30,12 @@ impl Memory {
         Ok((cl.ldw(addr)?, path))
     }
 
+    /// write instruction at address
+    pub fn write_inst(&mut self, addr: u32, val: u64) -> Result<()> {
+        let _ = self.trie.update(addr, |cl| cl.sdw(addr, val))?;
+        Ok(())
+    }
+
     /// perform load according to `width`
     pub fn load(&self, width: Width, addr: u32) -> Result<(u32, Path)> {
         let (cl, path) = self.trie.query(addr);

nvm/src/memory/cacheline.rs

@@ -180,6 +180,17 @@ impl CacheLine {
         }
         Ok(())
     }
+
+    /// store 64-bit value at addr
+    pub fn sdw(&mut self, addr: u32, val: u64) -> Result<()> {
+        if (addr & 7) != 0 {
+            return Err(Misaligned(addr));
+        }
+        unsafe {
+            self.dwords[((addr >> 3) & 3) as usize] = val;
+        }
+        Ok(())
+    }
 }
 
 #[cfg(test)]