fix zkDSL compiler crash (edge case when the end value of the non-unrolled loop index is defined by a const malloc)

Author: TomWambsgans

crates/lean_compiler/src/a_simplify_lang.rs

@@ -814,7 +814,20 @@ fn simplify_lines(
 
                 let start_simplified =
                     simplify_expr(start, &mut res, counters, array_manager, const_malloc, const_arrays);
-                let end_simplified = simplify_expr(end, &mut res, counters, array_manager, const_malloc, const_arrays);
+                let mut end_simplified =
+                    simplify_expr(end, &mut res, counters, array_manager, const_malloc, const_arrays);
+                if let SimpleExpr::ConstMallocAccess { malloc_label, offset } = end_simplified.clone() {
+                    // we use an auxilary variable to store the end value (const malloc inside non-unrolled loops does not work)
+                    let aux_end_var = format!("@aux_end_{}", counters.aux_vars);
+                    counters.aux_vars += 1;
+                    res.push(SimpleLine::Assignment {
+                        var: aux_end_var.clone().into(),
+                        operation: HighLevelOperation::Add,
+                        arg0: SimpleExpr::ConstMallocAccess { malloc_label, offset },
+                        arg1: SimpleExpr::zero(),
+                    });
+                    end_simplified = SimpleExpr::Var(aux_end_var);
+                }
 
                 for (simplified, original) in [
                     (start_simplified.clone(), start.clone()),

crates/lean_compiler/tests/test_compiler.rs

@@ -766,3 +766,21 @@ fn test_const_array() {
     "#;
     compile_and_run(program.to_string(), (&[], &[]), DEFAULT_NO_VEC_RUNTIME_MEMORY, false);
 }
+
+#[test]
+fn test_const_malloc_end_iterator_loop() {
+    let program = r#"
+    fn main() {
+        x = malloc(2);
+        x[0] = 3;
+        x[1] = 5;
+        for i in 0..2 unroll {
+            for j in 0..x[i] {
+                print(i, j);
+            }
+        }
+        return;
+    }
+    "#;
+    compile_and_run(program.to_string(), (&[], &[]), DEFAULT_NO_VEC_RUNTIME_MEMORY, false);
+}