masm: sorting boundary constraints

air-script/tests/aux_trace/aux_trace.masm

@@ -30,19 +30,19 @@ proc.compute_evaluate_boundary_constraints
     padw mem_loadw.4294900001 movdn.3 movdn.3 drop drop push.1 push.0 ext2sub
     # Multiply by the composition coefficient
     padw mem_loadw.4294900203 movdn.3 movdn.3 drop drop ext2mul
-    # boundary constraint 0 for aux
+    # boundary constraint 2 for aux
     padw mem_loadw.4294900073 movdn.3 movdn.3 drop drop push.1 push.0 ext2sub
     # Multiply by the composition coefficient
     padw mem_loadw.4294900203 drop drop ext2mul
-    # boundary constraint 1 for aux
-    padw mem_loadw.4294900073 movdn.3 movdn.3 drop drop push.1 push.0 ext2sub
+    # boundary constraint 3 for aux
+    padw mem_loadw.4294900074 movdn.3 movdn.3 drop drop padw mem_loadw.4294900150 movdn.3 movdn.3 drop drop ext2sub
     # Multiply by the composition coefficient
     padw mem_loadw.4294900204 movdn.3 movdn.3 drop drop ext2mul
-    # boundary constraint 2 for aux
-    padw mem_loadw.4294900074 movdn.3 movdn.3 drop drop padw mem_loadw.4294900150 movdn.3 movdn.3 drop drop ext2sub
+    # boundary constraint 4 for aux
+    padw mem_loadw.4294900073 movdn.3 movdn.3 drop drop push.1 push.0 ext2sub
     # Multiply by the composition coefficient
     padw mem_loadw.4294900204 drop drop ext2mul
-    # boundary constraint 3 for aux
+    # boundary constraint 5 for aux
     padw mem_loadw.4294900074 movdn.3 movdn.3 drop drop push.1 push.0 ext2sub
     # Multiply by the composition coefficient
     padw mem_loadw.4294900205 movdn.3 movdn.3 drop drop ext2mul

air-script/tests/constants/constants.masm

@@ -38,11 +38,11 @@ proc.compute_evaluate_boundary_constraints
     padw mem_loadw.4294900003 movdn.3 movdn.3 drop drop push.1 push.0 push.0 push.0 ext2add push.1 push.0 ext2sub push.1 push.0 ext2add push.2 push.0 ext2sub push.2 push.0 ext2add push.0 push.0 ext2sub ext2sub
     # Multiply by the composition coefficient
     padw mem_loadw.4294900204 movdn.3 movdn.3 drop drop ext2mul
-    # boundary constraint 0 for aux
+    # boundary constraint 4 for aux
     padw mem_loadw.4294900073 movdn.3 movdn.3 drop drop push.1 push.0 push.0 push.0 push.2 push.0 ext2mul ext2add ext2sub
     # Multiply by the composition coefficient
     padw mem_loadw.4294900204 drop drop ext2mul
-    # boundary constraint 1 for aux
+    # boundary constraint 5 for aux
     padw mem_loadw.4294900073 movdn.3 movdn.3 drop drop push.1 push.0 push.1 push.0 push.1 push.0 ext2mul ext2sub ext2sub
     # Multiply by the composition coefficient
     padw mem_loadw.4294900205 movdn.3 movdn.3 drop drop ext2mul

codegen/masm/src/codegen.rs

