Merge remote-tracking branch 'origin/master' into matthias/flat-struc…

…ture-again

ceno_emul/src/platform.rs

@@ -1,4 +1,4 @@
-use std::ops::Range;
+use std::{collections::HashSet, ops::Range};
 
 use crate::addr::{Addr, RegIdx};
 
@@ -10,20 +10,23 @@ use crate::addr::{Addr, RegIdx};
 #[derive(Clone, Debug)]
 pub struct Platform {
     pub rom: Range<Addr>,
-    pub ram: Range<Addr>,
+    // This is an `Option` to allow `const` here.
+    pub prog_data: Option<HashSet<Addr>>,
+    pub stack: Range<Addr>,
+    pub heap: Range<Addr>,
     pub public_io: Range<Addr>,
     pub hints: Range<Addr>,
-    pub stack_top: Addr,
     /// If true, ecall instructions are no-op instead of trap. Testing only.
     pub unsafe_ecall_nop: bool,
 }
 
 pub const CENO_PLATFORM: Platform = Platform {
     rom: 0x2000_0000..0x3000_0000,
-    ram: 0x8000_0000..0xFFFF_0000,
+    prog_data: None,
+    stack: 0xB0000000..0xC0000000,
+    heap: 0x8000_0000..0xFFFF_0000,
     public_io: 0x3000_1000..0x3000_2000,
     hints: 0x4000_0000..0x5000_0000,
-    stack_top: 0xC0000000,
     unsafe_ecall_nop: false,
 };
 
@@ -34,8 +37,15 @@ impl Platform {
         self.rom.contains(&addr)
     }
 
