Merge pull request #29019 from ProvableHQ/mohammadfawaz/28961

A few fixes to options in type checking and option lowering

Author: d0cd

compiler/passes/src/destructuring/ast.rs

@@ -24,6 +24,64 @@ impl AstReconstructor for DestructuringVisitor<'_> {
     type AdditionalInput = ();
     type AdditionalOutput = Vec<Statement>;
 
+    /// Reconstructs a binary expression, expanding equality and inequality over
+    /// tuples into elementwise comparisons. When both sides are tuples and the
+    /// operator is `==` or `!=`, it generates per-element comparisons and folds
+    /// them with AND/OR; otherwise the expression is rebuilt normally.
+    ///
+    /// Example: `(a, b) == (c, d)` → `(a == c) && (b == d)`
+    /// Example: `(a, b, c) != (x, y, z)` → `(a != x) || (b != y) || (c != z)`
+    fn reconstruct_binary(
+        &mut self,
+        input: BinaryExpression,
+        _additional: &Self::AdditionalInput,
+    ) -> (Expression, Self::AdditionalOutput) {
+        let (left, mut statements) = self.reconstruct_expression_tuple(input.left);
+        let (right, statements2) = self.reconstruct_expression_tuple(input.right);
+        statements.extend(statements2);
+
+        use BinaryOperation::*;
+
+        // Tuple equality / inequality expansion
+        if let (Expression::Tuple(tuple_left), Expression::Tuple(tuple_right)) = (&left, &right)
+            && matches!(input.op, Eq | Neq)
+        {
+            assert_eq!(tuple_left.elements.len(), tuple_right.elements.len());
+
+            // Directly build elementwise (l OP r)
+            let pieces: Vec<Expression> = tuple_left
+                .elements
+                .iter()
+                .zip(&tuple_right.elements)
+                .map(|(l, r)| {
+                    let expr: Expression = BinaryExpression {
+                        op: input.op,
+                        left: l.clone(),
+                        right: r.clone(),
+                        span: Default::default(),
+                        id: self.state.node_builder.next_id(),
+                    }
+                    .into();
+
+                    self.state.type_table.insert(expr.id(), Type::Boolean);
+                    expr
+                })
+                .collect();
+
+            // Fold appropriately
+            let op = match input.op {
+                Eq => BinaryOperation::And,
+                Neq => BinaryOperation::Or,
+                _ => unreachable!(),
+            };
+
+            return (self.fold_with_op(op, pieces.into_iter()), statements);
+        }
+
+        // Fallback
+        (BinaryExpression { op: input.op, left, right, ..input }.into(), Default::default())
+    }
+
     /// Replaces a tuple access expression with the appropriate expression.
     fn reconstruct_tuple_access(
         &mut self,
@@ -66,9 +124,12 @@ impl AstReconstructor for DestructuringVisitor<'_> {
         mut input: TernaryExpression,
         _additional: &(),
     ) -> (Expression, Self::AdditionalOutput) {
-        let (if_true, mut statements) = self.reconstruct_expression_tuple(std::mem::take(&mut input.if_true));
-        let (if_false, statements2) = self.reconstruct_expression_tuple(std::mem::take(&mut input.if_false));
+        let (condition, mut statements) =
+            self.reconstruct_expression(std::mem::take(&mut input.condition), &Default::default());
+        let (if_true, statements2) = self.reconstruct_expression_tuple(std::mem::take(&mut input.if_true));
         statements.extend(statements2);
+        let (if_false, statements3) = self.reconstruct_expression_tuple(std::mem::take(&mut input.if_false));
+        statements.extend(statements3);
 
         match (if_true, if_false) {
             (Expression::Tuple(tuple_true), Expression::Tuple(tuple_false)) => {
@@ -77,20 +138,17 @@ impl AstReconstructor for DestructuringVisitor<'_> {
                     panic!("Should have tuple type");
                 };
 
-                // We'll be reusing `input.condition`, so assign it to a variable.
-                let cond = if let Expression::Path(..) = input.condition {
-                    input.condition
+                // We'll be reusing `condition`, so assign it to a variable.
+                let cond = if let Expression::Path(..) = condition {
+                    condition
                 } else {
                     let place = Identifier::new(
                         self.state.assigner.unique_symbol("cond", "$$"),
                         self.state.node_builder.next_id(),
                     );
 
-                    let definition = self.state.assigner.simple_definition(
-                        place,
-                        input.condition,
-                        self.state.node_builder.next_id(),
-                    );
+                    let definition =
+                        self.state.assigner.simple_definition(place, condition, self.state.node_builder.next_id());
 
                     statements.push(definition);
 
@@ -144,12 +202,59 @@ impl AstReconstructor for DestructuringVisitor<'_> {
             }
             (if_true, if_false) => {
                 // This isn't a tuple. Just rebuild it and otherwise leave it alone.
-                (TernaryExpression { if_true, if_false, ..input }.into(), statements)
+                (TernaryExpression { condition, if_true, if_false, ..input }.into(), statements)
             }
         }
     }
 
     /* Statements */
+    /// `assert_eq` and `assert_neq` comparing tuples should be expanded to as many asserts as
+    /// the length of each tuple.
+    fn reconstruct_assert(&mut self, input: AssertStatement) -> (Statement, Self::AdditionalOutput) {
+        match input.variant {
+            AssertVariant::Assert(expr) => {
+                // Simple assert, just reconstruct the expression.
+                let (expr, _) = self.reconstruct_expression(expr, &Default::default());
+                (AssertStatement { variant: AssertVariant::Assert(expr), ..input }.into(), Default::default())
+            }
+            AssertVariant::AssertEq(ref left, ref right) | AssertVariant::AssertNeq(ref left, ref right) => {
+                let (left, mut statements) = self.reconstruct_expression_tuple(left.clone());
+                let (right, statements2) = self.reconstruct_expression_tuple(right.clone());
+                statements.extend(statements2);
+
+                match (&left, &right) {
+                    (Expression::Tuple(tuple_left), Expression::Tuple(tuple_right)) => {
+                        // Ensure the tuple lengths match
+                        assert_eq!(tuple_left.elements.len(), tuple_right.elements.len());
+
+                        for (l, r) in tuple_left.elements.iter().zip(&tuple_right.elements) {
+                            let assert_variant = match input.variant {
+                                AssertVariant::AssertEq(_, _) => AssertVariant::AssertEq(l.clone(), r.clone()),
+                                AssertVariant::AssertNeq(_, _) => AssertVariant::AssertNeq(l.clone(), r.clone()),
+                                _ => unreachable!(),
+                            };
+
+                            let stmt = AssertStatement { variant: assert_variant, ..input.clone() }.into();
+                            statements.push(stmt);
+                        }
+
+                        // We don't need the original statement, just the ones we've created.
+                        (Statement::dummy(), statements)
+                    }
+                    _ => {
+                        // Not tuples, just keep the original assert
+                        let variant = match input.variant {
+                            AssertVariant::AssertEq(_, _) => AssertVariant::AssertEq(left, right),
+                            AssertVariant::AssertNeq(_, _) => AssertVariant::AssertNeq(left, right),
+                            _ => unreachable!(),
+                        };
+                        (AssertStatement { variant, ..input }.into(), Default::default())
+                    }
+                }
+            }
+        }
+    }
+
     /// Modify assignments to tuples to become assignments to the corresponding variables.
     ///
     /// There are two cases we handle:

compiler/passes/src/destructuring/visitor.rs

@@ -18,6 +18,8 @@ use crate::CompilerState;
 
 use leo_ast::{
     AstReconstructor,
+    BinaryExpression,
+    BinaryOperation,
     DefinitionPlace,
     DefinitionStatement,
     Expression,
@@ -104,6 +106,34 @@ impl DestructuringVisitor<'_> {
         }
     }
 
