Implement most of the recent round of PR feedback

crates/compiler/can/src/def.rs

@@ -137,7 +137,6 @@ pub(crate) struct CanDefs {
     dbgs: OrderDependentStatements,
     expects: OrderDependentStatements,
     expects_fx: OrderDependentStatements,
-    returns: OrderDependentStatements,
     def_ordering: DefOrdering,
     aliases: VecMap<Symbol, Alias>,
 }
@@ -304,7 +303,6 @@ impl PendingTypeDef<'_> {
 pub enum Declaration {
     Declare(Def),
     DeclareRec(Vec<Def>, IllegalCycleMark),
-    Return(Expr, Region),
     Builtin(Def),
     Expects(OrderDependentStatements),
     ExpectsFx(OrderDependentStatements),
@@ -319,7 +317,6 @@ impl Declaration {
         match self {
             Declare(_) => 1,
             DeclareRec(defs, _) => defs.len(),
-            Return(_, _) => 0,
             InvalidCycle { .. } => 0,
             Builtin(_) => 0,
             Expects(_) => 0,
@@ -343,7 +340,6 @@ impl Declaration {
                 expects.regions.first().unwrap(),
                 expects.regions.last().unwrap(),
             ),
-            Declaration::Return(_return_expr, return_region) => *return_region,
         }
     }
 }
@@ -1141,7 +1137,6 @@ 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());
 
@@ -1164,9 +1159,6 @@ 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,
@@ -1246,7 +1238,6 @@ fn canonicalize_value_defs<'a>(
     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(
@@ -1290,21 +1281,11 @@ 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,
     };
@@ -1713,17 +1694,10 @@ pub(crate) fn sort_top_level_can_defs(
         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();
 
@@ -1855,7 +1829,6 @@ pub(crate) fn sort_top_level_can_defs(
                                         None,
                                     );
                                 }
-                                // TODO: do I need to handle returns here?
                                 _ => {
                                     declarations.push_value_def(
                                         Loc::at(def.loc_pattern.region, symbol),
@@ -2002,7 +1975,6 @@ pub(crate) fn sort_can_defs(
         dbgs,
         expects,
         expects_fx,
-        returns,
         def_ordering,
         aliases,
     } = defs;
@@ -2140,12 +2112,6 @@ pub(crate) fn sort_can_defs(
         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)
 }
 
@@ -2773,7 +2739,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_or_return(declaration, loc_expr, var_store);
+        loc_expr = decl_to_let(declaration, loc_expr);
     }
 
     (loc_expr.value, output)
@@ -2800,11 +2766,7 @@ pub fn report_unused_imports(
     }
 }
 
-fn decl_to_let_or_return(
-    decl: Declaration,
-    loc_ret: Loc<Expr>,
-    var_store: &mut VarStore,
-) -> Loc<Expr> {
+fn decl_to_let<'a>(decl: Declaration, loc_ret: Loc<Expr>) -> Loc<Expr> {
     match decl {
         Declaration::Declare(def) => {
             let region = Region::span_across(&def.loc_pattern.region, &loc_ret.region);
@@ -2816,17 +2778,6 @@ fn decl_to_let_or_return(
             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)))
         }
@@ -3062,7 +3013,6 @@ enum PendingValue<'a> {
     Dbg(PendingExpectOrDbg<'a>),
     Expect(PendingExpectOrDbg<'a>),
     ExpectFx(PendingExpectOrDbg<'a>),
-    Return(&'a Loc<ast::Expr<'a>>),
     ModuleImport(PendingModuleImport<'a>),
     SignatureDefMismatch,
     InvalidIngestedFile,
@@ -3212,8 +3162,6 @@ 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();

crates/compiler/can/src/desugar.rs

@@ -177,12 +177,6 @@ fn desugar_value_def<'a>(
                 body_expr: desugar_expr(env, scope, stmt_expr),
             }
         }
-
-        Return(return_expr) => {
-            let desugared_return_expr = &*env.arena.alloc(desugar_expr(env, scope, return_expr));
-
-            Return(desugared_return_expr)
-        }
     }
 }
 
