Merge pull request #7193 from smores56/try-desugar

Implement the `try` keyword with desugaring

crates/compiler/can/src/desugar.rs

@@ -23,6 +23,10 @@ use roc_region::all::{Loc, Region};
 //
 // Thank you, Markus!
 
+/// Desugar a single binary operation.
+///
+/// When using this function, don't desugar `left` and `right` before calling so that
+/// we can properly desugar `|> try` expressions!
 fn new_op_call_expr<'a>(
     env: &mut Env<'a>,
     scope: &mut Scope,
@@ -33,10 +37,43 @@ fn new_op_call_expr<'a>(
     let region = Region::span_across(&left.region, &right.region);
 
     let value = match loc_op.value {
+        // Rewrite the Pizza operator into an Apply
         Pizza => {
-            // Rewrite the Pizza operator into an Apply
+            // Allow `left |> try (optional)` to desugar to `try left (optional)`
+            let right_without_spaces = without_spaces(&right.value);
+            match right_without_spaces {
+                Try => {
+                    let desugared_left = desugar_expr(env, scope, left);
+                    return desugar_try_expr(env, scope, desugared_left);
+                }
+                Apply(&Loc { value: Try, .. }, arguments, _called_via) => {
+                    let try_fn = desugar_expr(env, scope, arguments.first().unwrap());
+
+                    let mut args = Vec::with_capacity_in(arguments.len(), env.arena);
+                    args.push(desugar_expr(env, scope, left));
+                    args.extend(
+                        arguments
+                            .iter()
+                            .skip(1)
+                            .map(|a| desugar_expr(env, scope, a)),
+                    );
+
+                    return desugar_try_expr(
+                        env,
+                        scope,
+                        env.arena.alloc(Loc::at(
+                            right.region,
+                            Expr::Apply(try_fn, args.into_bump_slice(), CalledVia::Try),
+                        )),
+                    );
+                }
+                _ => {}
+            }
 
-            match &right.value {
+            let left = desugar_expr(env, scope, left);
+            let right = desugar_expr(env, scope, right);
+
+            match right.value {
                 Apply(function, arguments, _called_via) => {
                     let mut args = Vec::with_capacity_in(1 + arguments.len(), env.arena);
 
@@ -55,6 +92,9 @@ fn new_op_call_expr<'a>(
             }
         }
         binop => {
+            let left = desugar_expr(env, scope, left);
+            let right = desugar_expr(env, scope, right);
+
             // This is a normal binary operator like (+), so desugar it
             // into the appropriate function call.
             let (module_name, ident) = binop_to_function(binop);
@@ -73,6 +113,13 @@ fn new_op_call_expr<'a>(
     Loc { region, value }
 }
 
+fn without_spaces<'a>(expr: &'a Expr<'a>) -> &'a Expr<'a> {
+    match expr {
+        Expr::SpaceBefore(inner, _) | Expr::SpaceAfter(inner, _) => without_spaces(inner),
+        _ => expr,
+    }
+}
+
 fn desugar_value_def<'a>(
     env: &mut Env<'a>,
     scope: &mut Scope,
@@ -351,7 +398,8 @@ pub fn desugar_expr<'a>(
         | OptionalFieldInRecordBuilder { .. }
         | Tag(_)
         | OpaqueRef(_)
-        | Crash => loc_expr,
+        | Crash
+        | Try => loc_expr,
 
         Str(str_literal) => match str_literal {
             StrLiteral::PlainLine(_) => loc_expr,
@@ -870,6 +918,31 @@ pub fn desugar_expr<'a>(
                 })
             }
         }
+        Apply(
+            Loc {
+                value: Try,
+                region: _,
+            },
+            loc_args,
+            _called_via,
+        ) => {
+            let result_expr = if loc_args.len() == 1 {
+                desugar_expr(env, scope, loc_args[0])
+            } else {
+                let function = desugar_expr(env, scope, loc_args.first().unwrap());
+                let mut desugared_args = Vec::with_capacity_in(loc_args.len() - 1, env.arena);
+                for loc_arg in &loc_args[1..] {
+                    desugared_args.push(desugar_expr(env, scope, loc_arg));
+                }
+
+                env.arena.alloc(Loc::at(
+                    loc_expr.region,
+                    Expr::Apply(function, desugared_args.into_bump_slice(), CalledVia::Try),
+                ))
+            };
+
+            env.arena.alloc(desugar_try_expr(env, scope, result_expr))
+        }
         Apply(loc_fn, loc_args, called_via) => {
             let mut desugared_args = Vec::with_capacity_in(loc_args.len(), env.arena);
 
@@ -1042,6 +1115,72 @@ pub fn desugar_expr<'a>(
     }
 }
 