+    /// Folds an iterator of expressions into a left-associated chain using `op`.
+    ///
+    /// Given expressions `[e1, e2, e3]`, this produces `((e1 op e2) op e3)`.
+    /// Each intermediate node is assigned a fresh ID and recorded as `Boolean`
+    /// in the type table.
+    ///
+    /// Panics if the iterator is empty.
+    pub fn fold_with_op<I>(&mut self, op: BinaryOperation, pieces: I) -> Expression
+    where
+        I: Iterator<Item = Expression>,
+    {
+        pieces
+            .reduce(|left, right| {
+                let expr: Expression = BinaryExpression {
+                    op,
+                    left,
+                    right,
+                    span: Default::default(),
+                    id: self.state.node_builder.next_id(),
+                }
+                .into();
+
+                self.state.type_table.insert(expr.id(), Type::Boolean);
+                expr
+            })
+            .expect("fold_with_op called with empty iterator")
+    }
+
     // Given the `expression` of tuple type, make a definition assigning variable to its members.
     //
     // That is, `let (mem1, mem2, mem3...) = expression;`

compiler/passes/src/option_lowering/ast.rs

@@ -161,13 +161,19 @@ impl leo_ast::AstReconstructor for OptionLoweringVisitor<'_> {
                     span: Span::default(),
                     id: self.state.node_builder.next_id(),
                 };
+                let Some(Type::Optional(OptionalType { inner })) = self.state.type_table.get(&optional_expr.id())
+                else {
+                    panic!("guaranteed by type checking");
+                };
+                self.state.type_table.insert(val_access.id(), *inner);
 
                 let is_some_access = MemberAccess {
                     inner: reconstructed_optional_expr.clone(),
                     name: Identifier::new(Symbol::intern("is_some"), self.state.node_builder.next_id()),
                     span: Span::default(),
                     id: self.state.node_builder.next_id(),
                 };
+                self.state.type_table.insert(is_some_access.id(), Type::Boolean);
 
                 // Create assertion: ensure `is_some` is `true`.
                 let assert_stmt = AssertStatement {
@@ -206,13 +212,15 @@ impl leo_ast::AstReconstructor for OptionLoweringVisitor<'_> {
                     span: Span::default(),
                     id: self.state.node_builder.next_id(),
                 };