@@ -1,7 +1,8 @@
 use crate::constants::{AUX_TRACE, MAIN_TRACE};
 use crate::visitor::{
-    walk_boundary_constraints, walk_constant_bindings, walk_integrity_constraint_degrees,
-    walk_integrity_constraints, walk_periodic_columns, walk_public_inputs, AirVisitor,
+    walk_boundary_constraints_in_natural_order, walk_constant_bindings,
+    walk_integrity_constraint_degrees, walk_integrity_constraints, walk_periodic_columns,
+    walk_public_inputs, AirVisitor,
 };
 use ir::{
     constraints::{ConstraintRoot, Operation},
@@ -308,9 +309,7 @@ impl<'ast> CodeGenerator<'ast> {
     /// Emits code for the procedure `compute_evaluate_boundary_constraints`.
     fn gen_compute_evaluate_boundary_constraints(&mut self) -> Result<(), CodegenError> {
         self.writer.proc("compute_evaluate_boundary_constraints");
-        walk_boundary_constraints(self, self.ir, MAIN_TRACE)?;
-        self.boundary_contraints = 0; // reset counter for the aux trace
-        walk_boundary_constraints(self, self.ir, AUX_TRACE)?;
+        walk_boundary_constraints_in_natural_order(self, self.ir)?;
         self.writer.end();
 
         Ok(())
@@ -415,6 +414,8 @@ impl<'ast> AirVisitor<'ast> for CodeGenerator<'ast> {
         self.writer
             .header("Multiply by the composition coefficient");
 
+        // Note: The correctness of the load below relies on the integrity constraint being
+        // iterated first _and_ the boundary constraints being iterated in natural order.
         load_quadratic_element(
             &mut self.writer,
             self.config.composition_coef_address,

codegen/masm/src/visitor.rs

@@ -1,5 +1,6 @@
+use crate::constants::{AUX_TRACE, MAIN_TRACE};
 use ir::{
-    constraints::{ConstraintRoot, Operation},
+    constraints::{ConstraintDomain, ConstraintRoot, Operation},
     AccessType, AirIR, ConstantBinding, IntegrityConstraintDegree, NodeIndex, PeriodicColumn,
     PublicInput, TraceAccess, Value,
 };
@@ -103,13 +104,41 @@ pub fn walk_periodic_columns<'ast, V: AirVisitor<'ast>>(
     Ok(())
 }
 
-pub fn walk_boundary_constraints<'ast, V: AirVisitor<'ast>>(
+/// Walks the IR's boundary constraints.
+///
+/// The boundary constraints have an implicit natural order. Defined by:
+///
+/// - Trace segment: The main trace is sorted before the auxiliary trace.
+/// - Row: Constraints on the first row are sorted before the last row.
+/// - Column: Constraints on lower columns are sorted before higher columns.
+///
+/// The order above has two functions:
+///
+/// - It sorts the constraints so that the order of iteration matches the order in which the
+/// composition coefficients are defined.
+/// - It sorts the constraints so groups with the same divisor are iterated together.
+pub fn walk_boundary_constraints_in_natural_order<'ast, V: AirVisitor<'ast>>(
     visitor: &mut V,
     ir: &'ast AirIR,
-    trace_segment: u8,
 ) -> Result<(), V::Error> {
-    for boundary in ir.boundary_constraints(trace_segment) {
-        visitor.visit_boundary_constraint(boundary, trace_segment)?;
+    fn domain(boundary: &&ConstraintRoot) -> u8 {
+        match boundary.domain() {
+            ConstraintDomain::FirstRow => 0,
+            ConstraintDomain::LastRow => 1,
+            ConstraintDomain::EveryRow => panic!("EveryRow is not supported"),
+            ConstraintDomain::EveryFrame(_) => panic!("EveryFrame is not supported"),
+        }
+    }
+    for segment in [MAIN_TRACE, AUX_TRACE] {
+        let mut constraints: Vec<&'ast ConstraintRoot> =
+            ir.boundary_constraints(segment).iter().collect();
+
+        // TODO: Sort by the column index. Issue #315
+        constraints.sort_by_key(domain);
+
+        for boundary in constraints {
+            visitor.visit_boundary_constraint(boundary, segment)?;
+        }
     }
 
     Ok(())

codegen/masm/tests/test_boundary.rs

@@ -222,20 +222,34 @@ fn test_complex_boundary() {
     let program_outputs = process.execute(&program).expect("execution failed");
     let result_stack = program_outputs.stack();
 
-    // results are in stack-order
+    // Note: The order of the results is _not_ the same as the definition order in the AirScript.
+    // The order below is:
+    //
+    // 1. First row boundary constraints for the MAIN trace
+    // 2. Last row boundary constraints for the MAIN trace
+    // 3. First row boundary constraints for the AUX trace
+    // 4. Last row boundary constraints for the AUX trace
+    //
+    // Results are in stack-order.
     #[rustfmt::skip]
     let expected = to_stack_order(&[
+        // last row aux trace
         f - one,              // enf f.last = 1
+
+        // first row aux trace
         f - rand[0],          // enf f.first = $rand[0]
 
+        // last row main trace
+        b - public_inputs[3], // enf b.last = stack_outputs[1]
+        a - public_inputs[2], // enf a.last = stack_outputs[0]
+
+        // first row main trace
         e[1] - one,           // enf e[1].first = 1
         e[0],                 // enf e[0].first = 0
 
         d - one,              // enf d.first = 1
         c,                    // enf c.first = (B[0] - C[1][1]) * A
 
-        b - public_inputs[3], // enf b.last = stack_outputs[1]
-        a - public_inputs[2], // enf a.last = stack_outputs[0]
         b - public_inputs[1], // enf b.first = stack_inputs[1]
         a - public_inputs[0], // enf a.first = stack_inputs[0]
     ]);