+pub fn desugar_try_expr<'a>(
+    env: &mut Env<'a>,
+    scope: &mut Scope,
+    result_expr: &'a Loc<Expr<'a>>,
+) -> Loc<Expr<'a>> {
+    let region = result_expr.region;
+    let ok_symbol = env.arena.alloc(scope.gen_unique_symbol_name().to_string());
+    let err_symbol = env.arena.alloc(scope.gen_unique_symbol_name().to_string());
+
+    let ok_branch = env.arena.alloc(WhenBranch {
+        patterns: env.arena.alloc([Loc::at(
+            region,
+            Pattern::Apply(
+                env.arena.alloc(Loc::at(region, Pattern::Tag("Ok"))),
+                env.arena
+                    .alloc([Loc::at(region, Pattern::Identifier { ident: ok_symbol })]),
+            ),
+        )]),
+        value: Loc::at(
+            region,
+            Expr::Var {
+                module_name: "",
+                ident: ok_symbol,
+            },
+        ),
+        guard: None,
+    });
+
+    let err_branch = env.arena.alloc(WhenBranch {
+        patterns: env.arena.alloc([Loc::at(
+            region,
+            Pattern::Apply(
+                env.arena.alloc(Loc::at(region, Pattern::Tag("Err"))),
+                env.arena
+                    .alloc([Loc::at(region, Pattern::Identifier { ident: err_symbol })]),
+            ),
+        )]),
+        value: Loc::at(
+            region,
+            Expr::Return(
+                env.arena.alloc(Loc::at(
+                    region,
+                    Expr::Apply(
+                        env.arena.alloc(Loc::at(region, Expr::Tag("Err"))),
+                        &*env.arena.alloc([&*env.arena.alloc(Loc::at(
+                            region,
+                            Expr::Var {
+                                module_name: "",
+                                ident: err_symbol,
+                            },
+                        ))]),
+                        CalledVia::Try,
+                    ),
+                )),
+                None,
+            ),
+        ),
+        guard: None,
+    });
+
+    Loc::at(
+        region,
+        Expr::When(result_expr, &*env.arena.alloc([&*ok_branch, &*err_branch])),
+    )
+}
+
 fn desugar_str_segments<'a>(
     env: &mut Env<'a>,
     scope: &mut Scope,
@@ -1407,7 +1546,7 @@ fn desugar_bin_ops<'a>(
     let mut op_stack: Vec<Loc<BinOp>> = Vec::with_capacity_in(lefts.len(), env.arena);
 
     for (loc_expr, loc_op) in lefts {
-        arg_stack.push(desugar_expr(env, scope, loc_expr));
+        arg_stack.push(loc_expr);
         match run_binop_step(
             env,
             scope,
@@ -1421,7 +1560,7 @@ fn desugar_bin_ops<'a>(
         }
     }
 
-    let mut expr = desugar_expr(env, scope, right);
+    let mut expr = right;
 
     for (left, loc_op) in arg_stack.into_iter().zip(op_stack.into_iter()).rev() {
         expr = env

crates/compiler/can/src/expr.rs

@@ -1132,6 +1132,11 @@ pub fn canonicalize_expr<'a>(
         ast::Expr::TrySuffix { .. } => internal_error!(
             "a Expr::TrySuffix expression was not completely removed in desugar_value_def_suffixed"
         ),
+        ast::Expr::Try => {
+            // Treat remaining `try` keywords as normal variables so that we can continue to support `Result.try`
+            canonicalize_var_lookup(env, var_store, scope, "", "try", region)
+        }
+
         ast::Expr::Tag(tag) => {
             let variant_var = var_store.fresh();
             let ext_var = var_store.fresh();
@@ -2547,6 +2552,7 @@ pub fn is_valid_interpolation(expr: &ast::Expr<'_>) -> bool {
         | ast::Expr::RecordUpdater(_)
         | ast::Expr::Crash
         | ast::Expr::Dbg
+        | ast::Expr::Try
         | ast::Expr::Underscore(_)
         | ast::Expr::MalformedIdent(_, _)
         | ast::Expr::Tag(_)