@@ -324,12 +318,6 @@ pub fn desugar_value_def_suffixed<'a>(arena: &'a Bump, value_def: ValueDef<'a>)
         // TODO support desugaring of Dbg and ExpectFx
         Dbg { .. } | ExpectFx { .. } => value_def,
         ModuleImport { .. } | IngestedFileImport(_) => value_def,
-        Return(ret_expr) => match unwrap_suffixed_expression(arena, ret_expr, None) {
-            Ok(new_ret_expr) => ValueDef::Return(new_ret_expr),
-            Err(..) => {
-                internal_error!("Unable to desugar the suffix inside a Return value def");
-            }
-        },
 
         Stmt(..) => {
             internal_error!(

crates/compiler/can/src/expr.rs

@@ -1280,8 +1280,11 @@ pub fn canonicalize_expr<'a>(
             );
 
             if let Some(after_return) = after_return {
+                let region_with_return =
+                    Region::span_across(&return_expr.region, &after_return.region);
+
                 env.problem(Problem::StatementsAfterReturn {
-                    region: after_return.region,
+                    region: region_with_return,
                 });
             }
 
@@ -1649,6 +1652,17 @@ fn canonicalize_closure_body<'a>(
         }
     }
 
+    let final_expr = match &loc_body_expr.value {
+        Expr::LetRec(_, final_expr, _) | Expr::LetNonRec(_, final_expr) => &final_expr.value,
+        _ => &loc_body_expr.value,
+    };
+
+    if let Expr::Return { return_value, .. } = final_expr {
+        env.problem(Problem::ReturnAtEndOfFunction {
+            region: return_value.region,
+        });
+    }
+
     // store the references of this function in the Env. This information is used
     // when we canonicalize a surrounding def (if it exists)
     env.closures.insert(symbol, output.references.clone());

crates/compiler/can/src/suffixed.rs

@@ -680,7 +680,7 @@ pub fn unwrap_suffixed_expression_defs_help<'a>(
                 };
 
                 let maybe_suffixed_value_def = match current_value_def {
-                    Annotation(..) | Dbg{..} | Expect{..} | ExpectFx{..} | Return(_) | Stmt(..) | ModuleImport{..} | IngestedFileImport(_) => None,
+                    Annotation(..) | Dbg{..} | Expect{..} | ExpectFx{..} | Stmt(..) | ModuleImport{..} | IngestedFileImport(_) => None,
                     AnnotatedBody { body_pattern, body_expr, ann_type, ann_pattern, .. } => Some((body_pattern, body_expr, Some((ann_pattern, ann_type)))),
                     Body (def_pattern, def_expr) => Some((def_pattern, def_expr, None)),
                 };

crates/compiler/fmt/src/def.rs

@@ -423,7 +423,6 @@ impl<'a> Formattable for ValueDef<'a> {
             ModuleImport(module_import) => module_import.is_multiline(),
             IngestedFileImport(ingested_file_import) => ingested_file_import.is_multiline(),
             Stmt(loc_expr) => loc_expr.is_multiline(),
-            Return(loc_expr) => loc_expr.is_multiline(),
         }
     }
 
@@ -465,7 +464,6 @@ impl<'a> Formattable for ValueDef<'a> {
             ModuleImport(module_import) => module_import.format(buf, indent),
             IngestedFileImport(ingested_file_import) => ingested_file_import.format(buf, indent),
             Stmt(loc_expr) => loc_expr.format_with_options(buf, parens, newlines, indent),
-            Return(loc_expr) => loc_expr.format_with_options(buf, parens, newlines, indent),
         }
     }
 }

crates/compiler/load_internal/src/docs.rs

