From 5ba59fff1d4658bbe050565598acc051363a9f56 Mon Sep 17 00:00:00 2001
From: "Chris T." <chris@succinct.xyz>
Date: Tue, 13 Aug 2024 14:12:45 -0700
Subject: [PATCH] fix: memory checkpoints with syscalls (#1284)

---
 core/src/runtime/mod.rs     | 11 ++++++----
 core/src/runtime/syscall.rs |  8 ++++----
 core/src/runtime/utils.rs   | 40 ++++++++++++++++++-------------------
 3 files changed, 31 insertions(+), 28 deletions(-)

diff --git a/core/src/runtime/mod.rs b/core/src/runtime/mod.rs
index 17d6846607..08f47d475a 100644
--- a/core/src/runtime/mod.rs
+++ b/core/src/runtime/mod.rs
@@ -227,7 +227,7 @@ impl<'a> Runtime<'a> {
     }
 
     /// Get the current values of the registers.
-    pub fn registers(&self) -> [u32; 32] {
+    pub fn registers(&mut self) -> [u32; 32] {
         let mut registers = [0; 32];
         for i in 0..32 {
             let addr = Register::from_u32(i as u32) as u32;
@@ -235,13 +235,15 @@ impl<'a> Runtime<'a> {
                 Some(record) => record.value,
                 None => 0,
             };
+            self.touched_memory.insert(addr);
         }
         registers
     }
 
     /// Get the current value of a register.
-    pub fn register(&self, register: Register) -> u32 {
+    pub fn register(&mut self, register: Register) -> u32 {
         let addr = register as u32;
+        self.touched_memory.insert(addr);
         match self.state.memory.get(&addr) {
             Some(record) => record.value,
             None => 0,
@@ -249,7 +251,8 @@ impl<'a> Runtime<'a> {
     }
 
     /// Get the current value of a word.
-    pub fn word(&self, addr: u32) -> u32 {
+    pub fn word(&mut self, addr: u32) -> u32 {
+        self.touched_memory.insert(addr);
         match self.state.memory.get(&addr) {
             Some(record) => record.value,
             None => 0,
@@ -257,7 +260,7 @@ impl<'a> Runtime<'a> {
     }
 
     /// Get the current value of a byte.
-    pub fn byte(&self, addr: u32) -> u8 {
+    pub fn byte(&mut self, addr: u32) -> u8 {
         let word = self.word(addr - addr % 4);
         (word >> ((addr % 4) * 8)) as u8
     }
diff --git a/core/src/runtime/syscall.rs b/core/src/runtime/syscall.rs
index 7a6465ee30..35cff1bc3a 100644
--- a/core/src/runtime/syscall.rs
+++ b/core/src/runtime/syscall.rs
@@ -238,19 +238,19 @@ impl<'a, 'b> SyscallContext<'a, 'b> {
 
     /// Get the current value of a register, but doesn't use a memory record.
     /// This is generally unconstrained, so you must be careful using it.
-    pub fn register_unsafe(&self, register: Register) -> u32 {
+    pub fn register_unsafe(&mut self, register: Register) -> u32 {
         self.rt.register(register)
     }
 
-    pub fn byte_unsafe(&self, addr: u32) -> u8 {
+    pub fn byte_unsafe(&mut self, addr: u32) -> u8 {
         self.rt.byte(addr)
     }
 
-    pub fn word_unsafe(&self, addr: u32) -> u32 {
+    pub fn word_unsafe(&mut self, addr: u32) -> u32 {
         self.rt.word(addr)
     }
 
-    pub fn slice_unsafe(&self, addr: u32, len: usize) -> Vec<u32> {
+    pub fn slice_unsafe(&mut self, addr: u32, len: usize) -> Vec<u32> {
         let mut values = Vec::new();
         for i in 0..len {
             values.push(self.rt.word(addr + i as u32 * 4));
diff --git a/core/src/runtime/utils.rs b/core/src/runtime/utils.rs
index 41e132bccc..7c0ad541e7 100644
--- a/core/src/runtime/utils.rs
+++ b/core/src/runtime/utils.rs
@@ -1,7 +1,6 @@
 use std::io::Write;
 
 use super::{Instruction, Runtime};
-use crate::runtime::Register;
 
 pub const fn align(addr: u32) -> u32 {
     addr - addr % 4
@@ -42,30 +41,31 @@ impl<'a> Runtime<'a> {
 
         // If RUST_LOG is set to "trace", then log the current state of the runtime every cycle.
         let width = 12;
+        let registers = self.registers();
         log::trace!(
             "clk={} [pc=0x{:x?}] {:<width$?} |         x0={:<width$} x1={:<width$} x2={:<width$} x3={:<width$} x4={:<width$} x5={:<width$} x6={:<width$} x7={:<width$} x8={:<width$} x9={:<width$} x10={:<width$} x11={:<width$} x12={:<width$} x13={:<width$} x14={:<width$} x15={:<width$} x16={:<width$} x17={:<width$} x18={:<width$}",
             self.state.global_clk,
             self.state.pc,
             instruction,
-            self.register(Register::X0),
-            self.register(Register::X1),
-            self.register(Register::X2),
-            self.register(Register::X3),
-            self.register(Register::X4),
-            self.register(Register::X5),
-            self.register(Register::X6),
-            self.register(Register::X7),
-            self.register(Register::X8),
-            self.register(Register::X9),
-            self.register(Register::X10),
-            self.register(Register::X11),
-            self.register(Register::X12),
-            self.register(Register::X13),
-            self.register(Register::X14),
-            self.register(Register::X15),
-            self.register(Register::X16),
-            self.register(Register::X17),
-            self.register(Register::X18),
+            registers[0],
+            registers[1],
+            registers[2],
+            registers[3],
+            registers[4],
+            registers[5],
+            registers[6],
+            registers[7],
+            registers[8],
+            registers[9],
+            registers[10],
+            registers[11],
+            registers[12],
+            registers[13],
+            registers[14],
+            registers[15],
+            registers[16],
+            registers[17],
+            registers[18],
         );
 
         if !self.unconstrained && self.state.global_clk % 10_000_000 == 0 {