Merge pull request #730 from powdr-labs/vm-to-vm-witgen3

VM-to-VM witgen

compiler/tests/asm.rs

@@ -165,7 +165,6 @@ fn vm_to_block_multiple_interfaces() {
 }
 
 #[test]
-#[should_panic = "not implemented"]
 fn vm_to_vm() {
     let f = "vm_to_vm.asm";
     let i = [];
@@ -174,6 +173,33 @@ fn vm_to_vm() {
     gen_estark_proof(f, slice_to_vec(&i));
 }
 
+#[test]
+fn vm_to_vm_dynamic_trace_length() {
+    let f = "vm_to_vm_dynamic_trace_length.asm";
+    let i = [];
+    verify_asm::<GoldilocksField>(f, slice_to_vec(&i));
+    gen_halo2_proof(f, slice_to_vec(&i));
+    gen_estark_proof(f, slice_to_vec(&i));
+}
+
+#[test]
+fn vm_to_vm_to_block() {
+    let f = "vm_to_vm_to_block.asm";
+    let i = [];
+    verify_asm::<GoldilocksField>(f, slice_to_vec(&i));
+    gen_halo2_proof(f, slice_to_vec(&i));
+    gen_estark_proof(f, slice_to_vec(&i));
+}
+
+#[test]
+fn vm_to_vm_to_vm() {
+    let f = "vm_to_vm_to_vm.asm";
+    let i = [];
+    verify_asm::<GoldilocksField>(f, slice_to_vec(&i));
+    gen_halo2_proof(f, slice_to_vec(&i));
+    gen_estark_proof(f, slice_to_vec(&i));
+}
+
 #[test]
 fn test_mem_read_write() {
     let f = "mem_read_write.asm";

compiler/tests/pil.rs

@@ -221,12 +221,13 @@ fn test_single_line_blocks() {
     gen_estark_proof(f, Default::default());
 }
 
-#[test]
-fn test_two_block_machine_functions() {
-    let f = "two_block_machine_functions.pil";
-    verify_pil(f, None);
-    gen_estark_proof(f, Default::default());
-}
+// TODO: Add back once #693 is merged
+// #[test]
+// fn test_two_block_machine_functions() {
+//     let f = "two_block_machine_functions.pil";
+//     verify_pil(f, None);
+//     gen_estark_proof(f, Default::default());
+// }
 
 #[test]
 fn test_fixed_columns() {

executor/src/witgen/block_processor.rs

@@ -27,15 +27,19 @@ pub struct WithoutCalldata;
 pub struct OuterQuery<'a, T: FieldElement> {
     /// A local copy of the left-hand side of the outer query.
     /// This will be mutated while processing the block.
-    left: Left<'a, T>,
+    pub left: Left<'a, T>,
     /// The right-hand side of the outer query.
-    right: &'a SelectedExpressions<Expression<T>>,
+    pub right: &'a SelectedExpressions<Expression<T>>,
 }
 
 impl<'a, T: FieldElement> OuterQuery<'a, T> {
     pub fn new(left: Left<'a, T>, right: &'a SelectedExpressions<Expression<T>>) -> Self {
         Self { left, right }
     }
+
+    pub fn is_complete(&self) -> bool {
+        self.left.iter().all(|l| l.is_constant())
+    }
 }
 
 /// A basic processor that knows how to determine a unique satisfying witness
@@ -343,7 +347,7 @@ mod tests {
     fn name_to_poly_id<T: FieldElement>(fixed_data: &FixedData<T>) -> BTreeMap<String, PolyID> {
         let mut name_to_poly_id = BTreeMap::new();
         for (poly_id, col) in fixed_data.witness_cols.iter() {
-            name_to_poly_id.insert(col.name.clone(), poly_id);
+            name_to_poly_id.insert(col.poly.name.clone(), poly_id);
         }
         for (poly_id, col) in fixed_data.fixed_cols.iter() {
             name_to_poly_id.insert(col.name.clone(), poly_id);

executor/src/witgen/data_structures/finalizable_data.rs

@@ -52,6 +52,10 @@ impl<'a, T: FieldElement> FinalizableData<'a, T> {
         self.data.len()
     }
 
+    pub fn is_empty(&self) -> bool {
+        self.data.is_empty()
+    }
+
     pub fn push(&mut self, row: Row<'a, T>) {
         self.data.push(Entry::InProgress(row));
     }
@@ -86,6 +90,14 @@ impl<'a, T: FieldElement> FinalizableData<'a, T> {
         }
     }
 
+    pub fn last(&self) -> Option<&Row<'a, T>> {
+        match self.data.last() {
+            Some(Entry::InProgress(row)) => Some(row),
+            Some(Entry::Finalized(_, _)) => panic!("Last row already finalized."),
+            None => None,
+        }
+    }
+
     pub fn mutable_row_pair(&mut self, i: usize) -> (&mut Row<'a, T>, &mut Row<'a, T>) {
         let (before, after) = self.data.split_at_mut(i + 1);
         let current = before.last_mut().unwrap();

executor/src/witgen/generator.rs

@@ -7,18 +7,23 @@ use std::collections::{HashMap, HashSet};
 
 use crate::witgen::data_structures::finalizable_data::FinalizableData;
 use crate::witgen::rows::CellValue;
+use crate::witgen::EvalValue;
 
 use super::affine_expression::AffineExpression;
-use super::block_processor::BlockProcessor;
+use super::block_processor::{BlockProcessor, OuterQuery};
 use super::data_structures::column_map::WitnessColumnMap;
 use super::global_constraints::GlobalConstraints;
-use super::machines::Machine;
-
+use super::machines::{FixedLookup, Machine};
 use super::rows::{Row, RowFactory};
 use super::sequence_iterator::{DefaultSequenceIterator, ProcessingSequenceIterator};
 use super::vm_processor::VmProcessor;
 use super::{EvalResult, FixedData, MutableState, QueryCallback};
 
+struct ProcessResult<'a, T: FieldElement> {
+    eval_value: EvalValue<&'a AlgebraicReference, T>,
+    block: FinalizableData<'a, T>,
+}
+
 pub struct Generator<'a, T: FieldElement> {
     fixed_data: &'a FixedData<'a, T>,
     identities: Vec<&'a Identity<Expression<T>>>,
@@ -31,15 +36,63 @@ pub struct Generator<'a, T: FieldElement> {
 impl<'a, T: FieldElement> Machine<'a, T> for Generator<'a, T> {
     fn process_plookup<Q: QueryCallback<T>>(
         &mut self,
-        _mutable_state: &mut MutableState<'a, '_, T, Q>,
+        mutable_state: &mut MutableState<'a, '_, T, Q>,
         _kind: IdentityKind,
-        _left: &[AffineExpression<&'a AlgebraicReference, T>],
-        _right: &'a SelectedExpressions<Expression<T>>,
+        left: &[AffineExpression<&'a AlgebraicReference, T>],
+        right: &'a SelectedExpressions<Expression<T>>,
     ) -> Option<EvalResult<'a, T>> {
-        unimplemented!()
+        if right.selector != self.latch {
+            None
+        } else {
+            log::trace!("Start processing secondary VM '{}'", self.name());
+            log::trace!("Arguments:");
+            for (r, l) in right.expressions.iter().zip(left) {
+                log::trace!("  {r} = {l}");
+            }
+
+            let first_row = self
+                .data
+                .last()
+                .cloned()
+                .unwrap_or_else(|| self.compute_partial_first_row(mutable_state));
+
+            let outer_query = OuterQuery {
+                left: left.to_vec(),
+                right,
+            };
+            let ProcessResult { eval_value, block } =
+                self.process(first_row, 0, mutable_state, Some(outer_query));
+
+            if eval_value.is_complete() {
+                log::trace!("End processing VM '{}' (successfully)", self.name());
+                // Remove the last row of the previous block, as it is the first row of the current
+                // block.
+                self.data.pop();
+                self.data.extend(block);
+            } else {
+                log::trace!("End processing VM '{}' (incomplete)", self.name());
+            }
+            Some(Ok(eval_value))
+        }
     }
 
-    fn take_witness_col_values(&mut self) -> HashMap<String, Vec<T>> {
+    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>> {
+        log::debug!("Finalizing VM: {}", self.name());
+
+        // In this stage, we don't have access to other machines, as they might already be finalized.
+        let mut mutable_state_no_machines = MutableState {
+            fixed_lookup,
+            machines: [].into_iter().into(),
+            query_callback,
+        };
+
+        self.fill_remaining_rows(&mut mutable_state_no_machines);
+        self.fix_first_row();
+
         self.data
             .take_transposed()
             .map(|(id, (values, _))| (self.fixed_data.column_name(&id).to_string(), values))
@@ -66,11 +119,23 @@ impl<'a, T: FieldElement> Generator<'a, T> {
         }
     }
 
-    pub fn run<Q: QueryCallback<T>>(&mut self, mutable_state: &mut MutableState<'a, '_, T, Q>) {
+    pub fn name(&self) -> &str {
+        let first_witness = self.witnesses.iter().next().unwrap();
+        let first_witness_name = self.fixed_data.column_name(first_witness);
+        let namespace = first_witness_name
+            .rfind('.')
+            .map(|idx| &first_witness_name[..idx]);
+
+        // For machines compiled using Powdr ASM we'll always have a namespace, but as a last
+        // resort we'll use the first witness name.
+        namespace.unwrap_or(first_witness_name)
+    }
+
+    /// Runs the machine without any arguments from the first row.
+    pub fn run<'b, Q: QueryCallback<T>>(&mut self, mutable_state: &mut MutableState<'a, 'b, T, Q>) {
+        assert!(self.data.is_empty());
         let first_row = self.compute_partial_first_row(mutable_state);
-        self.data = self.process(first_row, 0, mutable_state);
-        self.fill_remaining_rows(mutable_state);
-        self.fix_first_row();
+        self.data = self.process(first_row, 0, mutable_state, None).block;
     }
 
     fn fill_remaining_rows<Q>(&mut self, mutable_state: &mut MutableState<'a, '_, T, Q>)
@@ -81,9 +146,15 @@ impl<'a, T: FieldElement> Generator<'a, T> {
             assert!(self.latch.is_some());
 
             let first_row = self.data.pop().unwrap();
+            let ProcessResult { block, eval_value } = self.process(
+                first_row,
+                self.data.len() as DegreeType,
+                mutable_state,
+                None,
+            );
+            assert!(eval_value.is_complete());
 
-            self.data
-                .extend(self.process(first_row, self.data.len() as DegreeType, mutable_state));
+            self.data.extend(block);
         }
     }
 
@@ -131,7 +202,8 @@ impl<'a, T: FieldElement> Generator<'a, T> {
         first_row: Row<'a, T>,
         row_offset: DegreeType,
         mutable_state: &mut MutableState<'a, '_, T, Q>,
-    ) -> FinalizableData<'a, T> {
+        outer_query: Option<OuterQuery<'a, T>>,
+    ) -> ProcessResult<'a, T> {
         log::trace!(
             "Running main machine from row {row_offset} with the following initial values in the first row:\n{}", first_row.render_values(false, None)
         );
@@ -147,11 +219,13 @@ impl<'a, T: FieldElement> Generator<'a, T> {
             self.witnesses.clone(),
             data,
             row_factory,
-            self.latch.clone(),
         );
-        let result = processor.run(mutable_state);
-        assert!(result.is_complete(), "Main machine did not complete");
-        processor.finish()
+        if let Some(outer_query) = outer_query {
+            processor = processor.with_outer_query(outer_query);
+        }
+        let eval_value = processor.run(mutable_state);
+        let block = processor.finish();
+        ProcessResult { eval_value, block }
     }
 
     /// At the end of the solving algorithm, we'll have computed the first row twice

executor/src/witgen/machines/block_machine.rs

@@ -1,6 +1,6 @@
 use std::collections::{HashMap, HashSet};
 
-use super::{EvalResult, FixedData};
+use super::{EvalResult, FixedData, FixedLookup};
 use crate::witgen::affine_expression::AffineExpression;
 
 use crate::witgen::block_processor::{BlockProcessor, OuterQuery};
@@ -164,7 +164,11 @@ impl<'a, T: FieldElement> Machine<'a, T> for BlockMachine<'a, T> {
         })
     }
 
