More efficient estimate of memory-checked RAM size (#616)

* More efficient estimate of memory-checked RAM size

* Silence clippy

* Use WORD_SIZE

Author: Yoichi Hirai

prover/src/machine.rs

@@ -126,13 +126,9 @@ impl<C: MachineChip + Sync> Machine<C> {
     ) -> Result<Proof, ProvingError> {
         let num_steps = trace.get_num_steps();
         let program_len = view.get_program_memory().program.len();
-        let memory_len = view.get_initial_memory().len()
-            + view.get_exit_code().len()
-            + view.get_public_output().len();
+        let tracked_ram_size = view.view_tracked_ram_size();
 
-        // [memory_len + 4 * num_steps] is the union bound for the size of touched-or-initialized memory
-        // The STARK trace needs to contain [num_steps], too, but it is never bigger than the above.
-        let log_size = Self::max_log_size(&[program_len, memory_len + 4 * num_steps])
+        let log_size = Self::max_log_size(&[num_steps, program_len, tracked_ram_size])
             .max(PreprocessedTraces::MIN_LOG_SIZE);
 
         let extensions_iter = BASE_EXTENSIONS.iter().chain(extensions);

sdk/Cargo.toml

@@ -12,6 +12,7 @@ categories = { workspace = true }
 [dependencies]
 serde.workspace = true
 
+nexus-common = { path = "../common" }
 nexus-core = { path = "../core" }
 nexus-sdk-macros = { path = "./macros" }
 

sdk/src/traits.rs

@@ -1,5 +1,6 @@
 use crypto::digest::{Digest, OutputSizeUser};
 use crypto_common::generic_array::{ArrayLength, GenericArray};
+use nexus_common::constants::WORD_SIZE;
 use serde::{de::DeserializeOwned, Serialize};
 use std::path::Path;
 
@@ -88,11 +89,15 @@ impl CheckedView for nexus_core::nvm::View {
             expected_public_output,
         );
 
+        let static_memory_size =
+            (&expected_elf.rom_image.len() + &expected_elf.ram_image.len()) * WORD_SIZE;
+
         Self::new(
             &Some(*memory_layout),
             &Vec::new(),
             &program_memory,
             &initial_memory,
+            memory_layout.tracked_ram_size(static_memory_size),
             &exit_code,
             &output_memory,
             &expected_ad.to_vec(),

vm/src/emulator/executor.rs

@@ -745,6 +745,12 @@ impl Emulator for HarvardEmulator {
             Vec::new()
         };
 
+        let input_size = rom_iter.len() + self.static_ram_image.len() + public_input.len();
+        let tracked_ram_size = self
+            .memory_stats
+            .get_tracked_ram_size(input_size as u32, output_memory.len() as u32)
+            as usize;
+
         View {
             memory_layout: None,
             debug_logs,
@@ -766,6 +772,7 @@ impl Emulator for HarvardEmulator {
                 .chain(ram_iter)
                 .chain(public_input)
                 .collect(),
+            tracked_ram_size,
             exit_code,
             output_memory,
             associated_data: Vec::new(),
@@ -1227,6 +1234,10 @@ impl Emulator for LinearEmulator {
             )
             .unwrap_or_default();
 
+        let tracked_ram_size = self
+            .memory_layout
+            .tracked_ram_size(self.initial_static_ram_image.len() + rom_initialization.len());
+
         View {
             memory_layout: Some(self.memory_layout),
             debug_logs,
@@ -1256,6 +1267,7 @@ impl Emulator for LinearEmulator {
                 .chain(ram_iter)
                 .chain(public_input_iter)
                 .collect(),
+            tracked_ram_size,
             exit_code,
             output_memory,
             associated_data,

vm/src/emulator/layout.rs

@@ -294,4 +294,38 @@ impl LinearMemoryLayout {
     pub fn stack_top(&self) -> u32 {
         self.stack_top - WORD_SIZE as u32
     }
+
+    pub fn tracked_ram_size(&self, static_memory_size: usize) -> usize {
+        let stack_size: usize =
+            self.stack_top
+                .checked_sub(self.stack_bottom)
+                .expect("stack top should be above stack bottom") as usize;
+        let heap_size = self
+            .heap
+            .checked_sub(self.heap_start())
+            .expect("heap should be above heap start") as usize;
+        let public_input_size =
+            self.public_input_end()
+                .checked_sub(self.public_input_start())
+                .expect("public input end should be above public input start") as usize;
+        let public_output_size = self
+            .public_output_end()
+            .checked_sub(self.public_output_start())
+            .expect("public output end should be above public output start")
+            as usize;
+        let exit_code_size = WORD_SIZE;
+        // program, registers, and ad are not under RAM checking. omitted.
+        let total = [
+            static_memory_size,
+            stack_size,
+            heap_size,
+            public_input_size,
+            public_output_size,
+            exit_code_size,
+        ]
+        .iter()
+        .try_fold(0usize, |acc, &val| acc.checked_add(val))
+        .expect("overflow");
+        total
+    }
 }

vm/src/emulator/memory_stats.rs

@@ -115,6 +115,17 @@ impl MemoryStats {
             ad_size,
         )
     }
+
+    /// Returns the total number of addresses under RAM memory checking.
+    pub fn get_tracked_ram_size(&self, input_size: u32, output_size: u32) -> u32 {
+        let heap_size = self.max_heap_access - self.heap_bottom;
+        let stack_size = self.stack_top - self.min_stack_access;
+        let total = [heap_size, stack_size, input_size, output_size]
+            .iter()
+            .try_fold(0u32, |acc, &val| acc.checked_add(val))
+            .expect("overflow");
+        total
+    }
 }
 
 #[cfg(test)]

vm/src/emulator/utils.rs

@@ -189,6 +189,8 @@ pub struct View {
     pub(crate) debug_logs: Vec<Vec<u8>>,
     pub(crate) program_memory: ProgramInfo,
     pub(crate) initial_memory: Vec<MemoryInitializationEntry>,
+    /// The number of all addresses under RAM memory checking
+    pub(crate) tracked_ram_size: usize,
     pub(crate) exit_code: Vec<PublicOutputEntry>,
     pub(crate) output_memory: Vec<PublicOutputEntry>,
     // todo: incorporate into initial memory
@@ -197,11 +199,13 @@ pub struct View {
 
 impl View {
     /// Construct a view out of its raw parts.
+    #[allow(clippy::too_many_arguments)] // extra thought needed what's the best approach to reduce args
     pub fn new(
         memory_layout: &Option<LinearMemoryLayout>,
         debug_logs: &Vec<Vec<u8>>,
         program_memory: &ProgramInfo,
         initial_memory: &Vec<MemoryInitializationEntry>,
+        tracked_ram_size: usize,
         exit_code: &Vec<PublicOutputEntry>,
         output_memory: &Vec<PublicOutputEntry>,
         associated_data: &Vec<u8>,
@@ -211,6 +215,7 @@ impl View {
             debug_logs: debug_logs.to_owned(),
             program_memory: program_memory.to_owned(),
             initial_memory: initial_memory.to_owned(),
+            tracked_ram_size,
             exit_code: exit_code.to_owned(),
             output_memory: output_memory.to_owned(),
             associated_data: associated_data.to_owned(),
@@ -247,6 +252,11 @@ impl View {
             .map(|layout| io_entries_into_vec(layout.public_output_start(), &self.output_memory))
     }
 
+    /// Return the number of all addresses under RAM memory checking.
+    pub fn view_tracked_ram_size(&self) -> usize {
+        self.tracked_ram_size
+    }
+
     /// Return the raw bytes of the associated data, if any.
     pub fn view_associated_data(&self) -> Option<Vec<u8>> {
         if self.memory_layout.is_some() {