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

@@ -134,31 +134,32 @@ impl Annotation {
 #[derive(Debug)]
 pub(crate) struct CanDefs {
     defs: Vec<Option<Def>>,
-    dbgs: ExpectsOrDbgs,
-    expects: ExpectsOrDbgs,
-    expects_fx: ExpectsOrDbgs,
+    dbgs: OrderDependentStatements,
+    expects: OrderDependentStatements,
+    expects_fx: OrderDependentStatements,
+    returns: OrderDependentStatements,
     def_ordering: DefOrdering,
     aliases: VecMap<Symbol, Alias>,
 }
 
 #[derive(Clone, Debug)]
-pub struct ExpectsOrDbgs {
-    pub conditions: Vec<Expr>,
+pub struct OrderDependentStatements {
+    pub expressions: Vec<Expr>,
     pub regions: Vec<Region>,
     pub preceding_comment: Vec<Region>,
 }
 
-impl ExpectsOrDbgs {
+impl OrderDependentStatements {
     fn with_capacity(capacity: usize) -> Self {
         Self {
-            conditions: Vec::with_capacity(capacity),
+            expressions: Vec::with_capacity(capacity),
             regions: Vec::with_capacity(capacity),
             preceding_comment: Vec::with_capacity(capacity),
         }
     }
 
     fn push(&mut self, loc_can_condition: Loc<Expr>, preceding_comment: Region) {
-        self.conditions.push(loc_can_condition.value);
+        self.expressions.push(loc_can_condition.value);
         self.regions.push(loc_can_condition.region);
         self.preceding_comment.push(preceding_comment);
     }
@@ -303,9 +304,10 @@ impl PendingTypeDef<'_> {
 pub enum Declaration {
     Declare(Def),
     DeclareRec(Vec<Def>, IllegalCycleMark),
+    Return(Expr, Region),
     Builtin(Def),
-    Expects(ExpectsOrDbgs),
-    ExpectsFx(ExpectsOrDbgs),
+    Expects(OrderDependentStatements),
+    ExpectsFx(OrderDependentStatements),
     /// If we know a cycle is illegal during canonicalization.
     /// Otherwise we will try to detect this during solving; see [`IllegalCycleMark`].
     InvalidCycle(Vec<CycleEntry>),
@@ -317,6 +319,7 @@ impl Declaration {
         match self {
             Declare(_) => 1,
             DeclareRec(defs, _) => defs.len(),
+            Return(_, _) => 0,
             InvalidCycle { .. } => 0,
             Builtin(_) => 0,
             Expects(_) => 0,
@@ -340,6 +343,7 @@ impl Declaration {
                 expects.regions.first().unwrap(),
                 expects.regions.last().unwrap(),
             ),
+            Declaration::Return(_return_expr, return_region) => *return_region,
         }
     }
 }
@@ -1137,6 +1141,7 @@ fn canonicalize_value_defs<'a>(
     let mut pending_dbgs = Vec::with_capacity(value_defs.len());
     let mut pending_expects = Vec::with_capacity(value_defs.len());
     let mut pending_expect_fx = Vec::with_capacity(value_defs.len());
+    let mut pending_returns = Vec::with_capacity(value_defs.len());
 
     let mut imports_introduced = Vec::with_capacity(value_defs.len());
 
@@ -1159,6 +1164,9 @@ fn canonicalize_value_defs<'a>(
             PendingValue::ExpectFx(pending_expect) => {
                 pending_expect_fx.push(pending_expect);
             }
+            PendingValue::Return(pending_return) => {
+                pending_returns.push(pending_return);
+            }
             PendingValue::ModuleImport(PendingModuleImport {
                 module_id,
                 region,
@@ -1235,9 +1243,10 @@ fn canonicalize_value_defs<'a>(
         def_ordering.insert_symbol_references(def_id as u32, &temp_output.references)
     }
 
-    let mut dbgs = ExpectsOrDbgs::with_capacity(pending_dbgs.len());
-    let mut expects = ExpectsOrDbgs::with_capacity(pending_expects.len());
-    let mut expects_fx = ExpectsOrDbgs::with_capacity(pending_expects.len());
+    let mut dbgs = OrderDependentStatements::with_capacity(pending_dbgs.len());
+    let mut expects = OrderDependentStatements::with_capacity(pending_expects.len());
+    let mut expects_fx = OrderDependentStatements::with_capacity(pending_expects.len());
+    let mut returns = OrderDependentStatements::with_capacity(pending_returns.len());
 
     for pending in pending_dbgs {
         let (loc_can_condition, can_output) = canonicalize_expr(
@@ -1281,11 +1290,21 @@ fn canonicalize_value_defs<'a>(
         output.union(can_output);
     }
 
+    for pending in pending_returns {
+        let (loc_return_expr, can_output) =
+            canonicalize_expr(env, var_store, scope, pending.region, &pending.value);
+
+        returns.push(loc_return_expr, Region::zero());
+
+        output.union(can_output);
+    }
+
     let can_defs = CanDefs {
         defs,
         dbgs,
         expects,
         expects_fx,
+        returns,
         def_ordering,
         aliases,
     };
@@ -1681,7 +1700,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,
@@ -1694,10 +1713,17 @@ pub(crate) fn sort_can_defs_new(
         dbgs: _,
         expects,
         expects_fx,
+        returns,
         def_ordering,
         aliases,
     } = defs;
 
+    for return_region in returns.regions {
+        env.problem(Problem::ReturnOutsideOfFunction {
+            region: return_region,
+        });
+    }
+
     // TODO: inefficient, but I want to make this what CanDefs contains in the future
     let mut defs: Vec<_> = defs.into_iter().map(|x| x.unwrap()).collect();
 
@@ -1712,7 +1738,7 @@ pub(crate) fn sort_can_defs_new(
     // because of the ordering of declarations, expects should come first because they are
     // independent, but can rely on all other top-level symbols in the module
     let it = expects
-        .conditions
+        .expressions
         .into_iter()
         .zip(expects.regions)
         .zip(expects.preceding_comment);
@@ -1725,7 +1751,7 @@ pub(crate) fn sort_can_defs_new(
     }
 
     let it = expects_fx
-        .conditions
+        .expressions
         .into_iter()
         .zip(expects_fx.regions)
         .zip(expects_fx.preceding_comment);
@@ -1829,6 +1855,7 @@ pub(crate) fn sort_can_defs_new(
                                         None,
                                     );
                                 }
+                                // TODO: do I need to handle returns here?
                                 _ => {
                                     declarations.push_value_def(
                                         Loc::at(def.loc_pattern.region, symbol),
@@ -1975,6 +2002,7 @@ pub(crate) fn sort_can_defs(
         dbgs,
         expects,
         expects_fx,
+        returns,
         def_ordering,
         aliases,
     } = defs;
@@ -2100,18 +2128,24 @@ pub(crate) fn sort_can_defs(
         }
     }
 
-    if !dbgs.conditions.is_empty() {
+    if !dbgs.expressions.is_empty() {
         declarations.push(Declaration::Expects(dbgs));
     }
 
-    if !expects.conditions.is_empty() {
+    if !expects.expressions.is_empty() {
         declarations.push(Declaration::Expects(expects));
     }
 
-    if !expects_fx.conditions.is_empty() {
+    if !expects_fx.expressions.is_empty() {
         declarations.push(Declaration::ExpectsFx(expects_fx));
     }
 
+    if !returns.expressions.is_empty() {
+        for (return_expr, return_region) in returns.expressions.into_iter().zip(returns.regions) {
+            declarations.push(Declaration::Return(return_expr, return_region));
+        }
+    }
+
     (declarations, output)
 }
 
@@ -2354,6 +2388,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 +2606,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.)
     //
@@ -2603,7 +2640,13 @@ fn canonicalize_pending_body<'a>(
 
                 // 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,
+                    Some(tail_symbol) if tail_symbol == *defined_symbol => {
+                        if closure_data.early_returns.is_empty() {
+                            Recursive::TailRecursive
+                        } else {
+                            Recursive::Recursive
+                        }
+                    }
                     _ => Recursive::NotRecursive,
                 };
 
@@ -2664,7 +2707,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);
@@ -2731,7 +2773,7 @@ pub fn can_defs_with_return<'a>(
     let mut loc_expr: Loc<Expr> = ret_expr;
 
     for declaration in declarations.into_iter().rev() {
-        loc_expr = decl_to_let(declaration, loc_expr);
+        loc_expr = decl_to_let_or_return(declaration, loc_expr, var_store);
     }
 
     (loc_expr.value, output)
@@ -2758,7 +2800,11 @@ pub fn report_unused_imports(
     }
 }
 
-fn decl_to_let(decl: Declaration, loc_ret: Loc<Expr>) -> Loc<Expr> {
+fn decl_to_let_or_return(
+    decl: Declaration,
+    loc_ret: Loc<Expr>,
+    var_store: &mut VarStore,
+) -> Loc<Expr> {
     match decl {
         Declaration::Declare(def) => {
             let region = Region::span_across(&def.loc_pattern.region, &loc_ret.region);
@@ -2770,6 +2816,17 @@ fn decl_to_let(decl: Declaration, loc_ret: Loc<Expr>) -> Loc<Expr> {
             let expr = Expr::LetRec(defs, Box::new(loc_ret), cycle_mark);
             Loc::at(region, expr)
         }
+        Declaration::Return(return_expr, return_region) => {
+            let region = Region::span_across(&return_region, &loc_ret.region);
+
+            let return_var = var_store.fresh();
+            let expr = Expr::Return {
+                return_value: Box::new(Loc::at(return_region, return_expr)),
+                return_var,
+            };
+
+            Loc::at(region, expr)
+        }
         Declaration::InvalidCycle(entries) => {
             Loc::at_zero(Expr::RuntimeError(RuntimeError::CircularDef(entries)))
         }
@@ -2780,7 +2837,7 @@ fn decl_to_let(decl: Declaration, loc_ret: Loc<Expr>) -> Loc<Expr> {
         Declaration::Expects(expects) => {
             let mut loc_ret = loc_ret;
 
-            let conditions = expects.conditions.into_iter().rev();
+            let conditions = expects.expressions.into_iter().rev();
             let condition_regions = expects.regions.into_iter().rev();
             let expect_regions = expects.preceding_comment.into_iter().rev();
 
@@ -3005,6 +3062,7 @@ enum PendingValue<'a> {
     Dbg(PendingExpectOrDbg<'a>),
     Expect(PendingExpectOrDbg<'a>),
     ExpectFx(PendingExpectOrDbg<'a>),
+    Return(&'a Loc<ast::Expr<'a>>),
     ModuleImport(PendingModuleImport<'a>),
     SignatureDefMismatch,
     InvalidIngestedFile,
@@ -3154,6 +3212,8 @@ fn to_pending_value_def<'a>(
             preceding_comment: *preceding_comment,
         }),
 
+        Return(return_expr) => PendingValue::Return(return_expr),
+
         ModuleImport(module_import) => {
             let qualified_module_name: QualifiedModuleName = module_import.name.value.into();
             let module_name = qualified_module_name.module.clone();