zkDSL: fix bug

Author: TomWambsgans

crates/lean_compiler/src/a_simplify_lang.rs

@@ -651,12 +651,28 @@ fn simplify_lines(
                 Expression::Binary { left, operation, right } => {
                     let left = simplify_expr(left, &mut res, counters, array_manager, const_malloc, const_arrays);
                     let right = simplify_expr(right, &mut res, counters, array_manager, const_malloc, const_arrays);
-                    res.push(SimpleLine::Assignment {
-                        var: var.clone().into(),
-                        operation: *operation,
-                        arg0: left,
-                        arg1: right,
-                    });
+                    // If both operands are constants, evaluate at compile time and assign the result
+                    if let (SimpleExpr::Constant(left_cst), SimpleExpr::Constant(right_cst)) = (&left, &right) {
+                        let result = ConstExpression::Binary {
+                            left: Box::new(left_cst.clone()),
+                            operation: *operation,
+                            right: Box::new(right_cst.clone()),
+                        }
+                        .try_naive_simplification();
+                        res.push(SimpleLine::Assignment {
+                            var: var.clone().into(),
+                            operation: HighLevelOperation::Add,
+                            arg0: SimpleExpr::Constant(result),
+                            arg1: SimpleExpr::zero(),
+                        });
+                    } else {
+                        res.push(SimpleLine::Assignment {
+                            var: var.clone().into(),
+                            operation: *operation,
+                            arg0: left,
+                            arg1: right,
+                        });
+                    }
                 }
                 Expression::MathExpr(_, _) => unreachable!(),
             },

crates/lean_compiler/tests/test_compiler.rs

@@ -1305,15 +1305,20 @@ fn test_2d_nested_array_with_expressions() {
     );
 }
 
-// #[test]
-// fn bug() {
-//     let program = r#"
-//     const ARR = [1];
-//     fn main() {
-//         x = ARR[0]**2;
-//         print(x);
-//         return;
-//     }
-//     "#;
-//     compile_and_run(&ProgramSource::Raw(program.to_string()), (&[], &[]), false);
-// }
+#[test]
+fn test_const_array_element_exponentiation() {
+    let program = r#"
+    const ARR = [[5]];
+    fn main() {
+        x = ARR[0][0]**2;
+        assert x == 25;
+        return;
+    }
+    "#;
+    compile_and_run(
+        &ProgramSource::Raw(program.to_string()),
+        (&[], &[]),
+        DEFAULT_NO_VEC_RUNTIME_MEMORY,
+        false,
+    );
+}