Implement the return keyword

crates/cli/src/main.rs

@@ -238,7 +238,7 @@ fn main() -> io::Result<()> {
                             current_block = String::new();
                         } else if in_roc_block {
                             current_block.push_str(&line);
-                            current_block.push_str("\n");
+                            current_block.push('\n');
                         }
                     }
 

crates/compiler/can/src/builtins.rs

@@ -446,6 +446,7 @@ fn defn_help(
         function_type: var_store.fresh(),
         closure_type: var_store.fresh(),
         return_type: ret_var,
+        early_returns: vec![],
         name: fn_name,
         captured_symbols: Vec::new(),
         recursive: Recursive::NotRecursive,

crates/compiler/can/src/constraint.rs

@@ -97,6 +97,7 @@ impl Constraints {
             Category::List,
             Category::Str,
             Category::Character,
+            Category::Return,
         ]);
 
         pattern_categories.extend([
@@ -150,6 +151,7 @@ impl Constraints {
     pub const CATEGORY_LIST: Index<Category> = Index::new(11);
     pub const CATEGORY_STR: Index<Category> = Index::new(12);
     pub const CATEGORY_CHARACTER: Index<Category> = Index::new(13);
+    pub const CATEGORY_RETURN: Index<Category> = Index::new(14);
 
     pub const PCATEGORY_RECORD: Index<PatternCategory> = Index::new(0);
     pub const PCATEGORY_EMPTYRECORD: Index<PatternCategory> = Index::new(1);

crates/compiler/can/src/copy.rs

@@ -456,6 +456,7 @@ fn deep_copy_expr_help<C: CopyEnv>(env: &mut C, copied: &mut Vec<Variable>, expr
             function_type,
             closure_type,
             return_type,
+            early_returns,
             name,
             captured_symbols,
             recursive,
@@ -465,6 +466,10 @@ fn deep_copy_expr_help<C: CopyEnv>(env: &mut C, copied: &mut Vec<Variable>, expr
             function_type: sub!(*function_type),
             closure_type: sub!(*closure_type),
             return_type: sub!(*return_type),
+            early_returns: early_returns
+                .iter()
+                .map(|(var, region)| (sub!(*var), *region))
+                .collect(),
             name: *name,
             captured_symbols: captured_symbols
                 .iter()
@@ -688,6 +693,14 @@ fn deep_copy_expr_help<C: CopyEnv>(env: &mut C, copied: &mut Vec<Variable>, expr
             lookups_in_cond: lookups_in_cond.to_vec(),
         },
 
+        Return {
+            return_value,
+            return_var,
+        } => Return {
+            return_value: Box::new(return_value.map(|e| go_help!(e))),
+            return_var: sub!(*return_var),
+        },
+
         Dbg {
             source_location,
             source,

crates/compiler/can/src/debug/pretty_print.rs

@@ -449,6 +449,7 @@ fn expr<'a>(c: &Ctx, p: EPrec, f: &'a Arena<'a>, e: &'a Expr) -> DocBuilder<'a,
         Dbg { .. } => todo!(),
         Expect { .. } => todo!(),
         ExpectFx { .. } => todo!(),
+        Return { .. } => todo!(),
         TypedHole(_) => todo!(),
         RuntimeError(_) => todo!(),
     }

crates/compiler/can/src/def.rs

@@ -1681,7 +1681,7 @@ impl DefOrdering {
 }
 
 #[inline(always)]
-pub(crate) fn sort_can_defs_new(
+pub(crate) fn sort_top_level_can_defs(
     env: &mut Env<'_>,
     scope: &mut Scope,
     var_store: &mut VarStore,
@@ -2354,6 +2354,7 @@ fn canonicalize_pending_value_def<'a>(
                         function_type: var_store.fresh(),
                         closure_type: var_store.fresh(),
                         return_type: var_store.fresh(),
+                        early_returns: scope.early_returns.clone(),
                         name: symbol,
                         captured_symbols: Vec::new(),
                         recursive: Recursive::NotRecursive,
@@ -2571,6 +2572,8 @@ fn canonicalize_pending_body<'a>(
         loc_value = value;
     }
 
+    let expr_var = var_store.fresh();
+
     // We treat closure definitions `foo = \a, b -> ...` differently from other body expressions,
     // because they need more bookkeeping (for tail calls, closure captures, etc.)
     //
@@ -2602,9 +2605,17 @@ fn canonicalize_pending_body<'a>(
                 env.tailcallable_symbol = outer_tailcallable;
 
                 // The closure is self tail recursive iff it tail calls itself (by defined name).
-                let is_recursive = match can_output.tail_call {
-                    Some(tail_symbol) if tail_symbol == *defined_symbol => Recursive::TailRecursive,
-                    _ => Recursive::NotRecursive,
+                let is_recursive = if can_output
+                    .early_tail_calls
+                    .iter()
+                    .chain(std::iter::once(&can_output.final_tail_call))
+                    .all(|tail_call| {
+                        matches!(tail_call, Some(tail_symbol) if tail_symbol == defined_symbol)
+                    })
+                {
+                    Recursive::TailRecursive
+                } else {
+                    Recursive::NotRecursive
                 };
 
                 closure_data.recursive = is_recursive;
@@ -2664,7 +2675,6 @@ fn canonicalize_pending_body<'a>(
         }
     };
 
-    let expr_var = var_store.fresh();
     let mut vars_by_symbol = SendMap::default();
 
     pattern_to_vars_by_symbol(&mut vars_by_symbol, &loc_can_pattern.value, expr_var);

crates/compiler/can/src/desugar.rs

@@ -315,7 +315,7 @@ pub fn desugar_value_def_suffixed<'a>(arena: &'a Bump, value_def: ValueDef<'a>)
             }
         },
 
-        // TODO support desugaring of Dbg, Expect, and ExpectFx
+        // TODO support desugaring of Dbg and ExpectFx
         Dbg { .. } | ExpectFx { .. } => value_def,
         ModuleImport { .. } | IngestedFileImport(_) => value_def,
 
@@ -1008,6 +1008,7 @@ pub fn desugar_expr<'a>(
         Expect(condition, continuation) => {
             let desugared_condition = &*env.arena.alloc(desugar_expr(env, scope, condition));
             let desugared_continuation = &*env.arena.alloc(desugar_expr(env, scope, continuation));
+
             env.arena.alloc(Loc {
                 value: Expect(desugared_condition, desugared_continuation),
                 region: loc_expr.region,
@@ -1019,14 +1020,22 @@ pub fn desugar_expr<'a>(
         }
         DbgStmt(condition, continuation) => {
             let desugared_condition = &*env.arena.alloc(desugar_expr(env, scope, condition));
-
             let desugared_continuation = &*env.arena.alloc(desugar_expr(env, scope, continuation));
 
             env.arena.alloc(Loc {
                 value: *desugar_dbg_stmt(env, desugared_condition, desugared_continuation),
                 region: loc_expr.region,
             })
         }
+        Return(return_value, after_return) => {
+            let desugared_return_value = &*env.arena.alloc(desugar_expr(env, scope, return_value));
+
+            env.arena.alloc(Loc {
+                // Do not desugar after_return since it isn't run anyway
+                value: Return(desugared_return_value, *after_return),
+                region: loc_expr.region,
+            })
+        }
 
         // note this only exists after desugaring
         LowLevelDbg(_, _, _) => loc_expr,