add an ExecMode::Witness to the riscv executor

cli-rs/src/main.rs

@@ -444,6 +444,7 @@ fn execute<F: FieldElement>(
             pipeline.data_callback().unwrap(),
             &[],
             None,
+            true,
             profiling,
         );
 

riscv-executor/src/lib.rs

@@ -724,7 +724,7 @@ mod builder {
             regs[pc_idx as usize] = PC_INITIAL_VAL.into();
 
             let submachines: HashMap<_, RefCell<Box<dyn Submachine<F>>>> =
-                if let ExecMode::Trace = mode {
+                if let ExecMode::Witness = mode {
                     [
                         (
                             MachineInstance::memory,
@@ -797,7 +797,7 @@ mod builder {
             lookup_args: &[F],
             extra: &[F],
         ) {
-            if let ExecMode::Trace = self.mode {
+            if let ExecMode::Witness = self.mode {
                 self.trace
                     .submachine_ops
                     .get_mut(m as usize)
@@ -814,7 +814,9 @@ mod builder {
             let cols_len = self.trace.known_cols.len() / KnownWitnessCol::count();
 
             // sanity check
-            assert!(self.trace.len <= cols_len);
+            if let ExecMode::Witness = self.mode {
+                assert!(self.trace.len <= cols_len);
+            }
 
             cols_len
         }
@@ -884,33 +886,30 @@ mod builder {
         /// raw set next value of register by register index instead of name
         fn set_reg_idx(&mut self, idx: u16, value: Elem<F>) {
             // Record register write in trace. Only for non-pc, non-assignment registers.
-            if let ExecMode::Trace = self.mode {
+            if let ExecMode::Trace | ExecMode::Witness = self.mode {
                 if idx != self.pc_idx {
                     self.trace.reg_writes[idx as usize] = Some(value.into_fe());
                 }
             }
-
             self.regs[idx as usize] = value;
         }
 
         pub fn set_col_idx(&mut self, col: KnownWitnessCol, idx: usize, value: Elem<F>) {
-            if let ExecMode::Trace = self.mode {
+            if let ExecMode::Witness = self.mode {
                 let idx = (KnownWitnessCol::count() * idx) + col as usize;
                 *self.trace.known_cols.get_mut(idx).unwrap() = value.into_fe();
             }
         }
 
         pub fn set_col(&mut self, col: KnownWitnessCol, value: Elem<F>) {
-            if let ExecMode::Trace = self.mode {
+            if let ExecMode::Witness = self.mode {
                 let idx = (self.trace.known_cols.len() - KnownWitnessCol::count()) + col as usize;
                 *self.trace.known_cols.get_mut(idx).unwrap() = value.into_fe();
             }
         }
 
         pub fn push_row(&mut self, pc: u32) {
-            if let ExecMode::Trace = self.mode {
-                let new_len = self.trace.known_cols.len() + KnownWitnessCol::count();
-                self.trace.known_cols.resize(new_len, F::zero());
+            if let ExecMode::Trace | ExecMode::Witness = self.mode {
                 self.trace.pc_trace.push(pc);
                 self.trace
                     .reg_trace
@@ -925,6 +924,10 @@ mod builder {
                         }
                     });
             }
+            if let ExecMode::Witness = self.mode {
+                let new_len = self.trace.known_cols.len() + KnownWitnessCol::count();
+                self.trace.known_cols.resize(new_len, F::zero());
+            }
         }
 
         /// advance to next row, returns the index to the statement that must be
@@ -953,13 +956,15 @@ mod builder {
         }
 
         pub(crate) fn set_mem(&mut self, addr: u32, val: u32, step: u32, identity_id: u64) {
-            if let ExecMode::Trace = self.mode {
+            if let ExecMode::Witness = self.mode {
                 self.submachine_op(
                     MachineInstance::memory,
                     identity_id,
                     &[1.into(), addr.into(), step.into(), val.into()],
                     &[],
                 );
+            }
+            if let ExecMode::Trace | ExecMode::Witness = self.mode {
                 self.trace.mem_ops.push(MemOperation {
                     row: self.trace.len,
                     kind: MemOperationKind::Write,
@@ -972,13 +977,15 @@ mod builder {
 
         pub(crate) fn get_mem(&mut self, addr: u32, step: u32, identity_id: u64) -> u32 {
             let val = *self.mem.get(&addr).unwrap_or(&0);
-            if let ExecMode::Trace = self.mode {
+            if let ExecMode::Witness = self.mode {
                 self.submachine_op(
                     MachineInstance::memory,
                     identity_id,
                     &[0.into(), addr.into(), step.into(), val.into()],
                     &[],
                 );
+            }
+            if let ExecMode::Trace | ExecMode::Witness = self.mode {
                 self.trace.mem_ops.push(MemOperation {
                     row: self.trace.len,
                     kind: MemOperationKind::Read,
@@ -1024,34 +1031,43 @@ mod builder {
 
             let pil = opt_pil.unwrap();
 
-            let main_degree = {
-                let range = namespace_degree_range(pil, "main");
-                std::cmp::max(
-                    self.main_columns_len().next_power_of_two() as u32,
-                    range.min as u32,
-                )
-            };
-
             let start = Instant::now();
 
             // turn register write operations into witness columns
-            let main_regs = self.trace.generate_registers_trace();
+            let mut cols = self
+                .trace
+                .generate_registers_trace()
+                .into_iter()
+                .collect::<HashMap<_, _>>();
             log::debug!(
                 "Generating register traces took {}s",
                 start.elapsed().as_secs_f64(),
             );
 
-            // This hashmap will be added to until we get the full witness.
-            let mut cols = self.generate_main_columns();
+            let main_degree = {
+                let range = namespace_degree_range(pil, "main");
+                std::cmp::max(self.trace.len.next_power_of_two() as u32, range.min as u32)
+            };
 
-            // add reg columns to trace
-            cols.extend(main_regs);
+            // fill up reg trace to degree
+            cols.values_mut().for_each(|v| {
+                let last = *v.last().unwrap();
+                v.resize(main_degree as usize, last);
+            });
 
-            // sanity check that program columns and main trace have the same length
-            assert_eq!(
-                cols.values().next().unwrap().len(),
-                program_columns[0].1.len(),
-            );
+            // trace mode doesn't generate a full witness
+            if let ExecMode::Trace = self.mode {
+                return Execution {
+                    trace_len: self.trace.len,
+                    memory: self.mem,
+                    memory_accesses: std::mem::take(&mut self.trace.mem_ops),
+                    trace: cols,
+                    register_memory: self.reg_mem.for_bootloader(),
+                };
+            }
+
+            // add reg columns to trace
+            cols.extend(self.generate_main_columns());
 
             // add program columns to main trace
             cols.extend(program_columns);
@@ -1307,7 +1323,7 @@ impl<F: FieldElement> Executor<'_, '_, F> {
     fn init(&mut self) {
         self.step = 4;
 
-        if let ExecMode::Trace = self.mode {
+        if let ExecMode::Witness = self.mode {
             for c in KnownFixedCol::all() {
                 self.cached_fixed_cols
                     .push(self.get_fixed(c.name()).unwrap().clone());
@@ -1384,30 +1400,33 @@ impl<F: FieldElement> Executor<'_, '_, F> {
     /// Gets the identity id for a link associated with a given instruction.
     /// idx is based on the order link appear in the assembly (assumed to be the same in the optimized pil).
     fn instr_link_id(&mut self, instr: Instruction, target: MachineInstance, idx: usize) -> u64 {
-        if let ExecMode::Fast = self.mode {
-            return 0; // we don't care about identity ids in fast mode
+        if let ExecMode::Witness = self.mode {
+            let entries = self
+                .pil_instruction_links
+                .get_mut(instr as usize * MachineInstance::count() + target as usize)
+                .unwrap()
+                .get_or_insert_with(|| {
+                    pil::find_instruction_links(&self.pil_links, instr.flag(), target.namespace())
+                });
+            entries.get(idx).unwrap().id()
+        } else {
+            // we don't care about identity ids in non witness mode
+            0
         }
-
-        let entries = self
-            .pil_instruction_links
-            .get_mut(instr as usize * MachineInstance::count() + target as usize)
-            .unwrap()
-            .get_or_insert_with(|| {
-                pil::find_instruction_links(&self.pil_links, instr.flag(), target.namespace())
-            });
-        entries.get(idx).unwrap().id()
     }
 
     /// Find the identity id of a link.
     fn link_id(&mut self, from: &'static str, target: &'static str, idx: usize) -> u64 {
-        if let ExecMode::Fast = self.mode {
-            return 0; // we don't care about identity ids in fast mode
+        if let ExecMode::Witness = self.mode {
+            let entries = self
+                .pil_other_links
+                .entry((from, target))
+                .or_insert_with(|| pil::find_links(&self.pil_links, from, target));
+            entries.get(idx).unwrap().id()
+        } else {
+            // we don't care about identity ids in fast mode
+            0
         }
-        let entries = self
-            .pil_other_links
-            .entry((from, target))
-            .or_insert_with(|| pil::find_links(&self.pil_links, from, target));
-        entries.get(idx).unwrap().id()
     }
 
     fn exec_instruction(&mut self, name: &str, args: &[Expression]) -> Option<Elem<F>> {
@@ -1420,7 +1439,7 @@ impl<F: FieldElement> Executor<'_, '_, F> {
 
         macro_rules! get_fixed {
             ($name:ident) => {
-                if let ExecMode::Trace = self.mode {
+                if let ExecMode::Witness = self.mode {
                     Elem::Field(
                         self.get_known_fixed(KnownFixedCol::$name, self.proc.get_pc().u() as usize),
                     )
@@ -2816,8 +2835,12 @@ pub struct Execution<F: FieldElement> {
 
 #[derive(Clone, Copy)]
 enum ExecMode {
+    /// Doesn't generate any execution information besides the length of the trace
     Fast,
+    /// Generate traces for memory and powdr asm registers
     Trace,
+    /// Generate the full witness
+    Witness,
 }
 
 /// Execute a Powdr/RISCV assembly program, without generating a witness.
@@ -2854,6 +2877,7 @@ pub fn execute<F: FieldElement>(
     prover_ctx: &Callback<F>,
     bootloader_inputs: &[F],
     max_steps_to_execute: Option<usize>,
+    full_witness: bool,
     profiling: Option<ProfilerOptions>,
 ) -> Execution<F> {
     log::info!("Executing (trace generation)...");
@@ -2866,7 +2890,11 @@ pub fn execute<F: FieldElement>(
         prover_ctx,
         bootloader_inputs,
         max_steps_to_execute.unwrap_or(usize::MAX),
-        ExecMode::Trace,
+        if full_witness {
+            ExecMode::Witness
+        } else {
+            ExecMode::Trace
+        },
         profiling,
     )
 }
@@ -3111,7 +3139,7 @@ fn execute_inner<F: FieldElement>(
 
     log::debug!("Program execution took {}s", start.elapsed().as_secs_f64());
 
-    if let ExecMode::Trace = mode {
+    if let ExecMode::Trace | ExecMode::Witness = mode {
         let sink_id = e.sink_id();
 
         // reset
@@ -3138,12 +3166,14 @@ fn execute_inner<F: FieldElement>(
         e.proc
             .set_col(KnownWitnessCol::_operation_id, sink_id.into());
 
-        let start = Instant::now();
-        program_columns = e.generate_program_columns();
-        log::debug!(
-            "Generating program columns took {}s",
-            start.elapsed().as_secs_f64()
-        );
+        if let ExecMode::Witness = mode {
+            let start = Instant::now();
+            program_columns = e.generate_program_columns();
+            log::debug!(
+                "Generating program columns took {}s",
+                start.elapsed().as_secs_f64()
+            );
+        }
     }
 
     e.proc.finish(opt_pil, program_columns)

riscv/src/continuations.rs

@@ -332,6 +332,7 @@ pub fn rust_continuations_dry_run<F: FieldElement>(
         // we only know them after the full trace has been generated.
         &default_input(&[]),
         None,
+        false,
         profiler_opt,
     );
 
@@ -436,6 +437,7 @@ pub fn rust_continuations_dry_run<F: FieldElement>(
             pipeline.data_callback().unwrap(),
             &bootloader_inputs,
             Some(num_rows),
+            false,
             // profiling was done when full trace was generated
             None,
         );

riscv/tests/common/mod.rs

@@ -71,6 +71,7 @@ pub fn verify_riscv_asm_string<T: FieldElement, S: serde::Serialize + Send + Syn
             pipeline.data_callback().unwrap(),
             &[],
             None,
+            true,
             None,
         );
         pipeline.rollback_from_witness();