unwrap suffixed when condition

crates/compiler/can/src/suffixed.rs

@@ -564,11 +564,33 @@ pub fn unwrap_suffixed_expression_if_then_else_help<'a>(
 }
 
 pub fn unwrap_suffixed_expression_when_help<'a>(
-    _arena: &'a Bump,
+    arena: &'a Bump,
     loc_expr: &'a Loc<Expr<'a>>,
     _maybe_def_pat: Option<&'a Loc<Pattern<'a>>>,
 ) -> Result<&'a Loc<Expr<'a>>, EUnwrapped<'a>> {
-    Ok(loc_expr)
+    match loc_expr.value {
+        Expr::When(condition, branches) => {
+
+            // TODO first unwrap any when branches values
+
+            // then unwrap the when condition
+            match unwrap_suffixed_expression(arena, condition, None) {
+                Ok(unwrapped_condition) => {
+                    let new_when = arena.alloc(Loc::at(loc_expr.region, Expr::When(unwrapped_condition, branches)));
+                    Ok(new_when)
+                }
+                Err(EUnwrapped::UnwrappedSubExpr { sub_arg, sub_pat, sub_new }) => {
+                    let new_when = arena.alloc(Loc::at(loc_expr.region, Expr::When(sub_new, branches)));
+                    let applied_task_await = apply_task_await(arena,loc_expr.region,sub_arg,sub_pat,new_when);
+                    Ok(applied_task_await)
+                }
+                Err(EUnwrapped::UnwrappedDefExpr(..))
+                | Err(EUnwrapped::Malformed) => Err(EUnwrapped::Malformed)
+            }
+
+        }
+        _ => internal_error!("unreachable, expected a When node to be passed into unwrap_suffixed_expression_defs_help"),
+    }
 }
 
 pub fn unwrap_suffixed_expression_defs_help<'a>(

crates/compiler/can/tests/test_suffixed.rs

@@ -665,4 +665,32 @@ mod suffixed_tests {
             r#"Defs { tags: [Index(2147483648)], regions: [@0-103], space_before: [Slice(start = 0, length = 0)], space_after: [Slice(start = 0, length = 0)], spaces: [], type_defs: [], value_defs: [Body(@0-4 Identifier { ident: "copy", suffixed: 0 }, @7-103 Closure([@8-9 Identifier { ident: "a", suffixed: 0 }, @10-11 Identifier { ident: "b", suffixed: 0 }], @36-42 Apply(@36-42 Var { module_name: "Task", ident: "await", suffixed: 0 }, [@36-42 Apply(@36-42 Var { module_name: "", ident: "line", suffixed: 0 }, [@37-42 Str(PlainLine("FOO"))], Space), @36-42 Closure([@36-42 RecordDestructure([])], @60-103 Apply(@60-103 Var { module_name: "", ident: "mapErr", suffixed: 0 }, [@60-72 Apply(@60-67 Var { module_name: "CMD", ident: "new", suffixed: 0 }, [@68-72 Str(PlainLine("cp"))], Space), @100-103 Tag("ERR")], BinOp(Pizza)))], BangSuffix)))] }"#,
         );
     }
+
+    /**
+     * Unwrap a when expression
+     ```roc
+     list =
+        when getList! is
+            [] -> "empty"
+            _ -> "non-empty"
+
+    list =
+        Task.await getList \#!a0 ->
+            when #!a0 is
+                [] -> "empty"
+                _ -> "non-empty"
+     ```
+     */
+    #[test]
+    fn when_simple() {
+        run_test(
+            r#"
+            list = 
+                when getList! is
+                    [] -> "empty"
+                    _ -> "non-empty"
+            "#,
+            r##"Defs { tags: [Index(2147483648)], regions: [@0-111], space_before: [Slice(start = 0, length = 0)], space_after: [Slice(start = 0, length = 0)], spaces: [], type_defs: [], value_defs: [Body(@0-4 Identifier { ident: "list", suffixed: 0 }, @24-111 Apply(@24-111 Var { module_name: "Task", ident: "await", suffixed: 0 }, [@29-37 Var { module_name: "", ident: "getList", suffixed: 0 }, @24-111 Closure([@29-37 Identifier { ident: "#!a0", suffixed: 0 }], @24-111 When(@29-37 Var { module_name: "", ident: "#!a0", suffixed: 0 }, [WhenBranch { patterns: [@61-63 List([])], value: @67-74 Str(PlainLine("empty")), guard: None }, WhenBranch { patterns: [@95-96 Underscore("")], value: @100-111 Str(PlainLine("non-empty")), guard: None }]))], BangSuffix))] }"##,
+        );
+    }
 }