Preserve trace, error & todo formatting.

crates/aiken-lang/src/expr.rs

@@ -5,7 +5,8 @@ use vec1::Vec1;
 use crate::{
     ast::{
         Annotation, Arg, AssignmentKind, BinOp, CallArg, Clause, DefinitionLocation, IfBranch,
-        Pattern, RecordUpdateSpread, Span, TypedRecordUpdateArg, UnOp, UntypedRecordUpdateArg,
+        Pattern, RecordUpdateSpread, Span, TraceKind, TypedRecordUpdateArg, UnOp,
+        UntypedRecordUpdateArg,
     },
     builtins::void,
     tipo::{ModuleValueConstructor, PatternConstructor, Type, ValueConstructor},
@@ -375,6 +376,7 @@ pub enum UntypedExpr {
     },
 
     Trace {
+        kind: TraceKind,
         location: Span,
         then: Box<Self>,
         text: Box<Self>,
@@ -435,6 +437,7 @@ impl UntypedExpr {
     pub fn todo(location: Span, reason: Option<Self>) -> Self {
         UntypedExpr::Trace {
             location,
+            kind: TraceKind::Todo,
             then: Box::new(UntypedExpr::ErrorTerm { location }),
             text: Box::new(reason.unwrap_or_else(|| UntypedExpr::String {
                 location,

crates/aiken-lang/src/format.rs

@@ -7,12 +7,13 @@ use crate::{
     ast::{
         Annotation, Arg, ArgName, AssignmentKind, BinOp, CallArg, ClauseGuard, Constant, DataType,
         Definition, Function, IfBranch, ModuleConstant, Pattern, RecordConstructor,
-        RecordConstructorArg, RecordUpdateSpread, Span, TypeAlias, TypedArg, TypedConstant, UnOp,
+        RecordConstructorArg, RecordUpdateSpread, Span, TraceKind, TypeAlias, TypedArg,
-        UnqualifiedImport, UntypedArg, UntypedClause, UntypedClauseGuard, UntypedDefinition,
+        TypedConstant, UnOp, UnqualifiedImport, UntypedArg, UntypedClause, UntypedClauseGuard,
-        UntypedModule, UntypedPattern, UntypedRecordUpdateArg, Use, CAPTURE_VARIABLE,
+        UntypedDefinition, UntypedModule, UntypedPattern, UntypedRecordUpdateArg, Use,
+        CAPTURE_VARIABLE,
     },
     docvec,
-    expr::UntypedExpr,
+    expr::{UntypedExpr, DEFAULT_ERROR_STR, DEFAULT_TODO_STR},
     parser::extra::{Comment, ModuleExtra},
     pretty::{break_, concat, flex_break, join, line, lines, nil, Document, Documentable},
     tipo::{self, Type},
@@ -702,16 +703,9 @@ impl<'comments> Formatter<'comments> {
                 ..
             } => self.assignment(pattern, value, None, Some(*kind), annotation),
 
-            UntypedExpr::Trace { text, then, .. } => "trace"
+            UntypedExpr::Trace {
-                .to_doc()
+                kind, text, then, ..
-                .append(wrap_args([(self.wrap_expr(text), false)]))
+            } => self.trace(kind, text, then),
-                .group()
-                .append(if self.pop_empty_lines(then.start_byte_index()) {
-                    lines(2)
-                } else {
-                    line()
-                })
-                .append(self.expr(then)),
 
             UntypedExpr::When {
                 subjects, clauses, ..
@@ -755,6 +749,41 @@ impl<'comments> Formatter<'comments> {
         }
     }
 
+    pub fn trace<'a>(
+        &mut self,
+        kind: &'a TraceKind,
+        text: &'a UntypedExpr,
+        then: &'a UntypedExpr,
+    ) -> Document<'a> {
+        let (keyword, default_text) = match kind {
+            TraceKind::Trace => ("trace", None),
+            TraceKind::Error => ("error", Some(DEFAULT_ERROR_STR.to_string())),
+            TraceKind::Todo => ("todo", Some(DEFAULT_TODO_STR.to_string())),
+        };
+
+        let body = match text {
+            UntypedExpr::String { value, .. } if Some(value) == default_text.as_ref() => {
+                keyword.to_doc()
+            }
+            _ => keyword
+                .to_doc()
+                .append(" ")
+                .append(self.wrap_expr(text))
+                .group(),
+        };
+
+        match kind {
+            TraceKind::Error | TraceKind::Todo => body,
+            TraceKind::Trace => body
+                .append(if self.pop_empty_lines(then.start_byte_index()) {
+                    lines(2)
+                } else {
+                    line()
+                })
+                .append(self.expr(then)),
+        }
+    }
+
     pub fn pattern_constructor<'a>(
         &mut self,
         name: &'a str,
@@ -1305,7 +1334,10 @@ impl<'comments> Formatter<'comments> {
 
     fn wrap_expr<'a>(&mut self, expr: &'a UntypedExpr) -> Document<'a> {
         match expr {
-            UntypedExpr::Trace { .. }
+            UntypedExpr::Trace {
+                kind: TraceKind::Trace,
+                ..
+            }
             | UntypedExpr::Sequence { .. }
             | UntypedExpr::Assignment { .. } => "{"
                 .to_doc()
@@ -1344,7 +1376,10 @@ impl<'comments> Formatter<'comments> {
 
     fn case_clause_value<'a>(&mut self, expr: &'a UntypedExpr) -> Document<'a> {
         match expr {
-            UntypedExpr::Trace { .. }
+            UntypedExpr::Trace {
+                kind: TraceKind::Trace,
+                ..
+            }
             | UntypedExpr::Sequence { .. }
             | UntypedExpr::Assignment { .. } => " {"
                 .to_doc()

crates/aiken-lang/src/parser.rs

@@ -7,7 +7,7 @@ pub mod lexer;
 pub mod token;
 
 use crate::{
-    ast::{self, BinOp, Span, UnOp, UntypedDefinition, CAPTURE_VARIABLE},
+    ast::{self, BinOp, Span, TraceKind, UnOp, UntypedDefinition, CAPTURE_VARIABLE},
     expr,
 };
 
@@ -586,6 +586,7 @@ pub fn expr_seq_parser() -> impl Parser<Token, expr::UntypedExpr, Error = ParseE
                 .ignore_then(expr_parser(r.clone()))
                 .then(r.clone())
                 .map_with_span(|(text, then_), span| expr::UntypedExpr::Trace {
+                    kind: TraceKind::Trace,
                     location: span,
                     then: Box::new(then_),
                     text: Box::new(text),
@@ -605,17 +606,22 @@ pub fn todo_parser(
 ) -> impl Parser<Token, expr::UntypedExpr, Error = ParseError> + '_ {
     just(keyword.clone())
         .ignore_then(expr_parser(r.clone()).or_not())
-        .map_with_span(move |text, span| expr::UntypedExpr::Trace {
+        .map_with_span(move |text, span| {
+            let (kind, value) = match keyword {
+                Token::ErrorTerm => (TraceKind::Error, expr::DEFAULT_ERROR_STR.to_string()),
+                Token::Todo => (TraceKind::Todo, expr::DEFAULT_TODO_STR.to_string()),
+                _ => unreachable!(),
+            };
+
+            expr::UntypedExpr::Trace {
+                kind,
                 location: span,
                 then: Box::new(expr::UntypedExpr::ErrorTerm { location: span }),
-            text: Box::new(text.unwrap_or_else(|| expr::UntypedExpr::String {
+                text: Box::new(text.unwrap_or(expr::UntypedExpr::String {
                     location: span,
-                value: match keyword {
+                    value,
-                    Token::ErrorTerm => expr::DEFAULT_ERROR_STR.to_string(),
-                    Token::Todo => expr::DEFAULT_TODO_STR.to_string(),
-                    _ => unreachable!(),
-                },
                 })),
+            }
         })
 }
 

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

@@ -320,7 +320,7 @@ fn test_nested_function_calls() {
             ,
             when output.datum is {
               InlineDatum(_) -> True
-              _ -> error("expected inline datum")
+              _ -> error "expected inline datum"
             },
           ]
           |> list.and
@@ -339,7 +339,7 @@ fn test_nested_function_calls() {
             ),
             when output.datum is {
               InlineDatum(_) -> True
-              _ -> error("expected inline datum")
+              _ -> error "expected inline datum"
             },
           ]
           |> list.and
@@ -348,3 +348,34 @@ fn test_nested_function_calls() {
 
     assert_fmt(src, expected);
 }
+
+#[test]
+fn format_trace_todo_error() {
+    let src = indoc! {r#"
+        fn foo_1() {
+          todo
+        }
+
+        fn foo_2() {
+          todo "my custom message"
+        }
+
+        fn foo_3() {
+          when x is {
+            Foo -> True
+            _ -> error
+          }
+        }
+
+        fn foo_4() {
+          if 14 == 42 {
+            error "I don't think so"
+          } else {
+            trace "been there"
+            True
+          }
+        }
+    "#};
+
+    assert_fmt(src, src);
+}

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

@@ -309,6 +309,7 @@ fn empty_function() {
         vec![ast::UntypedDefinition::Fn(Function {
             arguments: vec![],
             body: expr::UntypedExpr::Trace {
+                kind: ast::TraceKind::Todo,
                 location: Span::new((), 0..15),
                 text: Box::new(expr::UntypedExpr::String {
                     value: "aiken::todo".to_string(),
@@ -1787,6 +1788,7 @@ fn function_def() {
             doc: None,
             arguments: vec![],
             body: expr::UntypedExpr::Trace {
+                kind: ast::TraceKind::Todo,
                 location: Span::new((), 0..11),
                 text: Box::new(expr::UntypedExpr::String {
                     value: "aiken::todo".to_string(),
@@ -2630,12 +2632,16 @@ fn trace_expressions() {
                         annotation: None,
                     },
                     expr::UntypedExpr::Trace {
+                        kind: ast::TraceKind::Trace,
                         location: Span::new((), 32..128),
                         then: Box::new(expr::UntypedExpr::Trace {
+                            kind: ast::TraceKind::Trace,
                             location: Span::new((), 49..128),
                             then: Box::new(expr::UntypedExpr::Trace {
+                                kind: ast::TraceKind::Trace,
                                 location: Span::new((), 62..128),
                                 then: Box::new(expr::UntypedExpr::Trace {
+                                    kind: ast::TraceKind::Trace,
                                     location: Span::new((), 97..128),
                                     then: Box::new(expr::UntypedExpr::Var {
                                         location: Span::new((), 124..128),
@@ -2735,6 +2741,7 @@ fn parse_keyword_error() {
             ast::Definition::Fn(Function {
                 arguments: vec![],
                 body: expr::UntypedExpr::Trace {
+                    kind: ast::TraceKind::Error,
                     location: Span::new((), 13..36),
                     then: Box::new(expr::UntypedExpr::ErrorTerm {
                         location: Span::new((), 13..36),
@@ -2789,6 +2796,7 @@ fn parse_keyword_error() {
                             alternative_patterns: vec![],
                             guard: None,
                             then: expr::UntypedExpr::Trace {
+                                kind: ast::TraceKind::Error,
                                 location: Span::new((), 100..105),
                                 then: Box::new(expr::UntypedExpr::ErrorTerm {
                                     location: Span::new((), 100..105),
@@ -2833,6 +2841,7 @@ fn parse_keyword_todo() {
             ast::Definition::Fn(Function {
                 arguments: vec![],
                 body: expr::UntypedExpr::Trace {
+                    kind: ast::TraceKind::Todo,
                     location: Span::new((), 13..35),
                     then: Box::new(expr::UntypedExpr::ErrorTerm {
                         location: Span::new((), 13..35),
@@ -2887,6 +2896,7 @@ fn parse_keyword_todo() {
                             alternative_patterns: vec![],
                             guard: None,
                             then: expr::UntypedExpr::Trace {
+                                kind: ast::TraceKind::Todo,
                                 location: Span::new((), 99..103),
                                 then: Box::new(expr::UntypedExpr::ErrorTerm {
                                     location: Span::new((), 99..103),

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

@@ -5,7 +5,7 @@ use vec1::Vec1;
 use crate::{
     ast::{
         Annotation, Arg, ArgName, AssignmentKind, BinOp, CallArg, Clause, ClauseGuard, Constant,
-        RecordUpdateSpread, Span, TypedArg, TypedCallArg, TypedClause, TypedClauseGuard,
+        RecordUpdateSpread, Span, TraceKind, TypedArg, TypedCallArg, TypedClause, TypedClauseGuard,
         TypedConstant, TypedIfBranch, TypedMultiPattern, TypedRecordUpdateArg, UnOp, UntypedArg,
         UntypedClause, UntypedClauseGuard, UntypedConstant, UntypedIfBranch, UntypedMultiPattern,
         UntypedPattern, UntypedRecordUpdateArg,
@@ -299,7 +299,8 @@ impl<'a, 'b> ExprTyper<'a, 'b> {
                 location,
                 then,
                 text,
-            } => self.infer_trace(*then, location, *text),
+                kind,
+            } => self.infer_trace(kind, *then, location, *text),
 
             UntypedExpr::When {
                 location,
@@ -1855,6 +1856,7 @@ impl<'a, 'b> ExprTyper<'a, 'b> {
 
     fn infer_trace(
         &mut self,
+        kind: TraceKind,
         then: UntypedExpr,
         location: Span,
         text: UntypedExpr,
@@ -1865,12 +1867,12 @@ impl<'a, 'b> ExprTyper<'a, 'b> {
         let then = self.infer(then)?;
         let tipo = then.tipo();
 
-        // TODO: reinstate once we can distinguish traces
+        if let TraceKind::Todo = kind {
-        //
+            self.environment.warnings.push(Warning::Todo {
-        // self.environment.warnings.push(Warning::Todo {
+                location,
-        //     location,
+                tipo: tipo.clone(),
-        //     tipo: tipo.clone(),
+            })
-        // })
+        }
 
         Ok(TypedExpr::Trace {
             location,

examples/acceptance_tests/032/lib/tests.ak

@@ -10,10 +10,10 @@ fn concat(left: String, right: String) -> String {
 
 fn is_negative(i: Int) -> Bool {
   if i < 0 {
-    trace("is negative")
+    trace "is negative"
     True
   } else {
-    trace(concat("is", concat(" ", "non-negative")))
+    trace concat("is", concat(" ", "non-negative"))
     False
   }
 }