From b16880a1702c8b1729c4dc0376537a9b34c50c3f Mon Sep 17 00:00:00 2001
From: microproofs <kwhitemsg@gmail.com>
Date: Wed, 13 Mar 2024 19:08:53 -0400
Subject: [PATCH] feat(annotation): not passing annotation into lambda when
 backpassing

Co-authored-by: Lucas Rosa <x@rvcas.dev>
---
 crates/aiken-lang/src/expr.rs        | 10 ++++--
 crates/aiken-lang/src/tests/check.rs | 49 ++++++++++++++++++++++++++++
 crates/aiken-lang/src/tipo/expr.rs   |  6 ++--
 3 files changed, 59 insertions(+), 6 deletions(-)

diff --git a/crates/aiken-lang/src/expr.rs b/crates/aiken-lang/src/expr.rs
index d159dc23..b97d608f 100644
--- a/crates/aiken-lang/src/expr.rs
+++ b/crates/aiken-lang/src/expr.rs
@@ -1299,16 +1299,20 @@ impl UntypedExpr {
         )
     }
 
-    pub fn lambda(names: Vec<String>, expressions: Vec<UntypedExpr>, location: Span) -> Self {
+    pub fn lambda(
+        names: Vec<(String, Option<Annotation>)>,
+        expressions: Vec<UntypedExpr>,
+        location: Span,
+    ) -> Self {
         Self::Fn {
             location,
             fn_style: FnStyle::Plain,
             arguments: names
                 .into_iter()
-                .map(|name| Arg {
+                .map(|(name, annotation)| Arg {
                     location,
                     doc: None,
-                    annotation: None,
+                    annotation,
                     tipo: (),
                     arg_name: ArgName::Named {
                         label: name.clone(),
diff --git a/crates/aiken-lang/src/tests/check.rs b/crates/aiken-lang/src/tests/check.rs
index 670ca9b6..6aac1771 100644
--- a/crates/aiken-lang/src/tests/check.rs
+++ b/crates/aiken-lang/src/tests/check.rs
@@ -1744,3 +1744,52 @@ fn discarded_let_bindings() {
         _ => unreachable!("ast isn't a Fn"),
     }
 }
+
+#[test]
+fn backpassing_type_annotation() {
+    let source_code = r#"
+        pub type Foo {
+          foo: Int,
+        }
+
+        fn transition_fold4(
+          inputs,
+          callback,
+        ) {
+          when inputs is {
+            [] -> {
+              (Foo(1), inputs)
+            }
+            [input, ..remaining_inputs] -> {
+        
+              callback(input)(
+                fn(foo) {
+                  transition_fold4(
+                    remaining_inputs,
+                    callback,
+                  )
+                },
+              )
+            }
+          }
+        }
+
+        pub fn backpassing(x) {
+          let input: Foo <-
+            transition_fold4(
+              x,
+            )
+          
+          fn(g){
+            g(if input.foo == 1{
+              1
+            } else {
+              2
+            })
+          }
+          
+        }
+    "#;
+
+    assert!(check(parse(source_code)).is_ok())
+}
diff --git a/crates/aiken-lang/src/tipo/expr.rs b/crates/aiken-lang/src/tipo/expr.rs
index 4b7c78ae..b54e32c0 100644
--- a/crates/aiken-lang/src/tipo/expr.rs
+++ b/crates/aiken-lang/src/tipo/expr.rs
@@ -1768,7 +1768,7 @@ impl<'a, 'b> ExprTyper<'a, 'b> {
             // (which is perhaps something we should support?).
             match pattern {
                 Pattern::Var { name, .. } | Pattern::Discard { name, .. } if kind.is_let() => {
-                    names.push(name.clone());
+                    names.push((name.clone(), annotation));
                 }
                 _ => {
                     let name = format!("{}_{}", ast::BACKPASS_VARIABLE, index);
@@ -1782,7 +1782,7 @@ impl<'a, 'b> ExprTyper<'a, 'b> {
                                 name: name.clone(),
                             }
                             .into(),
-                            patterns: AssignmentPattern::new(pattern, annotation).into(),
+                            patterns: AssignmentPattern::new(pattern, annotation.clone()).into(),
                             // erase backpassing while preserving assignment kind.
                             kind: match kind {
                                 AssignmentKind::Let { .. } => AssignmentKind::let_(),
@@ -1791,7 +1791,7 @@ impl<'a, 'b> ExprTyper<'a, 'b> {
                         },
                     );
 
-                    names.push(name);
+                    names.push((name, annotation));
                 }
             }
         }