+                self.state.type_table.insert(val_access.id(), *expected_inner_type.clone());
 
                 let is_some_access = MemberAccess {
                     inner: reconstructed_optional_expr,
                     name: Identifier::new(Symbol::intern("is_some"), self.state.node_builder.next_id()),
                     span: Span::default(),
                     id: self.state.node_builder.next_id(),
                 };
+                self.state.type_table.insert(is_some_access.id(), Type::Boolean);
 
                 // s.is_some ? s.val : fallback
                 let ternary_expr = TernaryExpression {
@@ -222,6 +230,7 @@ impl leo_ast::AstReconstructor for OptionLoweringVisitor<'_> {
                     span: Span::default(),
                     id: self.state.node_builder.next_id(),
                 };
+                self.state.type_table.insert(ternary_expr.id(), *expected_inner_type);
 
                 stmts1.extend(stmts2);
                 (ternary_expr.into(), stmts1)
@@ -504,25 +513,28 @@ impl leo_ast::AstReconstructor for OptionLoweringVisitor<'_> {
         mut input: TupleExpression,
         additional: &Option<Type>,
     ) -> (Expression, Self::AdditionalOutput) {
-        let mut all_stmts = Vec::new();
-        let mut new_elements = Vec::with_capacity(input.elements.len());
-
-        // Extract tuple element types if additional type info is Some(Type::Tuple).
+        // Determine the expected tuple element types
         let expected_types = additional
-            .as_ref()
-            .and_then(|ty| {
-                let mut ty = ty.clone();
-
-                // Unwrap Optional if any.
+            .clone()
+            .or_else(|| self.state.type_table.get(&input.id))
+            .and_then(|mut ty| {
+                // Unwrap Optional if any
                 if let Type::Optional(inner) = ty {
                     ty = *inner.inner;
                 }
-                // Expect Tuple type.
-                if let Type::Tuple(tuple_ty) = ty { Some(tuple_ty.elements.clone()) } else { None }
+
+                // Expect Tuple type
+                match ty {
+                    Type::Tuple(tuple_ty) => Some(tuple_ty.elements.clone()),
+                    _ => None,
+                }
             })
             .expect("guaranteed by type checking");
 
-        // Zip elements with expected types and reconstruct with expected type.
+        let mut all_stmts = Vec::new();
+        let mut new_elements = Vec::with_capacity(input.elements.len());
+
+        // Zip elements with expected types and reconstruct with expected type
         for (element, expected_ty) in input.elements.into_iter().zip(expected_types) {
             let (expr, mut stmts) = self.reconstruct_expression(element, &Some(expected_ty));
             all_stmts.append(&mut stmts);

compiler/passes/src/test_passes.rs

@@ -81,7 +81,7 @@ This structure ensures that **adding a new compiler pass test is minimal**:
 */
 
 use crate::*;
-use leo_ast::{Ast, NetworkName, NodeBuilder};
+use leo_ast::NetworkName;
 use leo_errors::{BufferEmitter, Handler};
 use leo_parser::parse_ast;
 use leo_span::{create_session_if_not_set_then, source_map::FileName, with_session_globals};
@@ -110,21 +110,11 @@ macro_rules! compiler_passes {
     };
 }
 
-/// Parse a Leo source program into an AST, returning errors via the handler.
-fn parse_program(source: &str, handler: &Handler) -> Result<Ast, ()> {
-    let node_builder = NodeBuilder::default();
-    let filename = FileName::Custom("test".into());
-
-    // Add the source to the session's source map
-    let source_file = with_session_globals(|s| s.source_map.new_source(source, filename));
-
-    handler.extend_if_error(parse_ast(handler.clone(), &node_builder, &source_file, &[], NetworkName::TestnetV0))
-}
-
 /// Macro to generate a single runner function for a compiler pass.
 ///
 /// Each runner:
 /// - Sets up a BufferEmitter and Handler for error/warning reporting.
+/// - Parse the test into an AST.
 /// - Runs the first three fixed passes: PathResolution, SymbolTableCreation, TypeChecking.
 /// - Runs the specified compiler pass.
 /// - Returns the resulting AST or formatted errors/warnings.
@@ -135,9 +125,16 @@ macro_rules! make_runner {
             let handler = Handler::new(buf.clone());
 
             create_session_if_not_set_then(|_| {
-                // Parse program into AST
-                let mut state = match parse_program(source, &handler) {
-                    Ok(ast) => CompilerState { ast, handler: handler.clone(), ..Default::default() },
+                let mut state = CompilerState { handler: handler.clone(), ..Default::default() };
+
+                state.ast = match handler.extend_if_error(parse_ast(
+                    handler.clone(),
+                    &state.node_builder,
+                    &with_session_globals(|s| s.source_map.new_source(source, FileName::Custom("test".into()))),
+                    &[],
+                    NetworkName::TestnetV0,
+                )) {
+                    Ok(ast) => ast,
                     Err(()) => return format!("{}{}", buf.extract_errs(), buf.extract_warnings()),
                 };