@@ -281,9 +281,6 @@ fn generate_entry_docs(
                 ValueDef::IngestedFileImport { .. } => {
                     // Don't generate docs for ingested file imports
                 }
-                ValueDef::Return { .. } => {
-                    // Don't generate docs for `return`s
-                }
 
                 ValueDef::Stmt(loc_expr) => {
                     if let roc_parse::ast::Expr::Var {

crates/compiler/parse/src/ast.rs

@@ -837,8 +837,6 @@ pub enum ValueDef<'a> {
     IngestedFileImport(IngestedFileImport<'a>),
 
     Stmt(&'a Loc<Expr<'a>>),
-
-    Return(&'a Loc<Expr<'a>>),
 }
 
 impl<'a> ValueDef<'a> {
@@ -1090,7 +1088,6 @@ impl<'a, 'b> Iterator for RecursiveValueDefIter<'a, 'b> {
                             }
                         }
                         ValueDef::Stmt(loc_expr) => self.push_pending_from_expr(&loc_expr.value),
-                        ValueDef::Return(loc_expr) => self.push_pending_from_expr(&loc_expr.value),
                         ValueDef::Annotation(_, _) | ValueDef::IngestedFileImport(_) => {}
                     }
 
@@ -2737,7 +2734,6 @@ impl<'a> Malformed for ValueDef<'a> {
                 annotation,
             }) => path.is_malformed() || annotation.is_malformed(),
             ValueDef::Stmt(loc_expr) => loc_expr.is_malformed(),
-            ValueDef::Return(loc_expr) => loc_expr.is_malformed(),
         }
     }
 }

crates/compiler/parse/src/expr.rs

@@ -2664,8 +2664,9 @@ fn return_help<'a>(options: ExprParseOptions) -> impl Parser<'a, Stmt<'a>, ERetu
 
         let region = Region::span_across(&return_kw.region, &return_value.region);
 
-        let stmt = Stmt::ValueDef(ValueDef::Return(
+        let stmt = Stmt::Expr(Expr::Return(
             arena.alloc(Loc::at(region, return_value.value)),
+            None,
         ));
 
         Ok((MadeProgress, stmt, state))
@@ -3057,7 +3058,6 @@ fn stmts_to_expr<'a>(
                     CalledVia::Space,
                 )
             }
-            Stmt::ValueDef(ValueDef::Return(return_value)) => Expr::Return(return_value, None),
             Stmt::ValueDef(ValueDef::Expect { .. }) => {
                 return Err(EExpr::Expect(
                     EExpect::Continuation(
@@ -3090,6 +3090,20 @@ fn stmts_to_defs<'a>(
     while i < stmts.len() {
         let sp_stmt = stmts[i];
         match sp_stmt.item.value {
+            Stmt::Expr(Expr::Return(return_value, _after_return)) => {
+                if i == stmts.len() - 1 {
+                    last_expr = Some(Loc::at_zero(Expr::Return(return_value, None)));
+                } else {
+                    let rest = stmts_to_expr(&stmts[i + 1..], arena)?;
+                    last_expr = Some(Loc::at_zero(Expr::Return(
+                        return_value,
+                        Some(arena.alloc(rest)),
+                    )));
+                }
+
+                // don't re-process the rest of the statements, they got consumed by the early return
+                break;
+            }
             Stmt::Expr(e) => {
                 if is_expr_suffixed(&e) && i + 1 < stmts.len() {
                     defs.push_value_def(
@@ -3112,20 +3126,6 @@ fn stmts_to_defs<'a>(
                     last_expr = Some(sp_stmt.item.with_value(e));
                 }
             }
-            Stmt::ValueDef(ValueDef::Return(return_value)) => {
-                if i == stmts.len() - 1 {
-                    last_expr = Some(Loc::at_zero(Expr::Return(return_value, None)));
-                } else {
-                    let rest = stmts_to_expr(&stmts[i + 1..], arena)?;
-                    last_expr = Some(Loc::at_zero(Expr::Return(
-                        return_value,
-                        Some(arena.alloc(rest)),
-                    )));
-                }
-
-                // don't re-process the rest of the statements, they got consumed by the early return
-                break;
-            }
             Stmt::Backpassing(pats, call) => {
                 if last_expr.is_some() {
                     return Err(EExpr::StmtAfterExpr(sp_stmt.item.region.start()));

crates/compiler/parse/src/normalize.rs

@@ -440,7 +440,6 @@ impl<'a> Normalize<'a> for ValueDef<'a> {
                 IngestedFileImport(ingested_file_import.normalize(arena))
             }
             Stmt(loc_expr) => Stmt(arena.alloc(loc_expr.normalize(arena))),
-            Return(loc_expr) => Return(arena.alloc(loc_expr.normalize(arena))),
         }
     }
 }