Memory machine: Only use the necessary size (#1680)

Builds on #1666 

Now, the memory machine is downsized to the smallest possible size:
```
$ MAX_DEGREE_LOG=10 cargo run -r pil test_data/asm/dynamic_vadcop.asm -o output -f --field bn254 --prove-with halo2-mock-composite
...
Running main machine for 1024 rows
[00:00:00 (ETA: 00:00:00)] ░░░░░░░░░░░░░░░░░░░░ 0% - Starting...                         Found loop with period 1 starting at row 100
[00:00:00 (ETA: 00:00:00)] ████████████████████ 100% - 280269 rows/s, 4205k identities/s, 94% progress
Resizing variable length machine 'Secondary machine 0: main_arith (BlockMachine)': 1024 -> 32 (rounded up from 3)
Resizing variable length machine 'Secondary machine 1: main_memory (DoubleSortedWitnesses)': 1024 -> 32 (rounded up from 4)
Witness generation took 0.005317875s
Writing output/commits.bin.
Instantiating a composite backend with 5 machines:
...
== Proving machine: main (size 1024), stage 0
==> Proof stage computed in 15.581958ms
== Proving machine: main__rom (size 16), stage 0
==> Proof stage computed in 800.708µs
== Proving machine: main_arith (size 32), stage 0
==> Proof stage computed in 506.25µs
== Proving machine: main_byte2 (size 65536), stage 0
==> Proof stage computed in 84.450916ms
== Proving machine: main_memory (size 32), stage 0
==> Proof stage computed in 719.375µs
Proof generation took 0.102423415s
Proof size: 88 bytes
Writing output/dynamic_vadcop_proof.bin.
```

executor/src/witgen/machines/double_sorted_witness_machine.rs

@@ -4,9 +4,10 @@ use std::iter::once;
 use itertools::Itertools;
 
 use super::Machine;
+use crate::constant_evaluator::{MAX_DEGREE_LOG, MIN_DEGREE_LOG};
 use crate::witgen::rows::RowPair;
 use crate::witgen::util::try_to_simple_poly;
-use crate::witgen::{EvalResult, FixedData, MutableState, QueryCallback};
+use crate::witgen::{EvalError, EvalResult, FixedData, MutableState, QueryCallback};
 use crate::witgen::{EvalValue, IncompleteCause};
 use crate::Identity;
 use powdr_number::{DegreeType, FieldElement};
@@ -52,10 +53,13 @@ pub struct DoubleSortedWitnesses<'a, T: FieldElement> {
     //witness_positions: HashMap<String, usize>,
     /// (addr, step) -> value
     trace: BTreeMap<(T, T), Operation<T>>,
+    /// A map addr -> value, the current content of the memory.
     data: BTreeMap<T, T>,
     is_initialized: BTreeMap<T, bool>,
     namespace: String,
     name: String,
+    /// The set of witness columns that are actually part of this machine.
+    witness_cols: HashSet<PolyID>,
     /// If the machine has the `m_diff_upper` and `m_diff_lower` columns, this is the base of the
     /// two digits.
     diff_columns_base: Option<u64>,
@@ -153,6 +157,7 @@ impl<'a, T: FieldElement> DoubleSortedWitnesses<'a, T> {
                 let diff_columns_base = Some(max.to_degree() + 1);
                 Some(Self {
                     name,
+                    witness_cols: witness_cols.clone(),
                     namespace,
                     fixed: fixed_data,
                     degree,
@@ -170,6 +175,7 @@ impl<'a, T: FieldElement> DoubleSortedWitnesses<'a, T> {
         } else {
             Some(Self {
                 name,
+                witness_cols: witness_cols.clone(),
                 namespace,
                 fixed: fixed_data,
                 degree,
@@ -263,6 +269,22 @@ impl<'a, T: FieldElement> Machine<'a, T> for DoubleSortedWitnesses<'a, T> {
             is_bootloader_write.push(0.into());
             set_selector(None);
         }
+
+        if self.fixed.is_variable_size(&self.witness_cols) {
+            let current_size = addr.len();
+            assert!(current_size <= 1 << *MAX_DEGREE_LOG);
+            let new_size = current_size.next_power_of_two() as DegreeType;
+            let new_size = new_size.max(1 << MIN_DEGREE_LOG);
+            log::info!(
+                "Resizing variable length machine '{}': {} -> {} (rounded up from {})",
+                self.name,
+                self.degree,
+                new_size,
+                current_size
+            );
+            self.degree = new_size;
+        }
+
         while addr.len() < self.degree as usize {
             addr.push(*addr.last().unwrap());
             step.push(*step.last().unwrap() + T::from(1));
@@ -456,6 +478,10 @@ impl<'a, T: FieldElement> DoubleSortedWitnesses<'a, T> {
             assignments = assignments.report_side_effect();
         }
 
+        if self.trace.len() >= (self.degree as usize) {
+            return Err(EvalError::RowsExhausted(self.name.clone()));
+        }
+
         Ok(assignments)
     }
 }

pipeline/tests/asm.rs

@@ -1,3 +1,5 @@
+use std::collections::BTreeMap;
+
 use powdr_backend::BackendType;
 use powdr_executor::constant_evaluator::get_uniquely_sized;
 use powdr_number::{Bn254Field, FieldElement, GoldilocksField};
@@ -201,17 +203,26 @@ fn vm_to_block_array() {
 }
 
 #[test]
-fn vm_to_block_different_length() {
-    let f = "asm/vm_to_block_different_length.asm";
+fn dynamic_vadcop() {
+    let f = "asm/dynamic_vadcop.asm";
+
+    let mut pipeline_gl = make_simple_prepared_pipeline(f);
+    let witness = pipeline_gl.compute_witness().unwrap();
+    let witness_by_name = witness
+        .iter()
+        .map(|(k, v)| (k.as_str(), v))
+        .collect::<BTreeMap<_, _>>();
+
+    // Spot-check some witness columns to have the expected length.
+    assert_eq!(witness_by_name["main.X"].len(), 128);
+    assert_eq!(witness_by_name["main_arith.y"].len(), 32);
+    assert_eq!(witness_by_name["main_memory.m_addr"].len(), 32);
+
     // Because machines have different lengths, this can only be proven
     // with a composite proof.
-    run_pilcom_with_backend_variant(make_simple_prepared_pipeline(f), BackendVariant::Composite)
-        .unwrap();
+    run_pilcom_with_backend_variant(pipeline_gl.clone(), BackendVariant::Composite).unwrap();
+    gen_estark_proof_with_backend_variant(pipeline_gl, BackendVariant::Composite);
     test_halo2_with_backend_variant(make_simple_prepared_pipeline(f), BackendVariant::Composite);
-    gen_estark_proof_with_backend_variant(
-        make_simple_prepared_pipeline(f),
-        BackendVariant::Composite,
-    );
 }
 
 #[test]

test_data/asm/vm_to_block_different_length.asm → test_data/asm/dynamic_vadcop.asm

@@ -1,18 +1,46 @@
+use std::machines::memory::Memory;
+
+// Copy of std::machines::range::Byte2 which sets the correct degree.
+machine Byte2 with
+    latch: latch,
+    operation_id: operation_id,
+    degree: 65536
+{
+    operation check<0> BYTE2 -> ;
+
+    let BYTE2: col = |i| i & 0xffff;
+    col fixed latch = [1]*;
+    col fixed operation_id = [0]*;
+}
+
 machine Main {
     Arith arith;
 
+    col fixed STEP(i) { i };
+    Byte2 byte2;
+    Memory memory(byte2);
+
     reg pc[@pc];
     reg X[<=];
     reg Y[<=];
     reg A;
+    reg B;
     reg Z[<=]; // we declare this assignment register last to test that the ordering does not matter
 
     instr add X, Y -> Z link => Z = arith.add(X, Y);
     instr mul X, Y -> Z link => Z = arith.mul(X, Y);
+    instr mload X -> Y link ~> Y = memory.mload(X, STEP);
+    instr mstore X, Y -> link ~> memory.mstore(X, STEP, Y);
     instr assert_eq X, Y { X = Y }
 
     function main {
-        A <== add(2, 1);
+        mstore 0, 1;
+        mstore 1, 2;
+
+        A <== mload(0);
+        B <== mload(1);
+
+        A <== add(A, B);
         A <== mul(A, 9);
         assert_eq A, 27;
         return;