Merge pull request #1298 from o1-labs/feature/add-environment

Add initial `witness::Env` for MIPS VM

optimism/Cargo.toml

@@ -13,7 +13,8 @@ license = "Apache-2.0"
 path = "src/lib.rs"
 
 [dependencies]
-kimchi = { path = "../kimchi", version = "0.1.0" }
+ark-bn254 = { version = "0.3.0" }
+kimchi = { path = "../kimchi", version = "0.1.0", features = [ "bn254" ] }
 poly-commitment = { path = "../poly-commitment", version = "0.1.0" }
 groupmap = { path = "../groupmap", version = "0.1.0" }
 mina-curves = { path = "../curves", version = "0.1.0" }

optimism/src/lib.rs

@@ -1 +1,2 @@
 pub mod cannon;
+pub mod mips;

optimism/src/main.rs

@@ -1,5 +1,8 @@
 use clap::{arg, value_parser, Arg, ArgAction, Command};
-use kimchi_optimism::cannon::{State, VmConfiguration};
+use kimchi_optimism::{
+    cannon::{State, VmConfiguration},
+    mips::witness,
+};
 use std::{fs::File, io::BufReader, process::ExitCode};
 
 fn cli() -> VmConfiguration {
@@ -116,7 +119,7 @@ pub fn main() -> ExitCode {
 
     let reader = BufReader::new(file);
     // Read the JSON contents of the file as an instance of `State`.
-    let _state: State = serde_json::from_reader(reader).expect("Error reading input state file");
+    let state: State = serde_json::from_reader(reader).expect("Error reading input state file");
 
     if let Some(host_program) = configuration.host {
         println!("Launching host program {}", host_program.name);
@@ -127,6 +130,14 @@ pub fn main() -> ExitCode {
             .expect("Could not spawn host process");
     };
 
+    let page_size = 1 << 12;
+
+    let mut env = witness::Env::<ark_bn254::Fq>::create(page_size, state);
+
+    while !env.halt {
+        env.step();
+    }
+
     // TODO: Logic
     ExitCode::FAILURE
 }

optimism/src/mips/mod.rs

@@ -0,0 +1,2 @@
+pub mod registers;
+pub mod witness;

optimism/src/mips/registers.rs

@@ -0,0 +1,47 @@
+use serde::{Deserialize, Serialize};
+use std::ops::{Index, IndexMut};
+
+pub const NUM_REGISTERS: usize = 34;
+
+#[derive(Clone, Default, Debug, Serialize, Deserialize)]
+pub struct Registers<T> {
+    pub general_purpose: [T; 32],
+    pub hi: T,
+    pub lo: T,
+}
+
+impl<T> Registers<T> {
+    pub fn iter(&self) -> impl Iterator<Item = &T> {
+        self.general_purpose.iter().chain([&self.hi, &self.lo])
+    }
+}
+
+impl<T: Clone> Index<usize> for Registers<T> {
+    type Output = T;
+
+    fn index(&self, index: usize) -> &Self::Output {
+        if index < 32 {
+            &self.general_purpose[index]
+        } else if index == 32 {
+            &self.hi
+        } else if index == 33 {
+            &self.lo
+        } else {
+            panic!("Index out of bounds");
+        }
+    }
+}
+
+impl<T: Clone> IndexMut<usize> for Registers<T> {
+    fn index_mut(&mut self, index: usize) -> &mut Self::Output {
+        if index < 32 {
+            &mut self.general_purpose[index]
+        } else if index == 32 {
+            &mut self.hi
+        } else if index == 33 {
+            &mut self.lo
+        } else {
+            panic!("Index out of bounds");
+        }
+    }
+}

optimism/src/mips/witness.rs

@@ -0,0 +1,99 @@
+use crate::{cannon::State, mips::registers::Registers};
+use ark_ff::Field;
+use std::array;
+
+pub const NUM_GLOBAL_LOOKUP_TERMS: usize = 1;
+pub const NUM_DECODING_LOOKUP_TERMS: usize = 2;
+pub const NUM_INSTRUCTION_LOOKUP_TERMS: usize = 5;
+pub const NUM_LOOKUP_TERMS: usize =
+    NUM_GLOBAL_LOOKUP_TERMS + NUM_DECODING_LOOKUP_TERMS + NUM_INSTRUCTION_LOOKUP_TERMS;
+pub const SCRATCH_SIZE: usize = 25;
+
+#[derive(Clone)]
+pub struct SyscallEnv {
+    pub heap: u32, // Heap pointer (actually unused in Cannon as of [2023-10-18])
+    pub preimage_offset: u32,
+    pub preimage_key: Vec<u8>,
+    pub last_hint: Option<Vec<u8>>,
+}
+
+impl SyscallEnv {
+    pub fn create(state: &State) -> Self {
+        SyscallEnv {
+            heap: state.heap,
+            preimage_key: state.preimage_key.as_bytes().to_vec(), // Might not be correct
+            preimage_offset: state.preimage_offset,
+            last_hint: state.last_hint.clone(),
+        }
+    }
+}
+
+#[derive(Clone)]
+pub struct Env<Fp> {
+    pub instruction_counter: usize,
+    pub memory: Vec<(u32, Vec<u8>)>,
+    pub memory_write_index: Vec<(u32, Vec<usize>)>,
+    pub registers: Registers<u32>,
+    pub registers_write_index: Registers<usize>,
+    pub instruction_pointer: u32,
+    pub scratch_state_idx: usize,
+    pub scratch_state: [Fp; SCRATCH_SIZE],
+    pub halt: bool,
+    pub syscall_env: SyscallEnv,
+}
+
+fn fresh_scratch_state<Fp: Field, const N: usize>() -> [Fp; N] {
+    array::from_fn(|_| Fp::zero())
+}
+
+impl<Fp: Field> Env<Fp> {
+    pub fn create(page_size: usize, state: State) -> Self {
+        let initial_instruction_pointer = state.pc;
+
+        let syscall_env = SyscallEnv::create(&state);
+
+        let mut initial_memory: Vec<(u32, Vec<u8>)> = state
+            .memory
+            .into_iter()
+            // Check that the conversion from page data is correct
+            .map(|page| (page.index, page.data))
+            .collect();
+
+        for (_address, initial_memory) in initial_memory.iter_mut() {
+            initial_memory.extend((0..(page_size - initial_memory.len())).map(|_| 0u8));
+            assert_eq!(initial_memory.len(), page_size);
+        }
+
+        let memory_offsets = initial_memory
+            .iter()
+            .map(|(offset, _)| *offset)
+            .collect::<Vec<_>>();
+
+        let initial_registers = Registers {
+            lo: state.lo,
+            hi: state.hi,
+            general_purpose: state.registers,
+        };
+
+        Env {
+            instruction_counter: state.step as usize,
+            memory: initial_memory.clone(),
+            memory_write_index: memory_offsets
+                .iter()
+                .map(|offset| (*offset, vec![0usize; page_size]))
+                .collect(),
+            registers: initial_registers.clone(),
+            registers_write_index: Registers::default(),
+            instruction_pointer: initial_instruction_pointer,
+            scratch_state_idx: 0,
+            scratch_state: fresh_scratch_state(),
+            halt: state.exited,
+            syscall_env,
+        }
+    }
+
+    pub fn step(&mut self) {
+        // TODO
+        self.halt = true;
+    }
+}