fix unwrapping of trailing expr

crates/compiler/can/src/desugar.rs

@@ -9,8 +9,8 @@ use roc_module::called_via::{BinOp, CalledVia};
 use roc_module::ident::ModuleName;
 use roc_parse::ast::Expr::{self, *};
 use roc_parse::ast::{
-    AssignedField, Collection, Pattern, RecordBuilderField, StrLiteral, StrSegment, ValueDef,
+    wrap_in_task_ok, AssignedField, Collection, Pattern, RecordBuilderField, StrLiteral,
-    WhenBranch,
+    StrSegment, ValueDef, WhenBranch,
 };
 use roc_region::all::{LineInfo, Loc, Region};
 
@@ -189,33 +189,16 @@ pub fn desugar_value_def_suffixed<'a>(arena: &'a Bump, value_def: ValueDef<'a>)
                     sub_arg,
                     sub_pat,
                     sub_new,
-                }) => {
+                }) => Body(
-                    let ok_wrapped_return = arena.alloc(Loc::at(
+                    loc_pattern,
+                    apply_task_await(
+                        arena,
                         loc_expr.region,
-                        Expr::Apply(
+                        sub_arg,
-                            arena.alloc(Loc::at(
+                        sub_pat,
-                                loc_expr.region,
+                        wrap_in_task_ok(arena, sub_new),
-                                Expr::Var {
+                    ),
-                                    module_name: ModuleName::TASK,
+                ),
-                                    ident: "ok",
-                                    suffixed: 0,
-                                },
-                            )),
-                            arena.alloc([sub_new]),
-                            CalledVia::BangSuffix,
-                        ),
-                    ));
-                    Body(
-                        loc_pattern,
-                        apply_task_await(
-                            arena,
-                            loc_expr.region,
-                            sub_arg,
-                            sub_pat,
-                            ok_wrapped_return,
-                        ),
-                    )
-                }
                 Err(..) => Body(
                     loc_pattern,
                     arena.alloc(Loc::at(loc_expr.region, MalformedSuffixed(loc_expr))),
@@ -248,7 +231,13 @@ pub fn desugar_value_def_suffixed<'a>(arena: &'a Bump, value_def: ValueDef<'a>)
                     ann_type,
                     comment,
                     body_pattern,
-                    body_expr: apply_task_await(arena, body_expr.region, sub_arg, sub_pat, sub_new),
+                    body_expr: apply_task_await(
+                        arena,
+                        body_expr.region,
+                        sub_arg,
+                        sub_pat,
+                        wrap_in_task_ok(arena, sub_new),
+                    ),
                 },
                 Err(..) => AnnotatedBody {
                     ann_pattern,

crates/compiler/can/src/suffixed.rs

@@ -6,7 +6,7 @@ use roc_error_macros::internal_error;
 use roc_module::called_via::CalledVia;
 use roc_module::ident::ModuleName;
 use roc_parse::ast::Expr::{self, *};
-use roc_parse::ast::{is_loc_expr_suffixed, Pattern, ValueDef};
+use roc_parse::ast::{is_loc_expr_suffixed, wrap_in_task_ok, Pattern, ValueDef};
 use roc_region::all::{Loc, Region};
 use std::cell::Cell;
 
@@ -398,28 +398,12 @@ pub fn unwrap_suffixed_expression_if_then_else_help<'a>(
                             sub_pat,
                             sub_new,
                         }) => {
-                            let ok_wrapped_return = arena.alloc(Loc::at(
-                                loc_expr.region,
-                                Expr::Apply(
-                                    arena.alloc(Loc::at(
-                                        loc_expr.region,
-                                        Expr::Var {
-                                            module_name: ModuleName::TASK,
-                                            ident: "ok",
-                                            suffixed: 0,
-                                        },
-                                    )),
-                                    arena.alloc([sub_new]),
-                                    CalledVia::BangSuffix,
-                                ),
-                            ));
-
                             let unwrapped_expression = apply_task_await(
                                 arena,
                                 sub_arg.region,
                                 sub_arg,
                                 sub_pat,
-                                ok_wrapped_return,
+                                wrap_in_task_ok(arena, sub_new),
                             );
 
                             let mut new_if_thens = Vec::new_in(arena);
@@ -557,28 +541,12 @@ pub fn unwrap_suffixed_expression_if_then_else_help<'a>(
                     sub_pat,
                     sub_new,
                 }) => {
-                    let ok_wrapped_return = arena.alloc(Loc::at(
-                        loc_expr.region,
-                        Expr::Apply(
-                            arena.alloc(Loc::at(
-                                loc_expr.region,
-                                Expr::Var {
-                                    module_name: ModuleName::TASK,
-                                    ident: "ok",
-                                    suffixed: 0,
-                                },
-                            )),
-                            arena.alloc([sub_new]),
-                            CalledVia::BangSuffix,
-                        ),
-                    ));
-
                     let unwrapped_final_else = apply_task_await(
                         arena,
                         sub_arg.region,
                         sub_arg,
                         sub_pat,
-                        ok_wrapped_return,
+                        wrap_in_task_ok(arena, sub_new),
                     );
 
                     let new_if = arena.alloc(Loc::at(
@@ -651,7 +619,8 @@ pub fn unwrap_suffixed_expression_defs_help<'a>(
                                     let next_expr = match unwrap_suffixed_expression(arena,loc_ret,maybe_def_pat) {
                                         Ok(next_expr) => next_expr,
                                         Err(EUnwrapped::UnwrappedSubExpr { sub_arg, sub_pat, sub_new }) => {
-                                            apply_task_await(arena,def_expr.region,sub_arg,sub_pat,sub_new)
+                                            // We need to apply Task.ok here as the defs final expression was unwrapped
+                                            apply_task_await(arena,def_expr.region,sub_arg,sub_pat,wrap_in_task_ok(arena, sub_new))
                                         }
                                         Err(EUnwrapped::UnwrappedDefExpr(..)) | Err(EUnwrapped::Malformed) => {
                                             // TODO handle case when we have maybe_def_pat so can return an unwrapped up
@@ -766,11 +735,13 @@ pub fn apply_task_await<'a>(
 ) -> &'a Loc<Expr<'a>> {
     // If the pattern and the new are the same then we don't need to unwrap anything
     // e.g. `Task.await foo \{} -> Task.ok {}` is the same as `foo`
-    if is_pattern_empty_record(loc_pat) && is_expr_task_ok(loc_new) {
+    if is_matching_empty_record(loc_pat, loc_new) {
         return loc_arg;
     }
 
-    if is_intermediate_answer_with_same_ident(loc_pat, loc_new) {
+    // If the pattern and the new are matching answers then we don't need to unwrap anything
+    // e.g. `Task.await foo \#!a1 -> Task.ok #!a1` is the same as `foo`
+    if is_matching_intermediate_answer(loc_pat, loc_new) {
         return loc_arg;
     }
 
@@ -804,40 +775,39 @@ pub fn apply_task_await<'a>(
     ))
 }
 
-fn is_pattern_empty_record<'a>(loc_pat: &'a Loc<Pattern<'a>>) -> bool {
+fn extract_wrapped_task_ok_value<'a>(loc_expr: &'a Loc<Expr<'a>>) -> Option<&'a Loc<Expr<'a>>> {
-    match loc_pat.value {
+    match loc_expr.value {
+        Expr::Apply(function, arguments, _) => match function.value {
+            Var {
+                module_name, ident, ..
+            } if module_name == ModuleName::TASK && ident == "ok" => arguments.first().copied(),
+            _ => None,
+        },
+        _ => None,
+    }
+}
+
+fn is_matching_empty_record<'a>(
+    loc_pat: &'a Loc<Pattern<'a>>,
+    loc_expr: &'a Loc<Expr<'a>>,
+) -> bool {
+    let is_empty_record = match extract_wrapped_task_ok_value(loc_expr) {
+        Some(task_expr) => match task_expr.value {
+            Expr::Record(collection) => collection.is_empty(),
+            _ => false,
+        },
+        None => false,
+    };
+
+    let is_pattern_empty_record = match loc_pat.value {
         Pattern::RecordDestructure(collection) => collection.is_empty(),
         _ => false,
-    }
+    };
+
+    is_empty_record && is_pattern_empty_record
 }
 
-fn is_expr_task_ok<'a>(loc_expr: &'a Loc<Expr<'a>>) -> bool {
+fn is_matching_intermediate_answer<'a>(
-    match loc_expr.value {
-        Expr::Apply(function, arguments, _) => {
-            let is_task_ok = match function.value {
-                Var {
-                    module_name, ident, ..
-                } => module_name == ModuleName::TASK && ident == "ok",
-                _ => false,
-            };
-
-            let is_arg_empty_record = arguments
-                .first()
-                .map(|arg_loc_expr| match arg_loc_expr.value {
-                    Expr::Record(collection) => collection.is_empty(),
-                    _ => false,
-                })
-                .unwrap_or(false);
-
-            is_task_ok && is_arg_empty_record
-        }
-        _ => false,
-    }
-}
-
-/// Detect when we are applying an identical Pattern and Ident
-/// e.g. `Task.await foo \#a!0 -> #a!0` is the same as `foo`
-fn is_intermediate_answer_with_same_ident<'a>(
     loc_pat: &'a Loc<Pattern<'a>>,
     loc_expr: &'a Loc<Expr<'a>>,
 ) -> bool {
@@ -845,14 +815,17 @@ fn is_intermediate_answer_with_same_ident<'a>(
         Pattern::Identifier { ident, .. } => Some(ident),
         _ => None,
     };