+    pub fn is_prog_data(&self, addr: Addr) -> bool {
+        self.prog_data
+            .as_ref()
+            .map(|set| set.contains(&(addr & !0x3)))
+            .unwrap_or(false)
+    }
+
     pub fn is_ram(&self, addr: Addr) -> bool {
-        self.ram.contains(&addr)
+        self.stack.contains(&addr) || self.heap.contains(&addr) || self.is_prog_data(addr)
     }
 
     pub fn is_pub_io(&self, addr: Addr) -> bool {
@@ -66,11 +76,11 @@ impl Platform {
     // Permissions.
 
     pub fn can_read(&self, addr: Addr) -> bool {
-        self.is_rom(addr) || self.is_ram(addr) || self.is_pub_io(addr) || self.is_hints(addr)
+        self.can_write(addr)
     }
 
     pub fn can_write(&self, addr: Addr) -> bool {
-        self.is_ram(addr)
+        self.is_ram(addr) || self.is_pub_io(addr) || self.is_hints(addr)
     }
 
     // Environment calls.
@@ -110,8 +120,8 @@ mod tests {
     fn test_no_overlap() {
         let p = CENO_PLATFORM;
         // ROM and RAM do not overlap.
-        assert!(!p.is_rom(p.ram.start));
-        assert!(!p.is_rom(p.ram.end - WORD_SIZE as Addr));
+        assert!(!p.is_rom(p.heap.start));
+        assert!(!p.is_rom(p.heap.end - WORD_SIZE as Addr));
         assert!(!p.is_ram(p.rom.start));
         assert!(!p.is_ram(p.rom.end - WORD_SIZE as Addr));
         // Registers do not overlap with ROM or RAM.

ceno_emul/src/tracer.rs

@@ -198,7 +198,7 @@ impl StepRecord {
             Some(value),
             Some(Change::new(value, value)),
             Some(WriteOp {
-                addr: CENO_PLATFORM.ram.start.into(),
+                addr: CENO_PLATFORM.heap.start.into(),
                 value: Change {
                     before: value,
                     after: value,

ceno_emul/src/vm_state.rs

@@ -122,7 +122,7 @@ impl EmuContext for VMState {
             tracing::warn!("ecall ignored: syscall_id={}", function);
             self.store_register(Instruction::RD_NULL as RegIdx, 0)?;
             // Example ecall effect - any writable address will do.
-            let addr = (self.platform.stack_top - WORD_SIZE as u32).into();
+            let addr = (self.platform.stack.end - WORD_SIZE as u32).into();
             self.store_memory(addr, self.peek_memory(addr))?;
             self.set_pc(ByteAddr(self.pc) + PC_STEP_SIZE);
             Ok(true)

ceno_emul/tests/test_elf.rs

@@ -1,5 +1,9 @@
+use std::{collections::HashSet, sync::Arc};
+
 use anyhow::Result;
-use ceno_emul::{ByteAddr, CENO_PLATFORM, EmuContext, InsnKind, Platform, StepRecord, VMState};
+use ceno_emul::{
+    ByteAddr, CENO_PLATFORM, EmuContext, InsnKind, Platform, Program, StepRecord, VMState,
+};
 
 #[test]
 fn test_ceno_rt_mini() -> Result<()> {
@@ -27,7 +31,7 @@ fn test_ceno_rt_mem() -> Result<()> {
     let mut state = VMState::new_from_elf(CENO_PLATFORM, program_elf)?;
     let _steps = run(&mut state)?;
 
-    let value = state.peek_memory(CENO_PLATFORM.ram.start.into());
+    let value = state.peek_memory(CENO_PLATFORM.heap.start.into());
     assert_eq!(value, 6765, "Expected Fibonacci 20, got {}", value);
     Ok(())
 }
@@ -60,7 +64,12 @@ fn test_ceno_rt_alloc() -> Result<()> {
 #[test]
 fn test_ceno_rt_io() -> Result<()> {
     let program_elf = ceno_examples::ceno_rt_io;
-    let mut state = VMState::new_from_elf(CENO_PLATFORM, program_elf)?;
+    let program = Program::load_elf(program_elf, u32::MAX)?;
+    let platform = Platform {
+        prog_data: Some(program.image.keys().copied().collect::<HashSet<u32>>()),
+        ..CENO_PLATFORM
+    };
+    let mut state = VMState::new(platform, Arc::new(program));
     let _steps = run(&mut state)?;
 
     let all_messages = read_all_messages(&state);

ceno_zkvm/benches/fibonacci.rs

@@ -4,10 +4,10 @@ use std::{
     time::{Duration, Instant},
 };
 
-use ceno_emul::{CENO_PLATFORM, Platform, Program, WORD_SIZE};
+use ceno_emul::{Platform, Program};
 use ceno_zkvm::{
     self,
-    e2e::{Checkpoint, run_e2e_with_checkpoint},
+    e2e::{Checkpoint, Preset, run_e2e_with_checkpoint, setup_platform},
 };
 use criterion::*;
 
@@ -28,29 +28,20 @@ type Pcs = BasefoldDefault<E>;
 type E = GoldilocksExt2;
 
 // Relevant init data for fibonacci run
-fn setup() -> (Program, Platform, u32, u32) {
+fn setup() -> (Program, Platform) {
     let mut file_path = PathBuf::from(env!("CARGO_MANIFEST_DIR"));
     file_path.push("examples/fibonacci.elf");
     let stack_size = 32768;
     let heap_size = 2097152;
+    let pub_io_size = 16;
     let elf_bytes = fs::read(&file_path).expect("read elf file");
     let program = Program::load_elf(&elf_bytes, u32::MAX).unwrap();
-
-    let platform = Platform {
-        // The stack section is not mentioned in ELF headers, so we repeat the constant STACK_TOP here.
-        stack_top: 0x0020_0400,
-        rom: program.base_address
-            ..program.base_address + (program.instructions.len() * WORD_SIZE) as u32,
-        ram: 0x0010_0000..0xFFFF_0000,
-        unsafe_ecall_nop: true,
-        ..CENO_PLATFORM
-    };
-
-    (program, platform, stack_size, heap_size)
+    let platform = setup_platform(Preset::Sp1, &program, stack_size, heap_size, pub_io_size);
+    (program, platform)
 }
 
 fn fibonacci_prove(c: &mut Criterion) {
-    let (program, platform, stack_size, heap_size) = setup();
+    let (program, platform) = setup();
     for max_steps in [1usize << 20, 1usize << 21, 1usize << 22] {
         // expand more input size once runtime is acceptable
         let mut group = c.benchmark_group(format!("fibonacci_max_steps_{}", max_steps));
@@ -68,8 +59,6 @@ fn fibonacci_prove(c: &mut Criterion) {
                         run_e2e_with_checkpoint::<E, Pcs>(
                             program.clone(),
                             platform.clone(),
-                            stack_size,
-                            heap_size,
                             vec![],
                             max_steps,
                             Checkpoint::PrepE2EProving,

ceno_zkvm/benches/fibonacci_witness.rs

@@ -1,9 +1,9 @@
 use std::{fs, path::PathBuf, time::Duration};
 
-use ceno_emul::{CENO_PLATFORM, Platform, Program, WORD_SIZE};
+use ceno_emul::{Platform, Program};
 use ceno_zkvm::{
     self,
-    e2e::{Checkpoint, run_e2e_with_checkpoint},
+    e2e::{Checkpoint, Preset, run_e2e_with_checkpoint, setup_platform},
 };
 use criterion::*;
 
@@ -23,29 +23,20 @@ type Pcs = BasefoldDefault<E>;
 type E = GoldilocksExt2;
 
 // Relevant init data for fibonacci run
-fn setup() -> (Program, Platform, u32, u32) {
+fn setup() -> (Program, Platform) {
     let mut file_path = PathBuf::from(env!("CARGO_MANIFEST_DIR"));
     file_path.push("examples/fibonacci.elf");
     let stack_size = 32768;
     let heap_size = 2097152;
+    let pub_io_size = 16;
     let elf_bytes = fs::read(&file_path).expect("read elf file");
     let program = Program::load_elf(&elf_bytes, u32::MAX).unwrap();
-
-    let platform = Platform {
-        // The stack section is not mentioned in ELF headers, so we repeat the constant STACK_TOP here.
-        stack_top: 0x0020_0400,
-        rom: program.base_address
-            ..program.base_address + (program.instructions.len() * WORD_SIZE) as u32,
-        ram: 0x0010_0000..0xFFFF_0000,
-        unsafe_ecall_nop: true,
-        ..CENO_PLATFORM
-    };
-
-    (program, platform, stack_size, heap_size)
+    let platform = setup_platform(Preset::Sp1, &program, stack_size, heap_size, pub_io_size);
+    (program, platform)
 }
 
 fn fibonacci_witness(c: &mut Criterion) {
-    let (program, platform, stack_size, heap_size) = setup();
+    let (program, platform) = setup();
 
     let max_steps = usize::MAX;
     let mut group = c.benchmark_group(format!("fib_wit_max_steps_{}", max_steps));
@@ -63,8 +54,6 @@ fn fibonacci_witness(c: &mut Criterion) {
                     run_e2e_with_checkpoint::<E, Pcs>(
                         program.clone(),
                         platform.clone(),
-                        stack_size,
-                        heap_size,
                         vec![],
                         max_steps,
                         Checkpoint::PrepWitnessGen,

ceno_zkvm/examples/riscv_opcodes.rs

@@ -76,7 +76,7 @@ fn main() {
         program_code,
         Default::default(),
     );
-    let mem_addresses = CENO_PLATFORM.ram.clone();
+    let mem_addresses = CENO_PLATFORM.heap.clone();
     let io_addresses = CENO_PLATFORM.public_io.clone();
 
     let mut fmt_layer = fmt::layer()

ceno_zkvm/src/bin/e2e.rs

@@ -1,9 +1,9 @@
-use ceno_emul::{CENO_PLATFORM, IterAddresses, Platform, Program, WORD_SIZE, Word};
+use ceno_emul::{IterAddresses, Program, WORD_SIZE, Word};
 use ceno_zkvm::{
-    e2e::{Checkpoint, run_e2e_with_checkpoint},
+    e2e::{Checkpoint, Preset, run_e2e_with_checkpoint, setup_platform},
     with_panic_hook,
 };
-use clap::{Parser, ValueEnum};
+use clap::Parser;
 use ff_ext::ff::Field;
 use goldilocks::{Goldilocks, GoldilocksExt2};
 use itertools::Itertools;
@@ -51,19 +51,14 @@ struct Args {
     heap_size: u32,
 }
 
-#[derive(Copy, Clone, Debug, PartialEq, Eq, PartialOrd, Ord, ValueEnum)]
-enum Preset {
-    Ceno,
-    Sp1,
-}
-
 fn main() {
     let args = {
         let mut args = Args::parse();
         args.stack_size = args.stack_size.next_multiple_of(WORD_SIZE as u32);
         args.heap_size = args.heap_size.next_multiple_of(WORD_SIZE as u32);
         args
     };
+    let pub_io_size = 16; // TODO: configure.
 
     // default filter
     let default_filter = EnvFilter::builder()
@@ -99,29 +94,16 @@ fn main() {
         .with(args.profiling.is_none().then_some(default_filter))
         .init();
 
-    let args = {
-        let mut args = Args::parse();
-        args.stack_size = args.stack_size.next_multiple_of(WORD_SIZE as u32);
-        args.heap_size = args.heap_size.next_multiple_of(WORD_SIZE as u32);
-        args
-    };
-
     tracing::info!("Loading ELF file: {}", &args.elf);
     let elf_bytes = fs::read(&args.elf).expect("read elf file");
     let program = Program::load_elf(&elf_bytes, u32::MAX).unwrap();
-
-    let platform = match args.platform {
-        Preset::Ceno => CENO_PLATFORM,
-        Preset::Sp1 => Platform {
-            // The stack section is not mentioned in ELF headers, so we repeat the constant STACK_TOP here.
-            stack_top: 0x0020_0400,
-            rom: program.base_address
-                ..program.base_address + (program.instructions.len() * WORD_SIZE) as u32,
-            ram: 0x0010_0000..0xFFFF_0000,
-            unsafe_ecall_nop: true,
-            ..CENO_PLATFORM
-        },
-    };
+    let platform = setup_platform(
+        args.platform,
+        &program,
+        args.stack_size,
+        args.heap_size,
+        pub_io_size,
+    );
     tracing::info!("Running on platform {:?} {:?}", args.platform, platform);
     tracing::info!(
         "Stack: {} bytes. Heap: {} bytes.",
@@ -146,8 +128,6 @@ fn main() {
     let (state, _) = run_e2e_with_checkpoint::<E, Pcs>(
         program,
         platform,
-        args.stack_size,
-        args.heap_size,
         hints,
         max_steps,
         Checkpoint::PrepSanityCheck,