Added pc to MemoryImage (risc0#534)

* Added program counter to MemoryImage * Removed PC from other structures and just use image pc
Auradine · Apr 26, 2023 · d44aeb4 · d44aeb4
1 parent 7a04a13
commit d44aeb4
Show file tree

Hide file tree

Showing 4 changed files with 14 additions and 11 deletions.
diff --git a/risc0/zkvm/src/binfmt/image.rs b/risc0/zkvm/src/binfmt/image.rs
@@ -112,6 +112,9 @@ pub struct MemoryImage {
 
     /// Metadata about the structure of the page table
     pub info: PageTableInfo,
+
+    /// Program Counter from [Program] entry point
+    pub pc: u32,
 }
 
 impl MemoryImage {
@@ -134,7 +137,11 @@ impl MemoryImage {
 
         // Compute the page table hashes except for the very last root hash.
         let info = PageTableInfo::new(PAGE_TABLE.start() as u32, page_size);
-        let mut img = Self { buf, info };
+        let mut img = Self {
+            buf,
+            info,
+            pc: program.entry,
+        };
         img.hash_pages();
         Ok(img)
     }
@@ -229,7 +236,10 @@ mod tests {
     fn check_integrity() {
         const PAGE_SIZE: u32 = 1024;
         let program = Program::load_elf(MULTI_TEST_ELF, TEXT.end() as u32).unwrap();
+        let prog_pc = program.entry;
         let image = MemoryImage::new(&program, PAGE_SIZE).unwrap();
+        assert_eq!(image.pc, prog_pc);
+
         // This is useful in case one needs to manually inspect the memory image.
         // std::fs::write("/tmp/test.img", &image.image).unwrap();
         image.check(STACK.start() as u32).unwrap();

diff --git a/risc0/zkvm/src/exec/mod.rs b/risc0/zkvm/src/exec/mod.rs
@@ -66,7 +66,6 @@ pub struct Executor<'a> {
     env: ExecutorEnv<'a>,
     pre_image: MemoryImage,
     monitor: MemoryMonitor,
-    pre_pc: u32,
     pc: u32,
     init_cycles: usize,
     fini_cycles: usize,
@@ -135,7 +134,6 @@ impl<'a> Executor<'a> {
             env,
             pre_image,
             monitor,
-            pre_pc: pc,
             pc,
             init_cycles,
             fini_cycles,
@@ -178,7 +176,6 @@ impl<'a> Executor<'a> {
                     self.segments.push(Segment::new(
                         pre_image,
                         post_image_id,
-                        self.pre_pc,
                         faults,
                         syscalls,
                         exit_code,
@@ -216,7 +213,7 @@ impl<'a> Executor<'a> {
         self.body_cycles = 0;
         self.insn_counter = 0;
         self.segment_cycle = self.init_cycles;
-        self.pre_pc = self.pc;
+        self.pre_image.pc = self.pc;
         self.monitor.clear_segment();
     }
 

diff --git a/risc0/zkvm/src/prove/mod.rs b/risc0/zkvm/src/prove/mod.rs
@@ -368,10 +368,10 @@ impl Segment {
 
     fn prepare_globals(&self) -> Vec<Elem> {
         let mut io = vec![Elem::INVALID; CircuitImpl::OUTPUT_SIZE];
-        log::debug!("run> pc: 0x{:08x}", self.pc);
+        log::debug!("run> pc: 0x{:08x}", self.pre_image.pc);
 
         // initialize PC
-        let pc_bytes = self.pc.to_le_bytes();
+        let pc_bytes = self.pre_image.pc.to_le_bytes();
         for i in 0..WORD_SIZE {
             io[i] = (pc_bytes[i] as u32).into();
         }

diff --git a/risc0/zkvm/src/session.rs b/risc0/zkvm/src/session.rs
@@ -80,7 +80,6 @@ pub struct Session {
 pub struct Segment {
     pub(crate) pre_image: MemoryImage,
     pub(crate) post_image_id: Digest,
-    pub(crate) pc: u32,
     pub(crate) faults: PageFaults,
     pub(crate) syscalls: Vec<SyscallRecord>,
     pub(crate) exit_code: ExitCode,
@@ -105,11 +104,9 @@ impl Session {
 
 impl Segment {
     /// Create a new [Segment] from its constituent components.
-    #[allow(clippy::too_many_arguments)]
     pub(crate) fn new(
         pre_image: MemoryImage,
         post_image_id: Digest,
-        pc: u32,
         faults: PageFaults,
         syscalls: Vec<SyscallRecord>,
         exit_code: ExitCode,
@@ -119,7 +116,6 @@ impl Segment {
         Self {
             pre_image,
             post_image_id,
-            pc,
             faults,
             syscalls,
             exit_code,