Allow to trace expressions (and not only string literals)

This however enforces that the argument unifies to a `String`. So this
  is more flexible than the previous form, but does fundamentally the
  same thing.

  Fixes #378.

crates/aiken-lang/src/air.rs

@@ -249,7 +249,6 @@ pub enum Air {
 
     Trace {
         scope: Vec<u64>,
-        text: Option<String>,
         tipo: Arc<Type>,
     },
 }

crates/aiken-lang/src/builder.rs

@@ -1685,12 +1685,12 @@ pub fn monomorphize(
                     needs_variant = true;
                 }
             }
-            Air::Trace { scope, text, tipo } => {
+            Air::Trace { scope, tipo } => {
                 if tipo.is_generic() {
                     let mut tipo = tipo.clone();
                     find_generics_to_replace(&mut tipo, &generic_types);
 
-                    new_air[index] = Air::Trace { scope, tipo, text };
+                    new_air[index] = Air::Trace { scope, tipo };
                     needs_variant = true;
                 }
             }

crates/aiken-lang/src/expr.rs

@@ -96,7 +96,7 @@ pub enum TypedExpr {
         location: Span,
         tipo: Arc<Type>,
         then: Box<Self>,
-        text: Option<String>,
+        text: Box<Self>,
     },
 
     When {
@@ -389,7 +389,7 @@ pub enum UntypedExpr {
     Trace {
         location: Span,
         then: Box<Self>,
-        text: Option<String>,
+        text: Box<Self>,
     },
 
     When {

crates/aiken-lang/src/format.rs

@@ -705,22 +705,10 @@ impl<'comments> Formatter<'comments> {
                 ..
             } => self.assignment(pattern, value, None, Some(*kind), annotation),
 
-            UntypedExpr::Trace {
+            UntypedExpr::Trace { text, then, .. } => "trace"
-                text: None, then, ..
-            } => "trace"
                 .to_doc()
-                .append(if self.pop_empty_lines(then.start_byte_index()) {
+                .append(wrap_args([(self.wrap_expr(text), false)]))
-                    lines(2)
+                .group()
-                } else {
-                    line()
-                })
-                .append(self.expr(then)),
-
-            UntypedExpr::Trace {
-                text: Some(l),
-                then,
-                ..
-            } => docvec!["trace(\"", l, "\")"]
                 .append(if self.pop_empty_lines(then.start_byte_index()) {
                     lines(2)
                 } else {

crates/aiken-lang/src/parser.rs

@@ -592,15 +592,14 @@ pub fn expr_seq_parser() -> impl Parser<Token, expr::UntypedExpr, Error = ParseE
         choice((
             just(Token::Trace)
                 .ignore_then(
-                    select! {Token::String {value} => value}
+                    expr_parser(r.clone())
-                        .delimited_by(just(Token::LeftParen), just(Token::RightParen))
+                        .delimited_by(just(Token::LeftParen), just(Token::RightParen)),
-                        .or_not(),
                 )
                 .then(r.clone())
                 .map_with_span(|(text, then_), span| expr::UntypedExpr::Trace {
                     location: span,
                     then: Box::new(then_),
-                    text,
+                    text: Box::new(text),
                 }),
             expr_parser(r.clone())
                 .then(r.repeated())

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

@@ -126,3 +126,35 @@ fn list_pattern_6() {
     "#;
     assert!(check(parse(source_code)).is_ok())
 }
+
+#[test]
+fn trace_strings() {
+    let source_code = r#"
+        fn bar() {
+            "BAR"
+        }
+
+        test foo() {
+            let msg1 = "FOO"
+            trace("INLINE")
+            trace(msg1)
+            trace(bar())
+            True
+        }
+    "#;
+    assert!(check(parse(source_code)).is_ok())
+}
+
+#[test]
+fn trace_non_strings() {
+    let source_code = r#"
+        test foo() {
+            trace(14 + 42)
+            True
+        }
+    "#;
+    assert!(matches!(
+        check(parse(source_code)),
+        Err((_, Error::CouldNotUnify { .. }))
+    ))
+}

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

@@ -2586,3 +2586,121 @@ fn scope_logical_expression() {
         })],
     )
 }
+
+#[test]
+fn trace_expressions() {
+    let code = indoc! {r#"
+          fn foo() {
+            let msg1 = "FOO"
+            trace "INLINE"
+            trace msg1
+            trace string.concat(msg1, "BAR")
+            trace ( 14 + 42 * 1337 )
+            Void
+          }
+        "#};
+    assert_definitions(
+        code,
+        vec![ast::Definition::Fn(Function {
+            arguments: vec![],
+            body: expr::UntypedExpr::Sequence {
+                location: Span::new((), 13..128),
+                expressions: vec![
+                    expr::UntypedExpr::Assignment {
+                        location: Span::new((), 13..29),
+                        value: Box::new(expr::UntypedExpr::String {
+                            location: Span::new((), 24..29),
+                            value: "FOO".to_string(),
+                        }),
+                        pattern: ast::Pattern::Var {
+                            location: Span::new((), 17..21),
+                            name: "msg1".to_string(),
+                        },
+                        kind: ast::AssignmentKind::Let,
+                        annotation: None,
+                    },
+                    expr::UntypedExpr::Trace {
+                        location: Span::new((), 32..128),
+                        then: Box::new(expr::UntypedExpr::Trace {
+                            location: Span::new((), 49..128),
+                            then: Box::new(expr::UntypedExpr::Trace {
+                                location: Span::new((), 62..128),
+                                then: Box::new(expr::UntypedExpr::Trace {
+                                    location: Span::new((), 97..128),
+                                    then: Box::new(expr::UntypedExpr::Var {
+                                        location: Span::new((), 124..128),
+                                        name: "Void".to_string(),
+                                    }),
+                                    text: Box::new(expr::UntypedExpr::BinOp {
+                                        location: Span::new((), 105..119),
+                                        name: ast::BinOp::AddInt,
+                                        left: Box::new(expr::UntypedExpr::Int {
+                                            location: Span::new((), 105..107),
+                                            value: "14".to_string(),
+                                        }),
+                                        right: Box::new(expr::UntypedExpr::BinOp {
+                                            location: Span::new((), 110..119),
+                                            name: ast::BinOp::MultInt,
+                                            left: Box::new(expr::UntypedExpr::Int {
+                                                location: Span::new((), 110..112),
+                                                value: "42".to_string(),
+                                            }),
+                                            right: Box::new(expr::UntypedExpr::Int {
+                                                location: Span::new((), 115..119),
+                                                value: "1337".to_string(),
+                                            }),
+                                        }),
+                                    }),
+                                }),
+                                text: Box::new(expr::UntypedExpr::Call {
+                                    arguments: vec![
+                                        ast::CallArg {
+                                            label: None,
+                                            location: Span::new((), 82..86),
+                                            value: expr::UntypedExpr::Var {
+                                                location: Span::new((), 82..86),
+                                                name: "msg1".to_string(),
+                                            },
+                                        },
+                                        ast::CallArg {
+                                            label: None,
+                                            location: Span::new((), 88..93),
+                                            value: expr::UntypedExpr::String {
+                                                location: Span::new((), 88..93),
+                                                value: "BAR".to_string(),
+                                            },
+                                        },
+                                    ],
+                                    fun: Box::new(expr::UntypedExpr::FieldAccess {
+                                        location: Span::new((), 68..81),
+                                        label: "concat".to_string(),
+                                        container: Box::new(expr::UntypedExpr::Var {
+                                            location: Span::new((), 68..74),
+                                            name: "string".to_string(),
+                                        }),
+                                    }),
+                                    location: Span::new((), 68..94),
+                                }),
+                            }),
+                            text: Box::new(expr::UntypedExpr::Var {
+                                location: Span::new((), 55..59),
+                                name: "msg1".to_string(),
+                            }),
+                        }),
+                        text: Box::new(expr::UntypedExpr::String {
+                            location: Span::new((), 38..46),
+                            value: "INLINE".to_string(),
+                        }),
+                    },
+                ],
+            },
+            doc: None,
+            location: Span::new((), 0..8),
+            name: "foo".to_string(),
+            public: false,
+            return_annotation: None,
+            return_type: (),
+            end_position: 129,
+        })],
+    )
+}

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

@@ -309,7 +309,7 @@ impl<'a, 'b> ExprTyper<'a, 'b> {
                 location,
                 then,
                 text,
-            } => self.infer_trace(*then, location, text),
+            } => self.infer_trace(*then, location, *text),
 
             UntypedExpr::When {
                 location,
@@ -1887,17 +1887,19 @@ impl<'a, 'b> ExprTyper<'a, 'b> {
         &mut self,
         then: UntypedExpr,
         location: Span,
-        text: Option<String>,
+        text: UntypedExpr,
     ) -> Result<TypedExpr, Error> {
-        let then = self.infer(then)?;
+        let text = self.infer(text)?;
+        self.unify(text.tipo(), string(), text.location(), false)?;
 
+        let then = self.infer(then)?;
         let tipo = then.tipo();
 
         Ok(TypedExpr::Trace {
             location,
             tipo,
             then: Box::new(then),
-            text,
+            text: Box::new(text),
         })
     }
 

crates/aiken-lang/src/uplc.rs

@@ -651,19 +651,21 @@ impl<'a> CodeGenerator<'a> {
 
                 ir_stack.append(&mut elems_air);
             }
+
             TypedExpr::Trace {
                 tipo, then, text, ..
             } => {
                 let mut scope = scope;
 
                 ir_stack.push(Air::Trace {
-                    text: text.clone(),
                     tipo: tipo.clone(),
                     scope: scope.clone(),
                 });
 
                 scope.push(self.id_gen.next());
+                self.build_ir(text, ir_stack, scope.clone());
 
+                scope.push(self.id_gen.next());
                 self.build_ir(then, ir_stack, scope);
             }
 
@@ -3504,13 +3506,12 @@ impl<'a> CodeGenerator<'a> {
                         label,
                     };
                 }
-                Air::Trace { tipo, scope, text } => {
+                Air::Trace { tipo, scope } => {
                     let mut replaced_type = tipo.clone();
                     replace_opaque_type(&mut replaced_type, self.data_types.clone());
 
                     ir_stack[index] = Air::Trace {
                         scope,
-                        text,
                         tipo: replaced_type,
                     };
                 }
@@ -5631,19 +5632,13 @@ impl<'a> CodeGenerator<'a> {
 
                 arg_stack.push(term);
             }
-            Air::Trace { text, .. } => {
+            Air::Trace { .. } => {
+                let text = arg_stack.pop().unwrap();
+
                 let term = arg_stack.pop().unwrap();
 
                 let term = apply_wrap(
-                    apply_wrap(
+                    apply_wrap(Term::Builtin(DefaultFunction::Trace).force_wrap(), text),
-                        Term::Builtin(DefaultFunction::Trace).force_wrap(),
-                        Term::Constant(
-                            UplcConstant::String(
-                                text.unwrap_or_else(|| "aiken::trace".to_string()),
-                            )
-                            .into(),
-                        ),
-                    ),
                     Term::Delay(term.into()),
                 )
                 .force_wrap();

examples/acceptance_tests/032/lib/tests.ak

@@ -1,9 +1,19 @@
+use aiken/builtin
+
+fn concat(left: String, right: String) -> String {
+  builtin.append_bytearray(
+    builtin.encode_utf8(left),
+    builtin.encode_utf8(right),
+  )
+  |> builtin.decode_utf8
+}
+
 fn is_negative(i: Int) -> Bool {
   if i < 0 {
     trace("is negative")
     True
   } else {
-    trace("is non-negative")
+    trace(concat("is", concat(" ", "non-negative")))
     False
   }
 }