Fix compiler crash around dangling expect/let in traces

Fixes #1029.

crates/aiken-lang/src/expr.rs

@@ -203,6 +203,28 @@ impl<T> From<Vec1Ref<T>> for Vec1<T> {
 }
 
 impl TypedExpr {
+    pub fn and_then(self, next: Self) -> Self {
+        if let TypedExpr::Trace {
+            tipo,
+            location,
+            then,
+            text,
+        } = self
+        {
+            return TypedExpr::Trace {
+                tipo,
+                location,
+                then: Box::new(then.and_then(next)),
+                text,
+            };
+        }
+
+        TypedExpr::Sequence {
+            location: self.location(),
+            expressions: vec![self, next],
+        }
+    }
+
     pub fn sequence(exprs: &[TypedExpr]) -> Self {
         TypedExpr::Sequence {
             location: Span::empty(),

crates/aiken-lang/src/tests/check.rs

@@ -3295,3 +3295,50 @@ fn softcasting_unused_let_binding() {
     let (warnings, _) = result.unwrap();
     assert!(warnings.is_empty(), "should not contain any warnings");
 }
+
+#[test]
+fn dangling_trace_let_standalone() {
+    let source_code = r#"
+        test foo() {
+          trace @"foo"
+          let True = True
+        }
+    "#;
+
+    assert!(matches!(
+        check_validator(parse(source_code)),
+        Err((_, Error::LastExpressionIsAssignment { .. }))
+    ))
+}
+
+#[test]
+fn dangling_trace_let_in_sequence() {
+    let source_code = r#"
+        test foo() {
+          let predicate = True
+          trace @"foo"
+          let result = predicate
+        }
+    "#;
+
+    assert!(matches!(
+        check_validator(parse(source_code)),
+        Err((_, Error::LastExpressionIsAssignment { .. }))
+    ))
+}
+
+#[test]
+fn dangling_trace_let_in_trace() {
+    let source_code = r#"
+        test foo() {
+          trace @"foo"
+          trace @"bar"
+          let result = True
+        }
+    "#;
+
+    assert!(matches!(
+        check_validator(parse(source_code)),
+        Err((_, Error::LastExpressionIsAssignment { .. }))
+    ))
+}

crates/aiken-lang/src/tipo/expr.rs

@@ -1565,10 +1565,7 @@ impl<'a, 'b> ExprTyper<'a, 'b> {
             let then = if let Some(filler) =
                 recover_from_no_assignment(assert_no_assignment(&then), then.location())?
             {
-                TypedExpr::Sequence {
-                    location,
-                    expressions: vec![scope.infer(then)?, filler],
-                }
+                scope.infer(then)?.and_then(filler)
             } else {
                 scope.infer(then)?
             };
@@ -1638,10 +1635,7 @@ impl<'a, 'b> ExprTyper<'a, 'b> {
         let typed_final_else = if let Some(filler) =
             recover_from_no_assignment(assert_no_assignment(&final_else), final_else.location())?
         {
-            TypedExpr::Sequence {
-                location: final_else.location(),
-                expressions: vec![self.infer(final_else)?, filler],
-            }
+            self.infer(final_else)?.and_then(filler)
         } else {
             self.infer(final_else)?
         };
@@ -1696,10 +1690,7 @@ impl<'a, 'b> ExprTyper<'a, 'b> {
                     assert_no_assignment(&branch.body),
                     branch.body.location(),
                 )? {
-                    TypedExpr::Sequence {
-                        location: branch.body.location(),
-                        expressions: vec![typer.infer(branch.body.clone())?, filler],
-                    }
+                    typer.infer(branch.body.clone())?.and_then(filler)
                 } else {
                     typer.infer(branch.body.clone())?
                 };
@@ -1720,10 +1711,7 @@ impl<'a, 'b> ExprTyper<'a, 'b> {
                     assert_no_assignment(&branch.body),
                     branch.body.location(),
                 )? {
-                    TypedExpr::Sequence {
-                        location: branch.body.location(),
-                        expressions: vec![self.infer(branch.body.clone())?, filler],
-                    }
+                    self.infer(branch.body.clone())?.and_then(filler)
                 } else {
                     self.infer(branch.body.clone())?
                 };
@@ -1815,10 +1803,7 @@ impl<'a, 'b> ExprTyper<'a, 'b> {
             body_infer.map_err(|e| e.with_unify_error_rigid_names(&body_rigid_names));
 
         let body = if let Some(filler) = recover_from_no_assignment(no_assignment, location)? {
-            TypedExpr::Sequence {
-                location,
-                expressions: vec![inferred_body?, filler],
-            }
+            inferred_body?.and_then(filler)
         } else {
             inferred_body?
         };
@@ -2206,8 +2191,12 @@ impl<'a, 'b> ExprTyper<'a, 'b> {
                         if let Some(filler) =
                             recover_from_no_assignment(no_assignment, typed_expression.location())?
                         {
-                            expressions.push(typed_expression);
-                            expressions.push(filler);
+                            match typed_expression.and_then(filler) {
+                                TypedExpr::Sequence {
+                                    expressions: seq, ..
+                                } => expressions.extend(seq),
+                                trace => expressions.push(trace),
+                            }
                         } else {
                             expressions.push(typed_expression);
                         }
@@ -2391,6 +2380,7 @@ impl<'a, 'b> ExprTyper<'a, 'b> {
             .collect::<Result<Vec<_>, Error>>()?;
 
         let then = self.infer(then)?;
+
         let tipo = then.tipo();
 
         if let TraceKind::Todo = kind {