From 4373f07db90c7f6b42d9403d6bd7cbdb90ade667 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Aur=C3=A9lien=20Nicolas?= <info@nau.re>
Date: Thu, 12 Dec 2024 02:57:02 +0100
Subject: [PATCH] sync-platform: rename rom to prog_code

---
 ceno_emul/src/platform.rs   | 14 +++++++-------
 ceno_emul/tests/test_elf.rs | 13 +++++++++++--
 ceno_zkvm/src/e2e.rs        |  2 +-
 3 files changed, 19 insertions(+), 10 deletions(-)

diff --git a/ceno_emul/src/platform.rs b/ceno_emul/src/platform.rs
index 0f7c5112d..c896f3b58 100644
--- a/ceno_emul/src/platform.rs
+++ b/ceno_emul/src/platform.rs
@@ -9,7 +9,7 @@ use crate::addr::{Addr, RegIdx};
 /// - codes of environment calls.
 #[derive(Clone, Debug)]
 pub struct Platform {
-    pub rom: Range<Addr>, // TODO: rename.
+    pub prog_code: Range<Addr>,
     pub ram: Range<Addr>, // TODO: remove.
     pub prog_data: Option<HashSet<Addr>>,
     pub stack: Range<Addr>,
@@ -21,7 +21,7 @@ pub struct Platform {
 }
 
 pub const CENO_PLATFORM: Platform = Platform {
-    rom: 0x2000_0000..0x3000_0000,
+    prog_code: 0x2000_0000..0x3000_0000,
     ram: 0x8000_0000..0xFFFF_0000,
     prog_data: None, // This is an `Option` to allow `const` here.
     stack: 0xB0000000..0xC0000000,
@@ -35,7 +35,7 @@ impl Platform {
     // Virtual memory layout.
 
     pub fn is_rom(&self, addr: Addr) -> bool {
-        self.rom.contains(&addr)
+        self.prog_code.contains(&addr)
     }
 
     pub fn is_ram(&self, addr: Addr) -> bool {
@@ -48,7 +48,7 @@ impl Platform {
     pub fn is_prog_data(&self, addr: Addr) -> bool {
         self.prog_data
             .as_ref()
-            .map(|set| set.contains(&addr))
+            .map(|set| set.contains(&(addr & !0x3)))
             .unwrap_or(false)
     }
 
@@ -74,7 +74,7 @@ impl Platform {
     // Startup.
 
     pub const fn pc_base(&self) -> Addr {
-        self.rom.start
+        self.prog_code.start
     }
 
     // Permissions.
@@ -131,8 +131,8 @@ mod tests {
         // ROM and RAM do not overlap.
         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));
+        assert!(!p.is_ram(p.prog_code.start));
+        assert!(!p.is_ram(p.prog_code.end - WORD_SIZE as Addr));
         // Registers do not overlap with ROM or RAM.
         for reg in [
             Platform::register_vma(0),
diff --git a/ceno_emul/tests/test_elf.rs b/ceno_emul/tests/test_elf.rs
index bb8addacf..b6e2b6356 100644
--- a/ceno_emul/tests/test_elf.rs
+++ b/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<()> {
@@ -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);
diff --git a/ceno_zkvm/src/e2e.rs b/ceno_zkvm/src/e2e.rs
index 3c2d0c690..0165276c0 100644
--- a/ceno_zkvm/src/e2e.rs
+++ b/ceno_zkvm/src/e2e.rs
@@ -202,7 +202,7 @@ pub fn setup_platform(
     };
 
     Platform {
-        rom: program.base_address
+        prog_code: program.base_address
             ..program.base_address + (program.instructions.len() * WORD_SIZE) as u32,
         prog_data: Some(prog_data),
         stack,