Preserve holes discard names in function captures

Fixes #1080

CHANGELOG.md

@@ -15,6 +15,7 @@
 
 - **aiken**: panic error when using `aiken uplc decode` on cbor encoded flat bytes. @rvcas
 - **aiken-lang**: comment formatting in pipelines leading to confusion. @rvcas
+- **aiken-lang**: preserve holes discard name in function captures (see [#1080](https://github.com/aiken-lang/aiken/issues/1080)). @KtorZ
 
 ## v1.1.9 - 2024-12-13
 

crates/aiken-lang/src/ast.rs

@@ -1107,19 +1107,6 @@ impl TypedArg {
         self.arg_name.get_name()
     }
 
-    pub fn is_capture(&self) -> bool {
-        if let ArgName::Named {
-            ref name, location, ..
-        } = self.arg_name
-        {
-            return name.starts_with(CAPTURE_VARIABLE)
-                && location == Span::empty()
-                && self.location == Span::empty();
-        }
-
-        false
-    }
-
     pub fn find_node(&self, byte_index: usize) -> Option<Located<'_>> {
         if self.arg_name.location().contains(byte_index) {
             Some(Located::Argument(&self.arg_name, self.tipo.clone()))

crates/aiken-lang/src/expr.rs

@@ -763,6 +763,10 @@ impl UntypedExpr {
         UntypedExpr::do_reify_constant(&mut IndexMap::new(), data_types, cst, tipo)
     }
 
+    pub fn is_discard(&self) -> bool {
+        matches!(self, UntypedExpr::Var { name, ..} if name.starts_with("_"))
+    }
+
     // Reify some opaque 'PlutusData' into an 'UntypedExpr', using a Type annotation. We also need
     // an extra map to lookup record & enum constructor's names as they're completely erased when
     // in their PlutusData form, and the Type annotation only contains type name.
@@ -1336,17 +1340,24 @@ impl UntypedExpr {
                     value: Some(value),
                     label,
                     location,
-                } => CallArg {
+                } if !value.is_discard() => CallArg {
                     value,
                     label,
                     location,
                 },
                 CallArg {
-                    value: None,
+                    value,
                     label,
                     location,
                 } => {
-                    let name = format!("{}__{index}", ast::CAPTURE_VARIABLE);
+                    let name = format!(
+                        "{}__{index}_{}",
+                        ast::CAPTURE_VARIABLE,
+                        match value {
+                            Some(UntypedExpr::Var { ref name, .. }) => name,
+                            _ => "_",
+                        }
+                    );
 
                     holes.push(ast::UntypedArg {
                         location: Span::empty(),
@@ -1354,7 +1365,7 @@ impl UntypedExpr {
                         doc: None,
                         by: ArgBy::ByName(ast::ArgName::Named {
                             label: name.clone(),
-                            name,
+                            name: name.clone(),
                             location: Span::empty(),
                         }),
                         is_validator_param: false,
@@ -1363,10 +1374,7 @@ impl UntypedExpr {
                     ast::CallArg {
                         label,
                         location,
-                        value: UntypedExpr::Var {
-                            location,
-                            name: format!("{}__{index}", ast::CAPTURE_VARIABLE),
-                        },
+                        value: UntypedExpr::Var { location, name },
                     }
                 }
             })

crates/aiken-lang/src/format.rs

@@ -992,7 +992,9 @@ impl<'comments> Formatter<'comments> {
                 }
             }
 
-            UntypedExpr::Var { name, .. } if name.contains(CAPTURE_VARIABLE) => "_".to_doc(),
+            UntypedExpr::Var { name, .. } if name.contains(CAPTURE_VARIABLE) => "_"
+                .to_doc()
+                .append(name.split('_').last().unwrap_or_default()),
 
             UntypedExpr::Var { name, .. } => name.to_doc(),
 
@@ -1560,8 +1562,14 @@ impl<'comments> Formatter<'comments> {
                 ..
             } => match args.as_slice() {
                 [first, second] if is_breakable_expr(&second.value) && first.is_capture_hole() => {
+                    let discard_name = match first.value {
+                        UntypedExpr::Var { ref name, .. } => name.split("_").last().unwrap_or("_"),
+                        _ => "",
+                    };
                     self.expr(fun, false)
-                        .append("(_, ")
+                        .append("(_")
+                        .append(discard_name)
+                        .append(", ")
                         .append(self.call_arg(second, false))
                         .append(")")
                         .group()

crates/aiken-lang/src/parser/chain/call.rs

@@ -1,11 +1,10 @@
-use chumsky::prelude::*;
-
 use super::Chain;
 use crate::{
     ast::CallArg,
     expr::UntypedExpr,
     parser::{token::Token, ParseError},
 };
+use chumsky::prelude::*;
 
 pub(crate) fn parser(
     expression: Recursive<'_, Token, UntypedExpr, ParseError>,
@@ -23,11 +22,11 @@ pub(crate) fn parser(
         select! { Token::Name { name } => name }
             .then_ignore(just(Token::Colon))
             .or_not()
-            .then_ignore(select! {Token::DiscardName {name} => name })
-            .map_with_span(|label, location| CallArg {
+            .then(select! {Token::DiscardName {name} => name })
+            .map_with_span(|(label, name), location| CallArg {
                 location,
                 label,
-                value: None,
+                value: Some(UntypedExpr::Var { location, name }),
             }),
     ))
     .separated_by(just(Token::Comma))

crates/aiken-lang/src/parser/expr/record.rs

@@ -1,5 +1,3 @@
-use chumsky::prelude::*;
-
 use crate::{
     ast,
     expr::UntypedExpr,
@@ -8,6 +6,7 @@ use crate::{
         token::Token,
     },
 };
+use chumsky::prelude::*;
 
 pub fn parser(
     r: Recursive<'_, Token, UntypedExpr, ParseError>,

crates/aiken-lang/src/parser/expr/snapshots/function_invoke.snap

@@ -50,8 +50,8 @@ Sequence {
                     UntypedArg {
                         by: ByName(
                             Named {
-                                name: "_capture__0",
-                                label: "_capture__0",
+                                name: "_capture__0__",
+                                label: "_capture__0__",
                                 location: 0..0,
                             },
                         ),
@@ -68,7 +68,7 @@ Sequence {
                             location: 45..46,
                             value: Var {
                                 location: 45..46,
-                                name: "_capture__0",
+                                name: "_capture__0__",
                             },
                         },
                         CallArg {

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

@@ -1436,3 +1436,16 @@ fn comment_in_pipeline() {
         "#
     );
 }
+
+#[test]
+fn capture_right_hand_side_assign() {
+    assert_format!(
+        r#"
+        fn foo() {
+            let (_aa, bb, _cc) = bar(foo: _a, _b, _)
+            let _ = baz(_d, [1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20,21,22,23])
+            bb
+        }
+        "#
+    );
+}

crates/aiken-lang/src/tests/snapshots/capture_right_hand_side_assign.snap

@@ -0,0 +1,14 @@
+---
+source: crates/aiken-lang/src/tests/format.rs
+description: "Code:\n\nfn foo() {\n    let (_aa, bb, _cc) = bar(foo: _a, _b, _)\n    let _ = baz(_d, [1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20,21,22,23])\n    bb\n}\n"
+---
+fn foo() {
+  let (_aa, bb, _cc) =
+    bar(foo: _a, _b, _)
+  let _ =
+    baz(_d, [
+      1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21,
+      22, 23,
+    ])
+  bb
+}