-    fn take_witness_col_values(&mut self) -> HashMap<String, Vec<T>> {
+    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>> {
         if self.data.len() < 2 * self.block_size {
             log::warn!(
                 "Filling empty blocks with zeros, because the block machine is never used. \
@@ -253,13 +257,10 @@ impl<'a, T: FieldElement> BlockMachine<'a, T> {
         let namespace = first_witness_name
             .rfind('.')
             .map(|idx| &first_witness_name[..idx]);
-        if let Some(namespace) = namespace {
-            namespace
-        } else {
-            // For machines compiled using Powdr ASM we'll always have a namespace, but as a last
-            // resort we'll use the first witness name.
-            first_witness_name
-        }
+
+        // For machines compiled using Powdr ASM we'll always have a namespace, but as a last
+        // resort we'll use the first witness name.
+        namespace.unwrap_or(first_witness_name)
     }
 
     fn process_plookup_internal<'b, Q: QueryCallback<T>>(

executor/src/witgen/machines/double_sorted_witness_machine.rs

@@ -5,7 +5,7 @@ use ast::parsed::SelectedExpressions;
 use itertools::Itertools;
 use num_traits::Zero;
 
-use super::Machine;
+use super::{FixedLookup, Machine};
 use crate::witgen::affine_expression::AffineExpression;
 use crate::witgen::util::is_simple_poly_of_name;
 use crate::witgen::{EvalResult, FixedData, MutableState, QueryCallback};
@@ -111,7 +111,11 @@ impl<'a, T: FieldElement> Machine<'a, T> for DoubleSortedWitnesses<T> {
         Some(self.process_plookup_internal(left, right))
     }
 
-    fn take_witness_col_values(&mut self) -> HashMap<String, Vec<T>> {
+    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>> {
         let mut addr = vec![];
         let mut step = vec![];
         let mut value = vec![];

executor/src/witgen/machines/mod.rs

@@ -41,7 +41,11 @@ pub trait Machine<'a, T: FieldElement>: Send + Sync {
     ) -> Option<EvalResult<'a, T>>;
 
     /// Returns the final values of the witness columns.
-    fn take_witness_col_values(&mut self) -> HashMap<String, Vec<T>>;
+    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>>;
 }
 
 /// All known implementations of [Machine].
@@ -61,7 +65,7 @@ impl<'a, T: FieldElement> Machine<'a, T> for KnownMachine<'a, T> {
         kind: IdentityKind,
         left: &[AffineExpression<&'a AlgebraicReference, T>],
         right: &'a SelectedExpressions<Expression<T>>,
-    ) -> Option<crate::witgen::EvalResult<'a, T>> {
+    ) -> Option<EvalResult<'a, T>> {
         match self {
             KnownMachine::SortedWitnesses(m) => m.process_plookup(mutable_state, kind, left, right),
             KnownMachine::DoubleSortedWitnesses(m) => {
@@ -72,12 +76,22 @@ impl<'a, T: FieldElement> Machine<'a, T> for KnownMachine<'a, T> {
         }
     }
 
-    fn take_witness_col_values(&mut self) -> std::collections::HashMap<String, Vec<T>> {
+    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>> {
         match self {
-            KnownMachine::SortedWitnesses(m) => m.take_witness_col_values(),
-            KnownMachine::DoubleSortedWitnesses(m) => m.take_witness_col_values(),
-            KnownMachine::BlockMachine(m) => m.take_witness_col_values(),
-            KnownMachine::Vm(m) => m.take_witness_col_values(),
+            KnownMachine::SortedWitnesses(m) => {
+                m.take_witness_col_values(fixed_lookup, query_callback)
+            }
+            KnownMachine::DoubleSortedWitnesses(m) => {
+                m.take_witness_col_values(fixed_lookup, query_callback)
+            }
+            KnownMachine::BlockMachine(m) => {
+                m.take_witness_col_values(fixed_lookup, query_callback)
+            }
+            KnownMachine::Vm(m) => m.take_witness_col_values(fixed_lookup, query_callback),
         }
     }
 }

executor/src/witgen/machines/sorted_witness_machine.rs

@@ -4,8 +4,8 @@ use ast::parsed::SelectedExpressions;
 use itertools::Itertools;
 
 use super::super::affine_expression::AffineExpression;
-use super::Machine;
 use super::{EvalResult, FixedData};
+use super::{FixedLookup, Machine};
 use crate::witgen::{
     expression_evaluator::ExpressionEvaluator, fixed_evaluator::FixedEvaluator,
     symbolic_evaluator::SymbolicEvaluator,
@@ -152,7 +152,12 @@ impl<'a, T: FieldElement> Machine<'a, T> for SortedWitnesses<'a, T> {
 
         Some(self.process_plookup_internal(left, right, rhs))
     }
-    fn take_witness_col_values(&mut self) -> HashMap<String, Vec<T>> {
+
+    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>> {
         let mut result = HashMap::new();
 
         let (mut keys, mut values): (Vec<_>, Vec<_>) =

executor/src/witgen/mod.rs

@@ -131,11 +131,15 @@ impl<'a, 'b, T: FieldElement, Q: QueryCallback<T>> WitnessGenerator<'a, 'b, T, Q
         generator.run(&mut mutable_state);
 
         // Get columns from machines
-        let main_columns = generator.take_witness_col_values();
+        let main_columns = generator
+            .take_witness_col_values(mutable_state.fixed_lookup, mutable_state.query_callback);
         let mut columns = mutable_state
             .machines
             .iter_mut()
-            .flat_map(|m| m.take_witness_col_values().into_iter())
+            .flat_map(|m| {
+                m.take_witness_col_values(mutable_state.fixed_lookup, mutable_state.query_callback)
+                    .into_iter()
+            })
             .chain(main_columns)
             .collect::<BTreeMap<_, _>>();
 
@@ -211,7 +215,7 @@ impl<'a, T: FieldElement> FixedData<'a, T> {
 
     fn column_name(&self, poly_id: &PolyID) -> &str {
         match poly_id.ptype {
-            PolynomialType::Committed => &self.witness_cols[poly_id].name,
+            PolynomialType::Committed => &self.witness_cols[poly_id].poly.name,
             PolynomialType::Constant => &self.fixed_cols[poly_id].name,
             PolynomialType::Intermediate => unimplemented!(),
         }
@@ -238,18 +242,17 @@ impl<'a, T> FixedColumn<'a, T> {
     }
 }
 
-#[derive(Debug)]
-pub struct Query<'a, T> {
-    /// The query expression
-    expr: &'a Expression<T>,
-    /// The polynomial that is referenced by the query
-    poly: AlgebraicReference,
-}
-
 #[derive(Debug)]
 pub struct WitnessColumn<'a, T> {
-    name: String,
-    query: Option<Query<'a, T>>,
+    /// A polynomial reference that points to this column in the "current" row
+    /// (i.e., the "next" flag is set to false).
+    /// This is needed in situations where we want to update a cell when the
+    /// update does not come from an identity (which also has an AlgebraicReference).
+    poly: AlgebraicReference,
+    /// The prover query expression, if any.
+    query: Option<&'a Expression<T>>,
+    /// A list of externally computed witness values, if any.
+    /// The length of this list must be equal to the degree.
     external_values: Option<Vec<T>>,
 }
 
@@ -265,24 +268,17 @@ impl<'a, T> WitnessColumn<'a, T> {
         } else {
             None
         };
-        let name = name.to_string();
-        let query = query.as_ref().map(|callback| {
-            let poly = AlgebraicReference {
-                poly_id: PolyID {
-                    id: id as u64,
-                    ptype: PolynomialType::Committed,
-                },
-                name: name.clone(),
-                next: false,
-                index: None,
-            };
-            Query {
-                poly,
-                expr: callback,
-            }
-        });
+        let poly = AlgebraicReference {
+            poly_id: PolyID {
+                id: id as u64,
+                ptype: PolynomialType::Committed,
+            },
+            name: name.to_string(),
+            next: false,
+            index: None,
+        };
         WitnessColumn {
-            name,
+            poly,
             query,
             external_values,
         }

executor/src/witgen/query_processor.rs

@@ -5,7 +5,7 @@ use ast::{
 };
 use number::FieldElement;
 
-use super::{rows::RowPair, Constraint, EvalValue, FixedData, IncompleteCause, Query};
+use super::{rows::RowPair, Constraint, EvalValue, FixedData, IncompleteCause};
 
 /// Computes value updates that result from a query.
 pub struct QueryProcessor<'a, 'b, T: FieldElement, QueryCallback: Send + Sync> {
@@ -32,8 +32,8 @@ where
         let column = &self.fixed_data.witness_cols[poly_id];
 
         if let Some(query) = column.query.as_ref() {
-            if rows.get_value(&query.poly).is_none() {
-                return self.process_witness_query(query, rows);
+            if rows.get_value(&column.poly).is_none() {
+                return self.process_witness_query(query, &column.poly, rows);
             }
         }
         // Either no query or the value is already known.
@@ -42,19 +42,20 @@ where
 
     fn process_witness_query(
         &mut self,
-        query: &'a Query<'_, T>,
+        query: &'a Expression<T>,
+        poly: &'a AlgebraicReference,
         rows: &RowPair<T>,
     ) -> EvalValue<&'a AlgebraicReference, T> {
-        let query_str = match self.interpolate_query(query.expr, rows) {
+        let query_str = match self.interpolate_query(query, rows) {
             Ok(query) => query,
             Err(incomplete) => return EvalValue::incomplete(incomplete),
         };
         if let Some(value) = (self.query_callback)(&query_str) {
-            EvalValue::complete(vec![(&query.poly, Constraint::Assignment(value))])
+            EvalValue::complete(vec![(poly, Constraint::Assignment(value))])
         } else {
             EvalValue::incomplete(IncompleteCause::NoQueryAnswer(
                 query_str,
-                query.poly.name.to_string(),
+                poly.name.to_string(),
             ))
         }
     }

executor/src/witgen/rows.rs

@@ -13,7 +13,7 @@ use super::{
     global_constraints::{GlobalConstraints, RangeConstraintSet},
     range_constraints::RangeConstraint,
     symbolic_witness_evaluator::{SymoblicWitnessEvaluator, WitnessColumnEvaluator},
-    EvalValue, FixedData,
+    FixedData,
 };
 
 #[derive(Clone, PartialEq, Debug)]
@@ -240,28 +240,6 @@ impl<'row, 'a, T: FieldElement> RowUpdater<'row, 'a, T> {
         self.get_cell_mut(poly).apply_update(c);
     }
 
-    /// Applies the updates to the underlying rows. Returns true if any updates
-    /// were applied.
-    ///
-    /// # Panics
-    /// Panics if any updates are redundant, as this indicates a bug that would
-    /// potentially cause infinite loops otherwise.
-    pub fn apply_updates(
-        &mut self,
-        updates: &EvalValue<&AlgebraicReference, T>,
-        source_name: impl Fn() -> String,
-    ) -> bool {
-        if updates.constraints.is_empty() {
-            return false;
-        }
-
-        log::trace!("    Updates from: {}", source_name());
-        for (poly, c) in &updates.constraints {
-            self.apply_update(poly, c)
-        }
-        true
-    }
-
     fn get_cell_mut<'b>(&'b mut self, poly: &AlgebraicReference) -> &'b mut Cell<'a, T> {
         match poly.next {
             false => &mut self.current[&poly.poly_id],

executor/src/witgen/vm_processor.rs

@@ -5,19 +5,24 @@ use itertools::Itertools;
 use number::{DegreeType, FieldElement};
 use parser_util::lines::indent;
 use std::cmp::max;
-use std::collections::HashSet;
+use std::collections::{BTreeMap, HashSet};
 use std::time::Instant;
 
 use crate::witgen::identity_processor::{self, IdentityProcessor};
 use crate::witgen::rows::RowUpdater;
+use crate::witgen::util::try_to_simple_poly;
+use crate::witgen::Constraint;
 use crate::witgen::IncompleteCause;
 
+use super::block_processor::OuterQuery;
 use super::data_structures::column_map::WitnessColumnMap;
 use super::data_structures::finalizable_data::FinalizableData;
 use super::query_processor::QueryProcessor;
 
-use super::rows::{Row, RowFactory, RowPair, UnknownStrategy};
-use super::{EvalError, EvalResult, EvalValue, FixedData, MutableState, QueryCallback};
+use super::rows::{CellValue, Row, RowFactory, RowPair, UnknownStrategy};
+use super::{
+    Constraints, EvalError, EvalResult, EvalValue, FixedData, MutableState, QueryCallback,
+};
 
 /// Maximal period checked during loop detection.
 const MAX_PERIOD: usize = 4;
@@ -60,7 +65,9 @@ pub struct VmProcessor<'a, T: FieldElement> {
     last_report: DegreeType,
     last_report_time: Instant,
     row_factory: RowFactory<'a, T>,
-    latch: Option<Expression<T>>,
+    outer_query: Option<OuterQuery<'a, T>>,
+    inputs: BTreeMap<PolyID, T>,
+    previously_set_inputs: BTreeMap<PolyID, DegreeType>,
 }
 
 impl<'a, T: FieldElement> VmProcessor<'a, T> {
@@ -71,7 +78,6 @@ impl<'a, T: FieldElement> VmProcessor<'a, T> {
         witnesses: HashSet<PolyID>,
         data: FinalizableData<'a, T>,
         row_factory: RowFactory<'a, T>,
-        latch: Option<Expression<T>>,
     ) -> Self {
         let (identities_with_next, identities_without_next): (Vec<_>, Vec<_>) = identities
             .iter()
@@ -95,7 +101,27 @@ impl<'a, T: FieldElement> VmProcessor<'a, T> {
             row_factory,
             last_report: 0,
             last_report_time: Instant::now(),
-            latch,
+            outer_query: None,
+            inputs: BTreeMap::new(),
+            previously_set_inputs: BTreeMap::new(),
+        }
+    }
+
+    pub fn with_outer_query(self, outer_query: OuterQuery<'a, T>) -> Self {
+        log::trace!("  Extracting inputs:");
+        let mut inputs = BTreeMap::new();
+        for (l, r) in outer_query.left.iter().zip(&outer_query.right.expressions) {
+            if let Some(right_poly) = try_to_simple_poly(r).map(|p| p.poly_id) {
+                if let Some(l) = l.constant_value() {
+                    log::trace!("    {} = {}", r, l);
+                    inputs.insert(right_poly, l);
+                }
+            }
+        }
+        Self {
+            outer_query: Some(outer_query),
+            inputs,
+            ..self
         }
     }
 
@@ -111,6 +137,8 @@ impl<'a, T: FieldElement> VmProcessor<'a, T> {
     ) -> EvalValue<&'a AlgebraicReference, T> {
         assert!(self.data.len() == 1);
 
+        let mut outer_assignments = vec![];
+
         // Are we in an infinite loop and can just re-use the old values?
         let mut looping_period = None;
         let mut loop_detection_log_level = log::Level::Info;
@@ -155,29 +183,26 @@ impl<'a, T: FieldElement> VmProcessor<'a, T> {
             // add and compute some values for the next row as well.
             if looping_period.is_none() && row_index != rows_left - 1 {
                 self.ensure_has_next_row(row_index);
-                self.compute_row(row_index, mutable_state);
+                outer_assignments.extend(self.compute_row(row_index, mutable_state).into_iter());
 
-                if let Some(latch) = self.latch.as_ref() {
-                    // Evaluate latch expression and return if it evaluates to 1.
-                    let row_pair = RowPair::from_single_row(
-                        self.row(row_index),
-                        row_index + self.row_offset,
-                        self.fixed_data,
-                        UnknownStrategy::Unknown,
-                    );
-                    let latch_value = row_pair
-                        .evaluate(latch)
-                        .ok()
-                        .and_then(|l| l.constant_value());
-
-                    if let Some(latch_value) = latch_value {
-                        if latch_value.is_one() {
-                            log::trace!("Machine returns!");
-                            return EvalValue::complete(vec![]);
+                // Evaluate latch expression and return if it evaluates to 1.
+                if let Some(latch) = self.latch_value(row_index) {
+                    if latch {
+                        log::trace!("Machine returns!");
+                        if self.outer_query.as_ref().unwrap().is_complete() {
+                            return EvalValue::complete(outer_assignments);
+                        } else {
+                            return EvalValue::incomplete_with_constraints(
+                                outer_assignments,
+                                IncompleteCause::BlockMachineLookupIncomplete,
+                            );
                         }
-                    } else {
-                        return EvalValue::incomplete(IncompleteCause::UnknownLatch);
                     }
+                } else if self.outer_query.is_some() {
+                    // If we have an outer query (and therefore a latch expression),
+                    // its value should be known at this point.
+                    // Probably, we don't have all the necessary inputs.
+                    return EvalValue::incomplete(IncompleteCause::UnknownLatch);
                 }
             };
         }
@@ -187,7 +212,7 @@ impl<'a, T: FieldElement> VmProcessor<'a, T> {
             self.fixed_data.degree + 1
         );
 
-        EvalValue::complete(vec![])
+        EvalValue::complete(outer_assignments)
     }
 
     /// Checks if the last rows are repeating and returns the period.
@@ -208,6 +233,21 @@ impl<'a, T: FieldElement> VmProcessor<'a, T> {
         })
     }
 
+    fn latch_value(&self, row_index: DegreeType) -> Option<bool> {
+        let row_pair = RowPair::from_single_row(
+            self.row(row_index),
+            row_index + self.row_offset,
+            self.fixed_data,
+            UnknownStrategy::Unknown,
+        );
+        self.outer_query
+            .as_ref()
+            .and_then(|outer_query| outer_query.right.selector.as_ref())
+            .and_then(|latch| row_pair.evaluate(latch).ok())
+            .and_then(|l| l.constant_value())
+            .map(|l| l.is_one())
+    }
+
     // Returns a reference to a given row
     fn row(&self, row_index: DegreeType) -> &Row<'a, T> {
         &self.data[row_index as usize]
@@ -224,7 +264,7 @@ impl<'a, T: FieldElement> VmProcessor<'a, T> {
         &mut self,
         row_index: DegreeType,
         mutable_state: &mut MutableState<'a, '_, T, Q>,
-    ) {
+    ) -> Constraints<&'a AlgebraicReference, T> {
         log::trace!("Row: {}", row_index + self.row_offset);
 
         log::trace!("  Going over all identities until no more progress is made");
@@ -234,52 +274,69 @@ impl<'a, T: FieldElement> VmProcessor<'a, T> {
             CompletableIdentities::new(self.identities_without_next_ref.iter().cloned());
         let mut identities_with_next_ref =
             CompletableIdentities::new(self.identities_with_next_ref.iter().cloned());
-        self.loop_until_no_progress(row_index, &mut identities_without_next_ref, mutable_state)
-            .and_then(|_| {
-                self.loop_until_no_progress(row_index, &mut identities_with_next_ref, mutable_state)
+        let outer_assignments = self
+            .loop_until_no_progress(row_index, &mut identities_without_next_ref, mutable_state)
+            .and_then(|outer_assignments| {
+                Ok(outer_assignments
+                    .into_iter()
+                    .chain(self.loop_until_no_progress(
+                        row_index,
+                        &mut identities_with_next_ref,
+                        mutable_state,
+                    )?)
+                    .collect::<Vec<_>>())
             })
             .map_err(|e| self.report_failure_and_panic_unsatisfiable(row_index, e))
             .unwrap();
 
         // Check that the computed row is "final" by asserting that all unknown values can
         // be set to 0.
-        log::trace!("  Checking that remaining identities hold when unknown values are set to 0");
-        let mut identity_processor = IdentityProcessor::new(self.fixed_data, mutable_state);
-        self.process_identities(
-            row_index,
-            &mut identities_without_next_ref,
-            UnknownStrategy::Zero,
-            &mut identity_processor,
-        )
-        .and_then(|_| {
+        // This check is only done for the primary machine, as secondary machines might simply
+        // not have all the inputs yet and therefore be underconstrained.
+        if self.outer_query.is_none() {
+            log::trace!(
+                "  Checking that remaining identities hold when unknown values are set to 0"
+            );
+            let mut identity_processor = IdentityProcessor::new(self.fixed_data, mutable_state);
             self.process_identities(
                 row_index,
-                &mut identities_with_next_ref,
+                &mut identities_without_next_ref,
                 UnknownStrategy::Zero,
                 &mut identity_processor,
             )
-        })
-        .map_err(|e| self.report_failure_and_panic_underconstrained(row_index, e))
-        .unwrap();
+            .and_then(|_| {
+                self.process_identities(
+                    row_index,
+                    &mut identities_with_next_ref,
+                    UnknownStrategy::Zero,
+                    &mut identity_processor,
+                )
+            })
+            .map_err(|e| self.report_failure_and_panic_underconstrained(row_index, e))
+            .unwrap();
+        }
 
         log::trace!(
             "{}",
-            self.row(row_index)
-                .render(&format!("===== Row {}", row_index), true, &self.witnesses)
+            self.row(row_index).render(
+                &format!("===== Row {}", row_index as DegreeType + self.row_offset),
+                true,
+                &self.witnesses
+            )
         );
+
+        outer_assignments
     }
 
     /// Loops over all identities and queries, until no further progress is made.
     /// @returns the "incomplete" identities, i.e. identities that contain unknown values.
-    fn loop_until_no_progress<Q>(
+    fn loop_until_no_progress<Q: QueryCallback<T>>(
         &mut self,
         row_index: DegreeType,
         identities: &mut CompletableIdentities<'a, T>,
         mutable_state: &mut MutableState<'a, '_, T, Q>,
-    ) -> Result<(), Vec<EvalError<T>>>
-    where
-        Q: FnMut(&str) -> Option<T> + Send + Sync,
-    {
+    ) -> Result<Constraints<&'a AlgebraicReference, T>, Vec<EvalError<T>>> {
+        let mut outer_assignments = vec![];
         loop {
             let mut identity_processor = IdentityProcessor::new(self.fixed_data, mutable_state);
             let mut progress = self.process_identities(
@@ -288,6 +345,16 @@ impl<'a, T: FieldElement> VmProcessor<'a, T> {
                 UnknownStrategy::Unknown,
                 &mut identity_processor,
             )?;
+            if let Some(true) = self.latch_value(row_index) {
+                let (outer_query_progress, new_outer_assignments) = self
+                    .process_outer_query(row_index, mutable_state)
+                    .map_err(|e| vec![e])?;
+                progress |= outer_query_progress;
+                outer_assignments.extend(new_outer_assignments);
+            }
+
+            progress |= self.set_inputs_if_unset(row_index);
+
             let mut updates = EvalValue::complete(vec![]);
             let mut query_processor =
                 QueryProcessor::new(self.fixed_data, mutable_state.query_callback);
@@ -309,7 +376,7 @@ impl<'a, T: FieldElement> VmProcessor<'a, T> {
                 break;
             }
         }
-        Ok(())
+        Ok(outer_assignments)
     }
 
     /// Loops over all identities once and updates the current row and next row.
@@ -383,18 +450,117 @@ impl<'a, T: FieldElement> VmProcessor<'a, T> {
         }
     }
 
+    fn process_outer_query<Q: QueryCallback<T>>(
+        &mut self,
+        row_index: DegreeType,
+        mutable_state: &mut MutableState<'a, '_, T, Q>,
+    ) -> Result<(bool, Constraints<&'a AlgebraicReference, T>), EvalError<T>> {
+        let OuterQuery { left, right } = self
+            .outer_query
+            .as_ref()
+            .expect("Asked to process outer query, but it was not set!");
+
+        let row_pair = RowPair::new(
+            &self.data[row_index as usize],
+            &self.data[row_index as usize + 1],
+            (row_index + self.row_offset) % self.fixed_data.degree,
+            self.fixed_data,
+            UnknownStrategy::Unknown,
+        );
+
+        let mut identity_processor = IdentityProcessor::new(self.fixed_data, mutable_state);
+        let updates = identity_processor
+            .process_link(left, right, &row_pair)
+            .map_err(|e| {
+                log::warn!("Error in outer query: {e}");
+                log::warn!("Some of the following entries could not be matched:");
+                for (l, r) in left.iter().zip(right.expressions.iter()) {
+                    if let Ok(r) = row_pair.evaluate(r) {
+                        log::warn!("  => {} = {}", l, r);
+                    }
+                }
+                e
+            })?;
+
+        let progress = self.apply_updates(row_index, &updates, || "outer query".to_string());
+
+        let outer_assignments = updates
+            .constraints
+            .into_iter()
+            .filter(|(poly, _)| !self.witnesses.contains(&poly.poly_id))
+            .collect::<Vec<_>>();
+
+        Ok((progress, outer_assignments))
+    }
+
+    /// Sets the inputs to the values given in [VmProcessor::inputs] if they are not already set.
+    /// Typically, inputs will have a constraint of the form: `((1 - instr__reset) * (_input' - _input)) = 0;`
+    /// So, once the value of `_input` is set, this function will do nothing until the next reset instruction.
+    /// However, if `_input` does become unconstrained, we need to undo all changes we've done so far.
+    /// For this reason, we keep track of all changes we've done to inputs in [VmProcessor::previously_set_inputs].
+    fn set_inputs_if_unset(&mut self, row_index: DegreeType) -> bool {
+        let mut input_updates = EvalValue::complete(vec![]);
+        for (poly_id, value) in self.inputs.iter() {
+            match &self.data[row_index as usize][poly_id].value {
+                CellValue::Known(_) => {}
+                CellValue::RangeConstraint(_) | CellValue::Unknown => {
+                    input_updates.combine(EvalValue::complete([(
+                        &self.fixed_data.witness_cols[poly_id].poly,
+                        Constraint::Assignment(*value),
+                    )]));
+                }
+            };
+        }
+
+        for (poly, _) in &input_updates.constraints {
+            let poly_id = poly.poly_id;
+            if let Some(start_row) = self.previously_set_inputs.remove(&poly_id) {
+                log::trace!(
+                    "    Resetting previously set inputs for column: {}",
+                    self.fixed_data.column_name(&poly_id)
+                );
+                for row_index in start_row..row_index {
+                    self.data[row_index as usize][&poly_id].value = CellValue::Unknown;
+                }
+            }
+        }
+        for (poly, _) in &input_updates.constraints {
+            self.previously_set_inputs.insert(poly.poly_id, row_index);
+        }
+        self.apply_updates(row_index, &input_updates, || "inputs".to_string())
+    }
+
     fn apply_updates(
         &mut self,
         row_index: DegreeType,
-        updates: &EvalValue<&AlgebraicReference, T>,
+        updates: &EvalValue<&'a AlgebraicReference, T>,
         source_name: impl Fn() -> String,
     ) -> bool {
         if updates.constraints.is_empty() {
             return false;
         }
+
+        log::trace!("    Updates from: {}", source_name());
+
+        // Build RowUpdater
+        // (a bit complicated, because we need two mutable
+        // references to elements of the same vector)
         let (current, next) = self.data.mutable_row_pair(row_index as usize);
         let mut row_updater = RowUpdater::new(current, next, row_index + self.row_offset);
-        row_updater.apply_updates(updates, source_name)
+
+        for (poly, c) in &updates.constraints {
+            if self.witnesses.contains(&poly.poly_id) {
+                row_updater.apply_update(poly, c);
+            } else if let Constraint::Assignment(v) = c {
+                let left = &mut self.outer_query.as_mut().unwrap().left;
+                log::trace!("      => {} (outer) = {}", poly, v);
+                for l in left.iter_mut() {
+                    l.assign(poly, *v);
+                }
+            };
+        }
+
+        true
     }
 
     fn report_failure_and_panic_unsatisfiable(

test_data/asm/vm_to_vm.asm

@@ -1,5 +1,7 @@
 machine Main {
 
+    degree 256;
+
     VM vm;
 
     reg pc[@pc];
@@ -14,9 +16,8 @@ machine Main {
 
     function main {
         A <== add(1, 1);
-        // TODO: uncomment the following two lines once we support having many calls to a machine
-        // A <== add(A, 1);
-        // A <== sub(A, 1);
+        A <== add(A, 1);
+        A <== sub(A, 1);
         assert_eq A, 2;
         return;
     }

test_data/asm/vm_to_vm_dynamic_trace_length.asm

@@ -0,0 +1,70 @@
+machine Main {
+
+    degree 256;
+
+    Pow pow;
+
+    reg pc[@pc];
+    reg X[<=];
+    reg Y[<=];
+    reg Z[<=];
+    reg A;
+
+    instr pow X, Y -> Z = pow.pow
+    instr assert_eq X, Y { X = Y }
+
+    function main {
+        A <== pow(7, 0);
+        assert_eq A, 1;
+
+        A <== pow(7, 1);
+        assert_eq A, 7;
+
+        A <== pow(7, 2);
+        assert_eq A, 49;
+
+        A <== pow(2, 2);
+        assert_eq A, 4;
+
+        A <== pow(2, 10);
+        assert_eq A, 1024;
+
+        return;
+    }
+}
+
+// Computes X^Y by multiplying X by itself Y times
+machine Pow {
+
+    reg pc[@pc];
+    reg X[<=];
+    reg Y[<=];
+    reg Z[<=];
+    reg A;
+    reg CNT;
+
+    col witness XInv;
+    col witness XIsZero;
+    XIsZero  = 1 - X * XInv;
+    XIsZero * X = 0;
+    XIsZero * (1 - XIsZero) = 0;
+
+    instr mul X, Y -> Z { X * Y = Z }
+    instr jmpz X, l: label { pc' = XIsZero * l + (1 - XIsZero) * (pc + 1) }
+    instr jmp l: label { pc' = l }
+
+    function pow x: field, y: field -> field {
+        A <=X= 1;
+        CNT <=X= y;
+
+        start::
+        jmpz CNT, done;
+        A <== mul(A, x);
+        CNT <=X= CNT - 1;
+        jmp start;
+
+        done::
+        return A;
+    }
+
+}

test_data/asm/vm_to_vm_to_block.asm

@@ -0,0 +1,66 @@
+machine Main {
+
+    degree 256;
+
+    Pythagoras pythagoras;
+
+    reg pc[@pc];
+    reg X[<=];
+    reg Y[<=];
+    reg Z[<=];
+    reg A;
+
+    instr pythagoras X, Y -> Z = pythagoras.pythagoras
+    instr assert_eq X, Y { X = Y }
+
+    function main {
+        A <== pythagoras(3, 4);
+        assert_eq A, 25;
+
+        A <== pythagoras(4, 3);
+        assert_eq A, 25;
+
+        A <== pythagoras(1, 2);
+        assert_eq A, 5;
+
+        return;
+    }
+}
+
+
+machine Pythagoras {
+
+    Arith arith;
+
+    reg pc[@pc];
+    reg X[<=];
+    reg Y[<=];
+    reg Z[<=];
+    reg A;
+    reg B;
+
+
+    instr add X, Y -> Z = arith.add
+    instr mul X, Y -> Z = arith.mul
+
+    function pythagoras a: field, b: field -> field {
+        A <== mul(a, a);
+        B <== mul(b, b);
+        A <== add(A, B);
+        return A;
+    }
+}
+
+machine Arith(latch, operation_id) {
+
+    operation add<0> x1, x2 -> y;
+    operation mul<1> x1, x2 -> y;
+
+    col fixed latch = [1]*;
+    col witness operation_id;
+    col witness x1;
+    col witness x2;
+    col witness y;
+
+    y = operation_id * (x1 * x2) + (1 - operation_id) * (x1 + x2);
+}

test_data/asm/vm_to_vm_to_vm.asm

@@ -0,0 +1,73 @@
+machine Main {
+
+    degree 256;
+
+    Pythagoras pythagoras;
+
+    reg pc[@pc];
+    reg X[<=];
+    reg Y[<=];
+    reg Z[<=];
+    reg A;
+
+    instr pythagoras X, Y -> Z = pythagoras.pythagoras
+    instr assert_eq X, Y { X = Y }
+
+    function main {
+        A <== pythagoras(3, 4);
+        assert_eq A, 25;
+
+        A <== pythagoras(4, 3);
+        assert_eq A, 25;
+
+        A <== pythagoras(1, 2);
+        assert_eq A, 5;
+
+        return;
+    }
+}
+
+machine Pythagoras {
+
+    Arith arith;
+
+    reg pc[@pc];
+    reg X[<=];
+    reg Y[<=];
+    reg Z[<=];
+    reg A;
+    reg B;
+
+
+    instr add X, Y -> Z = arith.add
+    instr mul X, Y -> Z = arith.mul
+
+    function pythagoras a: field, b: field -> field {
+        A <== mul(a, a);
+        B <== mul(b, b);
+        A <== add(A, B);
+        return A;
+    }
+}
+
+machine Arith {
+
+    reg pc[@pc];
+    reg X[<=];
+    reg Y[<=];
+    reg Z[<=];
+    reg A;
+
+    instr add X, Y -> Z { X + Y = Z }
+    instr mul X, Y -> Z { X * Y = Z }
+
+    function add x: field, y: field -> field {
+        A <== add(x, y);
+        return A;
+    }
+
+    function mul x: field, y: field -> field {
+        A <== mul(x, y);
+        return A;
+    }
+}