Write-Once Memory

compiler/src/verify.rs

@@ -10,7 +10,12 @@ use std::{
 
 use crate::compile_asm_string;
 
-pub fn verify_asm_string<T: FieldElement>(file_name: &str, contents: &str, inputs: Vec<T>) {
+pub fn verify_asm_string<T: FieldElement>(
+    file_name: &str,
+    contents: &str,
+    inputs: Vec<T>,
+    external_witness_values: Vec<(&str, Vec<T>)>,
+) {
     let temp_dir = mktemp::Temp::new_dir().unwrap();
     let (_, result) = compile_asm_string(
         file_name,
@@ -20,7 +25,7 @@ pub fn verify_asm_string<T: FieldElement>(file_name: &str, contents: &str, input
         &temp_dir,
         true,
         Some(BackendType::PilStarkCli),
-        vec![],
+        external_witness_values,
     )
     .unwrap();
 

compiler/tests/asm.rs

@@ -4,6 +4,14 @@ use std::fs;
 use test_log::test;
 
 fn verify_asm<T: FieldElement>(file_name: &str, inputs: Vec<T>) {
+    verify_asm_with_external_witness(file_name, inputs, vec![]);
+}
+
+fn verify_asm_with_external_witness<T: FieldElement>(
+    file_name: &str,
+    inputs: Vec<T>,
+    external_witness_values: Vec<(&str, Vec<T>)>,
+) {
     let file_name = format!(
         "{}/../test_data/asm/{file_name}",
         env!("CARGO_MANIFEST_DIR")
@@ -11,7 +19,7 @@ fn verify_asm<T: FieldElement>(file_name: &str, inputs: Vec<T>) {
 
     let contents = fs::read_to_string(&file_name).unwrap();
 
-    verify_asm_string(&file_name, &contents, inputs)
+    verify_asm_string(&file_name, &contents, inputs, external_witness_values);
 }
 
 fn gen_estark_proof(file_name: &str, inputs: Vec<GoldilocksField>) {
@@ -71,6 +79,24 @@ fn secondary_block_machine_add2() {
     gen_estark_proof(f, vec![]);
 }
 
+#[test]
+fn mem_write_once() {
+    let f = "mem_write_once.asm";
+    verify_asm::<GoldilocksField>(f, vec![]);
+    gen_halo2_proof(f, vec![]);
+    gen_estark_proof(f, vec![]);
+}
+
+#[test]
+fn mem_write_once_external_write() {
+    let f = "mem_write_once_external_write.asm";
+    let mut mem = vec![GoldilocksField::from(0); 256];
+    mem[17] = GoldilocksField::from(42);
+    mem[62] = GoldilocksField::from(123);
+    mem[255] = GoldilocksField::from(-1);
+    verify_asm_with_external_witness::<GoldilocksField>(f, vec![], vec![("main.v", mem)]);
+}
+
 #[test]
 fn block_machine_cache_miss() {
     let f = "block_machine_cache_miss.asm";

compiler/tests/powdr_std.rs

@@ -11,7 +11,7 @@ fn verify_asm<T: FieldElement>(file_name: &str, inputs: Vec<T>) {
 
     let contents = fs::read_to_string(&file_name).unwrap();
 
-    verify_asm_string(&file_name, &contents, inputs)
+    verify_asm_string(&file_name, &contents, inputs, vec![])
 }
 
 fn gen_estark_proof(file_name: &str, inputs: Vec<GoldilocksField>) {

executor/src/witgen/machines/block_machine.rs

@@ -15,6 +15,7 @@ use crate::witgen::{machines::Machine, EvalError, EvalValue, IncompleteCause};
 use crate::witgen::{MutableState, QueryCallback};
 use ast::analyzed::{
     AlgebraicExpression as Expression, AlgebraicReference, Identity, IdentityKind, PolyID,
+    PolynomialType,
 };
 use ast::parsed::SelectedExpressions;
 use number::{DegreeType, FieldElement};
@@ -72,6 +73,17 @@ impl<'a, T: FieldElement> BlockMachine<'a, T> {
         global_range_constraints: &GlobalConstraints<T>,
     ) -> Option<Self> {
         for id in connecting_identities {
+            for r in id.right.expressions.iter() {
+                if let Some(poly) = try_to_simple_poly(r) {
+                    if poly.poly_id.ptype == PolynomialType::Constant {
+                        // It does not really make sense to have constant polynomials on the RHS
+                        // of a block machine lookup, as all constant polynomials are periodic, so
+                        // it would always return the same value.
+                        return None;
+                    }
+                }
+            }
+
             // TODO we should check that the other constraints/fixed columns are also periodic.
             if let Some(block_size) = try_to_period(&id.right.selector, fixed_data) {
                 assert!(block_size <= fixed_data.degree as usize);

executor/src/witgen/machines/machine_extractor.rs

@@ -8,6 +8,7 @@ use super::FixedData;
 use super::KnownMachine;
 use crate::witgen::generator::Generator;
 use crate::witgen::global_constraints::GlobalConstraints;
+use crate::witgen::machines::write_once_memory::WriteOnceMemory;
 use ast::analyzed::{AlgebraicExpression as Expression, Identity, IdentityKind, PolyID};
 use ast::parsed::visitor::ExpressionVisitable;
 use ast::parsed::SelectedExpressions;
@@ -107,6 +108,11 @@ pub fn split_out_machines<'a, T: FieldElement>(
         {
             log::info!("Detected machine: memory");
             machines.push(KnownMachine::DoubleSortedWitnesses(machine));
+        } else if let Some(machine) =
+            WriteOnceMemory::try_new(fixed, &connecting_identities, &machine_identities)
+        {
+            log::info!("Detected machine: write-once memory");
+            machines.push(KnownMachine::WriteOnceMemory(machine));
         } else if let Some(machine) = BlockMachine::try_new(
             fixed,
             &connecting_identities,

executor/src/witgen/machines/mod.rs

@@ -9,6 +9,7 @@ use self::block_machine::BlockMachine;
 use self::double_sorted_witness_machine::DoubleSortedWitnesses;
 pub use self::fixed_lookup_machine::FixedLookup;
 use self::sorted_witness_machine::SortedWitnesses;
+use self::write_once_memory::WriteOnceMemory;
 use ast::analyzed::IdentityKind;
 
 use super::affine_expression::AffineExpression;
@@ -23,6 +24,7 @@ mod double_sorted_witness_machine;
 mod fixed_lookup_machine;
 pub mod machine_extractor;
 mod sorted_witness_machine;
+mod write_once_memory;
 
 /// A machine is a set of witness columns and identities where the columns
 /// are used on the right-hand-side of lookups. It can process plookups.
@@ -54,6 +56,7 @@ pub trait Machine<'a, T: FieldElement>: Send + Sync {
 pub enum KnownMachine<'a, T: FieldElement> {
     SortedWitnesses(SortedWitnesses<'a, T>),
     DoubleSortedWitnesses(DoubleSortedWitnesses<T>),
+    WriteOnceMemory(WriteOnceMemory<'a, T>),
     BlockMachine(BlockMachine<'a, T>),
     Vm(Generator<'a, T>),
 }
@@ -71,6 +74,7 @@ impl<'a, T: FieldElement> Machine<'a, T> for KnownMachine<'a, T> {
             KnownMachine::DoubleSortedWitnesses(m) => {
                 m.process_plookup(mutable_state, kind, left, right)
             }
+            KnownMachine::WriteOnceMemory(m) => m.process_plookup(mutable_state, kind, left, right),
             KnownMachine::BlockMachine(m) => m.process_plookup(mutable_state, kind, left, right),
             KnownMachine::Vm(m) => m.process_plookup(mutable_state, kind, left, right),
         }
@@ -88,6 +92,9 @@ impl<'a, T: FieldElement> Machine<'a, T> for KnownMachine<'a, T> {
             KnownMachine::DoubleSortedWitnesses(m) => {
                 m.take_witness_col_values(fixed_lookup, query_callback)
             }
+            KnownMachine::WriteOnceMemory(m) => {
+                m.take_witness_col_values(fixed_lookup, query_callback)
+            }
             KnownMachine::BlockMachine(m) => {
                 m.take_witness_col_values(fixed_lookup, query_callback)
             }

executor/src/witgen/machines/write_once_memory.rs

@@ -0,0 +1,234 @@
+use std::collections::{BTreeMap, HashMap};
+
+use ast::{
+    analyzed::{
+        AlgebraicExpression as Expression, AlgebraicReference, Identity, IdentityKind, PolyID,
+        PolynomialType,
+    },
+    parsed::SelectedExpressions,
+};
+use number::{DegreeType, FieldElement};
+
+use crate::witgen::{
+    affine_expression::AffineExpression, util::try_to_simple_poly, EvalError, EvalResult,
+    EvalValue, FixedData, IncompleteCause, MutableState, QueryCallback,
+};
+
+use super::{FixedLookup, Machine};
+
+/// A memory machine with a fixed address space, and each address can only have one
+/// value during the lifetime of the program.
+/// In the simplest case, it looks like this:
+/// ```pil
+/// let ADDR = |i| i;
+/// let v;
+/// // Stores a value, fails if the cell already has a value that's different
+/// instr mstore X, Y -> { {X, Y} in {ADDR, v} }
+/// // Loads a value. If the cell is empty, the prover can choose a value.
+/// // Note that this is the same lookup, only Y is considered an output instead
+/// // of an input.
+/// instr mload X -> Y { {X, Y} in {ADDR, v} }
+/// ```
+pub struct WriteOnceMemory<'a, T: FieldElement> {
+    /// The fixed data
+    fixed_data: &'a FixedData<'a, T>,
+    /// The right-hand side of the connecting identity
+    /// (if there are several, they must all be the same)
+    rhs: &'a SelectedExpressions<Expression<T>>,
+    /// The polynomials that are used as values (witness polynomials on the RHS)
+    value_polys: Vec<PolyID>,
+    /// A map from keys to row indices
+    key_to_index: BTreeMap<Vec<T>, DegreeType>,
+    /// The memory content
+    data: BTreeMap<DegreeType, Vec<Option<T>>>,
+}
+
+impl<'a, T: FieldElement> WriteOnceMemory<'a, T> {
+    pub fn try_new(
+        fixed_data: &'a FixedData<'a, T>,
+        connecting_identities: &[&'a Identity<Expression<T>>],
+        identities: &[&Identity<Expression<T>>],
+    ) -> Option<Self> {
+        if !identities.is_empty() {
+            return None;
+        }
+
+        let rhs = &connecting_identities[0].right;
+        if !connecting_identities.iter().all(|i| i.right == *rhs) {
+            return None;
+        }
+
+        if rhs.selector.is_some() {
+            return None;
+        }
+
+        let rhs_polys = rhs
+            .expressions
+            .iter()
+            .map(|e| try_to_simple_poly(e))
+            .collect::<Option<Vec<_>>>();
+
+        // Returns None if any RHS polynomial is a complex expression
+        let rhs_polys = rhs_polys?;
+
+        // Build a Vec<PolyID> for the key and value polynomials
+        let (key_polys, value_polys): (Vec<_>, Vec<_>) = rhs_polys
+            .into_iter()
+            .partition(|p| p.poly_id.ptype == PolynomialType::Constant);
+        let key_polys = key_polys
+            .into_iter()
+            .map(|p| {
+                assert!(!p.next);
+                p.poly_id
+            })
+            .collect::<Vec<_>>();
+        let value_polys = value_polys
+            .into_iter()
+            .map(|p| {
+                assert!(!p.next);
+                p.poly_id
+            })
+            .collect::<Vec<_>>();
+
+        let mut key_to_index = BTreeMap::new();
+        for row in 0..fixed_data.degree {
+            let key = key_polys
+                .iter()
+                .map(|k| fixed_data.fixed_cols[k].values[row as usize])
+                .collect::<Vec<_>>();
+            if key_to_index.insert(key, row).is_some() {
+                // Duplicate keys, can't be a write-once memory
+                return None;
+            }
+        }
+
+        Some(Self {
+            fixed_data,
+            rhs,
+            value_polys,
+            key_to_index,
+            data: BTreeMap::new(),
+        })
+    }
+
+    fn process_plookup_internal(
+        &mut self,
+        left: &[AffineExpression<&'a AlgebraicReference, T>],
+        right: &'a SelectedExpressions<Expression<T>>,
+    ) -> EvalResult<'a, T> {
+        let (key_expressions, value_expressions): (Vec<_>, Vec<_>) = left
+            .iter()
+            .zip(right.expressions.iter())
+            .partition(|(_, r)| {
+                try_to_simple_poly(r).unwrap().poly_id.ptype == PolynomialType::Constant
+            });
+        let key = key_expressions
+            .into_iter()
+            .map(|(k, _)| k.constant_value())
+            .collect::<Option<Vec<_>>>();
+        let value_expressions = value_expressions
+            .into_iter()
+            .map(|(v, _)| v)
+            .collect::<Vec<_>>();
+        let value = value_expressions
+            .iter()
+            .map(|v| v.constant_value())
+            .collect::<Vec<_>>();
+
+        log::trace!("Key: {:?}", key);
+        log::trace!("Value: {:?}", value);
+
+        let Some(key) = key else {
+            return Ok(EvalValue::incomplete(
+                IncompleteCause::NonConstantRequiredArgument("key"),
+            ));
+        };
+
+        let index = self.key_to_index.get(&key).cloned().ok_or_else(|| {
+            EvalError::from(format!("Key {:?} not found in write-once memory", key))
+        })?;
+
+        // If there is an externally provided memory value, use it
+        let external_witness_value = self
+            .value_polys
+            .iter()
+            .map(|p| self.fixed_data.external_witness(index, p))
+            .collect::<Vec<_>>();
+        let stored_value = match self.data.get(&index) {
+            Some(values) => values
+                .iter()
+                .zip(external_witness_value.iter())
+                .map(|(&stored, &external)| external.or(stored))
+                .collect(),
+            None => external_witness_value,
+        };
+
+        let mut updates = vec![];
+        let values = value_expressions
+            .into_iter()
+            .zip(stored_value.iter())
+            .map(|(l, r)| {
+                match (l.constant_value(), r) {
+                    // No value provided and none stored -> keep value as None
+                    (None, None) => Ok::<Option<T>, EvalError<T>>(None),
+                    // Value provided but none stored -> write, no updates
+                    (Some(l), None) => Ok(Some(l)),
+                    // Value stored -> keep stored value & either update LHS or assert equality
+                    (_, Some(r)) => {
+                        updates.extend((l.clone() - (*r).into()).solve()?.constraints);
+                        Ok(Some(*r))
+                    }
+                }
+            })
+            .collect::<Result<Vec<_>, _>>()?;
+
+        // Write values
+        let is_complete = !values.contains(&None);
+        self.data.insert(index, values);
+
+        match is_complete {
+            true => Ok(EvalValue::complete(updates)),
+            false => Ok(EvalValue::incomplete_with_constraints(
+                updates,
+                IncompleteCause::NonConstantRequiredArgument("value"),
+            )),
+        }
+    }
+}
+
+impl<'a, T: FieldElement> Machine<'a, T> for WriteOnceMemory<'a, T> {
+    fn process_plookup<'b, Q: QueryCallback<T>>(
+        &mut self,
+        _mutable_state: &'b mut MutableState<'a, 'b, T, Q>,
+        kind: IdentityKind,
+        left: &[AffineExpression<&'a AlgebraicReference, T>],
+        right: &'a SelectedExpressions<Expression<T>>,
+    ) -> Option<EvalResult<'a, T>> {
+        (right == self.rhs && kind == IdentityKind::Plookup)
+            .then(|| self.process_plookup_internal(left, right))
+    }
+
+    fn take_witness_col_values<'b, Q: QueryCallback<T>>(
+        &mut self,
+        _fixed_lookup: &'b mut FixedLookup<T>,
+        _query_callback: &'b mut Q,
+    ) -> HashMap<String, Vec<T>> {
+        self.value_polys
+            .iter()
+            .enumerate()
+            .map(|(value_index, poly)| {
+                let column = self.fixed_data.witness_cols[poly]
+                    .external_values
+                    .clone()
+                    .unwrap_or_else(|| {
+                        let mut column = vec![T::zero(); self.fixed_data.degree as usize];
+                        for (row, values) in self.data.iter() {
+                            column[*row as usize] = values[value_index].unwrap_or_default();
+                        }
+                        column
+                    });
+                (self.fixed_data.column_name(poly).to_string(), column)
+            })
+            .collect()
+    }
+}

riscv/tests/riscv.rs

@@ -145,7 +145,7 @@ fn verify_file(case: &str, inputs: Vec<GoldilocksField>, coprocessors: &CoProces
         riscv::compile_rust_to_riscv_asm(&format!("tests/riscv_data/{case}"), &temp_dir);
     let powdr_asm = riscv::compiler::compile(riscv_asm, coprocessors);
 
-    verify_asm_string(&format!("{case}.asm"), &powdr_asm, inputs);
+    verify_asm_string(&format!("{case}.asm"), &powdr_asm, inputs, vec![]);
 }
 
 #[test]

test_data/asm/mem_write_once.asm

@@ -0,0 +1,47 @@
+machine MemReadWrite {
+
+    degree 256;
+
+    reg pc[@pc];
+    reg X1[<=];
+    reg X2[<=];
+    reg Y1[<=];
+    reg Y2[<=];
+    reg A;
+    reg B;
+
+    // Write-once memory with key (ADDR1, ADDR2) and value (v1, v2)
+    let ADDR1 = |i| i;
+    let ADDR2 = |i| i + 1;
+    let v1;
+    let v2;
+    // Stores a value, fails if the cell already has a value that's different
+    instr mstore X1, X2, Y1, Y2 -> { {X1, X2, Y1, Y2} in {ADDR1, ADDR2, v1, v2} }
+    // Loads a value. If the cell is empty, the prover can choose a value.
+    instr mload X1, X2 -> Y1, Y2 { {X1, X2, Y1, Y2} in {ADDR1, ADDR2, v1, v2} }
+
+    instr assert_eq X1, Y1 { X1 = Y1 }
+
+    function main {
+        mstore 17, 18, 42, 43;
+        mstore 62, 63, 123, 1234;
+        mstore 255, 256, -1, -2;
+
+        // Setting the same value twice is not a problem
+        mstore 17, 18, 42, 43;
+
+        A, B <== mload(17, 18);
+        assert_eq A, 42;
+        assert_eq B, 43;
+
+        A, B <== mload(62, 63);
+        assert_eq A, 123;
+        assert_eq B, 1234;
+
+        A, B <== mload(255, 256);
+        assert_eq A, -1;
+        assert_eq B, -2;
+
+        return;
+    }
+}

test_data/asm/mem_write_once_external_write.asm

@@ -0,0 +1,33 @@
+// Very simple write-once memory, but without an mstore operation.
+// As a result, this only works if the content of the `v` column has
+// been provided externally.
+machine MemReadWrite {
+
+    degree 256;
+
+    reg pc[@pc];
+    reg X[<=];
+    reg Y[<=];
+    reg A;
+
+    // Write-once memory
+    let ADDR = |i| i;
+    let v;
+    // Loads a value. If the cell is empty, the prover can choose a value.
+    instr mload X -> Y { {X, Y} in {ADDR, v} }
+
+    instr assert_eq X, Y { X = Y }
+
+    function main {
+        A <== mload(17);
+        assert_eq A, 42;
+
+        A <== mload(62);
+        assert_eq A, 123;
+
+        A <== mload(255);
+        assert_eq A, -1;
+
+        return;
+    }
+}