From 49f2fd82803a7cbe2e815c8224135c456889cef6 Mon Sep 17 00:00:00 2001
From: grjte <grjte@protonmail.com>
Date: Thu, 26 Jan 2023 19:12:55 +0000
Subject: [PATCH] fixup codegen

---
 .../src/air/boundary_constraints.rs           | 54 +++++++++----------
 1 file changed, 24 insertions(+), 30 deletions(-)

diff --git a/codegen/winterfell/src/air/boundary_constraints.rs b/codegen/winterfell/src/air/boundary_constraints.rs
index eea684ce..011df8e9 100644
--- a/codegen/winterfell/src/air/boundary_constraints.rs
+++ b/codegen/winterfell/src/air/boundary_constraints.rs
@@ -13,27 +13,8 @@ pub(super) fn add_fn_get_assertions(impl_ref: &mut Impl, ir: &AirIR) {
         .arg_ref_self()
         .ret("Vec<Assertion<Felt>>");
 
-    // declare the result vector to be returned.
-    let last_row = "last_step";
-    get_assertions.line("let mut result = Vec::new();");
-    get_assertions.line(format!("let {last_row} = self.last_step();"));
-
     // add the boundary constraints
-    for constraint in ir.boundary_constraints(0) {
-        let (trace_access, expr_root) = ir
-            .constraint_graph()
-            .split_boundary_constraint(constraint.node_index())
-            .expect("failed to split boundary constraint");
-        debug_assert!(trace_access.trace_segment() == 0);
-
-        let assertion = format!(
-            "result.push(Assertion::single({}, {}, {}));",
-            trace_access.col_idx(),
-            domain_to_str(constraint.domain(), last_row),
-            expr_root.to_string(ir, ElemType::Felt)
-        );
-        get_assertions.line(assertion);
-    }
+    add_assertions(get_assertions, ir, 0);
 
     // return the result
     get_assertions.line("result");
@@ -50,30 +31,43 @@ pub(super) fn add_fn_get_aux_assertions(impl_ref: &mut Impl, ir: &AirIR) {
         .arg("aux_rand_elements", "&AuxTraceRandElements<E>")
         .ret("Vec<Assertion<E>>");
 
+    // add the boundary constraints
+    add_assertions(get_aux_assertions, ir, 1);
+
+    // return the result
+    get_aux_assertions.line("result");
+}
+
+/// Declares a result vector and adds assertions for boundary constraints to it for the specified
+/// trace segment
+fn add_assertions(func_body: &mut codegen::Function, ir: &AirIR, trace_segment: u8) {
+    let elem_type = if trace_segment == 0 {
+        ElemType::Felt
+    } else {
+        ElemType::E
+    };
+
     // declare the result vector to be returned.
     let last_row = "last_step";
-    get_aux_assertions.line("let mut result = Vec::new();");
-    get_aux_assertions.line(format!("let {last_row} = self.last_step();"));
+    func_body.line("let mut result = Vec::new();");
+    func_body.line(format!("let {last_row} = self.last_step();"));
 
-    // add the auxiliary boundary constraints
-    for constraint in ir.boundary_constraints(1) {
+    // add the boundary constraints
+    for constraint in ir.boundary_constraints(trace_segment) {
         let (trace_access, expr_root) = ir
             .constraint_graph()
             .split_boundary_constraint(constraint.node_index())
             .expect("failed to split boundary constraint");
-        debug_assert!(trace_access.trace_segment() == 1);
+        debug_assert!(trace_access.trace_segment() == trace_segment);
 
         let assertion = format!(
             "result.push(Assertion::single({}, {}, {}));",
             trace_access.col_idx(),
             domain_to_str(constraint.domain(), last_row),
-            expr_root.to_string(ir, ElemType::E)
+            expr_root.to_string(ir, elem_type)
         );
-        get_aux_assertions.line(assertion);
+        func_body.line(assertion);
     }
-
-    // return the result
-    get_aux_assertions.line("result");
 }
 
 /// Returns a string slice representing the provided constraint domain.