-    let exp_ident = match loc_expr.value {
+    let exp_ident = match extract_wrapped_task_ok_value(loc_expr) {
-        Expr::Var {
+        Some(task_expr) => match task_expr.value {
-            module_name, ident, ..
+            Expr::Var {
-        } if module_name == "" => Some(ident),
+                module_name, ident, ..
-        _ => None,
+            } if module_name.is_empty() && ident.starts_with('#') => Some(ident),
+            _ => None,
+        },
+        None => None,
     };
     match (pat_ident, exp_ident) {
-        (Some(a), Some(b)) => a.starts_with("#") && a == b,
+        (Some(a), Some(b)) => a == b,
         _ => false,
     }
 }

crates/compiler/parse/src/ast.rs

@@ -9,6 +9,7 @@ use bumpalo::Bump;
 use roc_collections::soa::{EitherIndex, Index, Slice};
 use roc_error_macros::internal_error;
 use roc_module::called_via::{BinOp, CalledVia, UnaryOp};
+use roc_module::ident::ModuleName;
 use roc_region::all::{Loc, Position, Region};
 
 #[derive(Debug, Clone, PartialEq, Eq)]
@@ -450,6 +451,24 @@ pub fn is_loc_expr_suffixed(loc_expr: &Loc<Expr>) -> bool {
     }
 }
 
+pub fn wrap_in_task_ok<'a>(arena: &'a Bump, loc_expr: &'a Loc<Expr<'a>>) -> &'a Loc<Expr<'a>> {
+    arena.alloc(Loc::at(
+        loc_expr.region,
+        Expr::Apply(
+            arena.alloc(Loc::at(
+                loc_expr.region,
+                Expr::Var {
+                    module_name: ModuleName::TASK,
+                    ident: "ok",
+                    suffixed: 0,
+                },
+            )),
+            arena.alloc([loc_expr]),
+            CalledVia::BangSuffix,
+        ),
+    ))
+}
+
 pub fn split_around<T>(items: &[T], target: usize) -> (&[T], &[T]) {
     let (before, rest) = items.split_at(target);
     let after = &rest[1..];