Merge remote-tracking branch 'origin/main' into expect-fx-codegen

crates/compiler/can/src/builtins.rs

@@ -444,7 +444,7 @@ fn no_region<T>(value: T) -> Loc<T> {
 #[inline(always)]
 fn tag(name: &'static str, args: Vec<Expr>, var_store: &mut VarStore) -> Expr {
     Expr::Tag {
-        variant_var: var_store.fresh(),
+        tag_union_var: var_store.fresh(),
         ext_var: var_store.fresh(),
         name: TagName(name.into()),
         arguments: args

crates/compiler/can/src/copy.rs

@@ -535,12 +535,12 @@ fn deep_copy_expr_help<C: CopyEnv>(env: &mut C, copied: &mut Vec<Variable>, expr
         },
 
         Tag {
-            variant_var,
+            tag_union_var: variant_var,
             ext_var,
             name,
             arguments,
         } => Tag {
-            variant_var: sub!(*variant_var),
+            tag_union_var: sub!(*variant_var),
             ext_var: sub!(*ext_var),
             name: name.clone(),
             arguments: arguments
@@ -1164,13 +1164,13 @@ mod test {
         let var2 = new_var(&mut subs, FlexVar(Some(b)));
 
         let expr = Expr::Tag {
-            variant_var: var1,
+            tag_union_var: var1,
             ext_var: Variable::EMPTY_TAG_UNION,
             name: TagName("F".into()),
             arguments: vec![(
                 var2,
                 Loc::at_zero(Expr::Tag {
-                    variant_var: var2,
+                    tag_union_var: var2,
                     ext_var: Variable::EMPTY_TAG_UNION,
                     name: TagName("G".into()),
                     arguments: vec![],
@@ -1185,7 +1185,7 @@ mod test {
 
         match expr {
             Expr::Tag {
-                variant_var,
+                tag_union_var: variant_var,
                 ext_var,
                 name,
                 mut arguments,
@@ -1219,7 +1219,7 @@ mod test {
 
                 match arg.value {
                     Expr::Tag {
-                        variant_var,
+                        tag_union_var: variant_var,
                         ext_var,
                         name,
                         arguments,
@@ -1250,13 +1250,13 @@ mod test {
         let var2 = new_var(&mut source, FlexVar(Some(b)));
 
         let expr = Expr::Tag {
-            variant_var: var1,
+            tag_union_var: var1,
             ext_var: Variable::EMPTY_TAG_UNION,
             name: TagName("F".into()),
             arguments: vec![(
                 var2,
                 Loc::at_zero(Expr::Tag {
-                    variant_var: var2,
+                    tag_union_var: var2,
                     ext_var: Variable::EMPTY_TAG_UNION,
                     name: TagName("G".into()),
                     arguments: vec![],
@@ -1271,7 +1271,7 @@ mod test {
 
         match expr {
             Expr::Tag {
-                variant_var,
+                tag_union_var: variant_var,
                 ext_var,
                 name,
                 mut arguments,
@@ -1300,7 +1300,7 @@ mod test {
 
                 match arg.value {
                     Expr::Tag {
-                        variant_var,
+                        tag_union_var: variant_var,
                         ext_var,
                         name,
                         arguments,

crates/compiler/can/src/def.rs

@@ -232,7 +232,7 @@ impl PendingTypeDef<'_> {
     }
 }
 
-// See github.com/rtfeldman/roc/issues/800 for discussion of the large_enum_variant check.
+// See github.com/roc-lang/roc/issues/800 for discussion of the large_enum_variant check.
 #[derive(Clone, Debug)]
 #[allow(clippy::large_enum_variant)]
 pub enum Declaration {

crates/compiler/can/src/expr.rs

@@ -184,7 +184,7 @@ pub enum Expr {
 
     // Sum Types
     Tag {
-        variant_var: Variable,
+        tag_union_var: Variable,
         ext_var: Variable,
         name: TagName,
         arguments: Vec<(Variable, Loc<Expr>)>,
@@ -780,12 +780,12 @@ pub fn canonicalize_expr<'a>(
                         return (fn_expr, output);
                     }
                     Tag {
-                        variant_var,
+                        tag_union_var: variant_var,
                         ext_var,
                         name,
                         ..
                     } => Tag {
-                        variant_var,
+                        tag_union_var: variant_var,
                         ext_var,
                         name,
                         arguments: args,
@@ -796,7 +796,7 @@ pub fn canonicalize_expr<'a>(
                         name,
                         ..
                     } => Tag {
-                        variant_var,
+                        tag_union_var: variant_var,
                         ext_var,
                         name,
                         arguments: args,
@@ -1422,7 +1422,7 @@ fn canonicalize_when_branch<'a>(
         if output.references.has_value_lookup(symbol) {
             pattern_bound_symbols_body_needs.insert(symbol);
         } else {
-            env.problem(Problem::UnusedDef(symbol, region));
+            env.problem(Problem::UnusedBranchDef(symbol, region));
         }
     }
 
@@ -1893,7 +1893,7 @@ pub fn inline_calls(var_store: &mut VarStore, scope: &mut Scope, expr: Expr) ->
         }
 
         Tag {
-            variant_var,
+            tag_union_var: variant_var,
             ext_var,
             name,
             arguments,

crates/compiler/can/src/lib.rs

@@ -1,5 +1,5 @@
 #![warn(clippy::dbg_macro)]
-// See github.com/rtfeldman/roc/issues/800 for discussion of the large_enum_variant check.
+// See github.com/roc-lang/roc/issues/800 for discussion of the large_enum_variant check.
 #![allow(clippy::large_enum_variant)]
 pub mod abilities;
 pub mod annotation;

crates/compiler/can/src/module.rs

@@ -696,15 +696,26 @@ pub fn canonicalize_module_defs<'a>(
     referenced_values.extend(env.qualified_value_lookups.iter().copied());
     referenced_types.extend(env.qualified_type_lookups.iter().copied());
 
+    let mut fix_closures_no_capture_symbols = VecSet::default();
+    let mut fix_closures_closure_captures = VecMap::default();
     for index in 0..declarations.len() {
         use crate::expr::DeclarationTag::*;
 
+        // For each declaration, we need to fixup the closures inside its def.
+        // Reuse the fixup buffer allocations from the previous iteration.
+        fix_closures_no_capture_symbols.clear();
+        fix_closures_closure_captures.clear();
+
         match declarations.declarations[index] {
             Value => {
                 // def pattern has no default expressions, so skip
                 let loc_expr = &mut declarations.expressions[index];
 
-                fix_values_captured_in_closure_expr(&mut loc_expr.value, &mut VecSet::default());
+                fix_values_captured_in_closure_expr(
+                    &mut loc_expr.value,
+                    &mut fix_closures_no_capture_symbols,
+                    &mut fix_closures_closure_captures,
+                );
             }
             Function(f_index) | Recursive(f_index) | TailRecursive(f_index) => {
                 let name = declarations.symbols[index].value;
@@ -722,30 +733,51 @@ pub fn canonicalize_module_defs<'a>(
                 for (_, _, loc_pat) in function_def.arguments.iter_mut() {
                     fix_values_captured_in_closure_pattern(
                         &mut loc_pat.value,
-                        &mut no_capture_symbols,
+                        &mut fix_closures_no_capture_symbols,
+                        &mut fix_closures_closure_captures,
                     );
                 }
 
-                fix_values_captured_in_closure_expr(&mut loc_expr.value, &mut no_capture_symbols);
+                fix_values_captured_in_closure_expr(
+                    &mut loc_expr.value,
+                    &mut fix_closures_no_capture_symbols,
+                    &mut fix_closures_closure_captures,
+                );
             }
             Destructure(d_index) => {
                 let destruct_def = &mut declarations.destructs[d_index.index()];
                 let loc_pat = &mut destruct_def.loc_pattern;
                 let loc_expr = &mut declarations.expressions[index];
 
-                fix_values_captured_in_closure_pattern(&mut loc_pat.value, &mut VecSet::default());
-                fix_values_captured_in_closure_expr(&mut loc_expr.value, &mut VecSet::default());
+                fix_values_captured_in_closure_pattern(
+                    &mut loc_pat.value,
+                    &mut fix_closures_no_capture_symbols,
+                    &mut fix_closures_closure_captures,
+                );
+                fix_values_captured_in_closure_expr(
+                    &mut loc_expr.value,
+                    &mut fix_closures_no_capture_symbols,
+                    &mut fix_closures_closure_captures,
+                );
             }
             MutualRecursion { .. } => {
                 // the declarations of this group will be treaded individually by later iterations
             }
             Expectation => {
                 let loc_expr = &mut declarations.expressions[index];
-                fix_values_captured_in_closure_expr(&mut loc_expr.value, &mut VecSet::default());
+                fix_values_captured_in_closure_expr(
+                    &mut loc_expr.value,
+                    &mut fix_closures_no_capture_symbols,
+                    &mut fix_closures_closure_captures,
+                );
             }
             ExpectationFx => {
                 let loc_expr = &mut declarations.expressions[index];
-                fix_values_captured_in_closure_expr(&mut loc_expr.value, &mut VecSet::default());
+                fix_values_captured_in_closure_expr(
+                    &mut loc_expr.value,
+                    &mut fix_closures_no_capture_symbols,
+                    &mut fix_closures_closure_captures,
+                );
             }
         }
     }
@@ -771,16 +803,26 @@ pub fn canonicalize_module_defs<'a>(
 fn fix_values_captured_in_closure_def(
     def: &mut crate::def::Def,
     no_capture_symbols: &mut VecSet<Symbol>,
+    closure_captures: &mut VecMap<Symbol, Vec<(Symbol, Variable)>>,
 ) {
     // patterns can contain default expressions, so much go over them too!
-    fix_values_captured_in_closure_pattern(&mut def.loc_pattern.value, no_capture_symbols);
+    fix_values_captured_in_closure_pattern(
+        &mut def.loc_pattern.value,
+        no_capture_symbols,
+        closure_captures,
+    );
 
-    fix_values_captured_in_closure_expr(&mut def.loc_expr.value, no_capture_symbols);
+    fix_values_captured_in_closure_expr(
+        &mut def.loc_expr.value,
+        no_capture_symbols,
+        closure_captures,
+    );
 }
 
 fn fix_values_captured_in_closure_defs(
     defs: &mut [crate::def::Def],
     no_capture_symbols: &mut VecSet<Symbol>,
+    closure_captures: &mut VecMap<Symbol, Vec<(Symbol, Variable)>>,
 ) {
     // recursive defs cannot capture each other
     for def in defs.iter() {
@@ -789,16 +831,38 @@ fn fix_values_captured_in_closure_defs(
         );
     }
 
-    // TODO mutually recursive functions should both capture the union of both their capture sets
-
     for def in defs.iter_mut() {
-        fix_values_captured_in_closure_def(def, no_capture_symbols);
+        fix_values_captured_in_closure_def(def, no_capture_symbols, closure_captures);
+    }
+
+    // Mutually recursive functions should both capture the union of all their capture sets
+    //
+    // Really unfortunate we make a lot of clones here, can this be done more efficiently?
+    let mut total_capture_set = Vec::default();
+    for def in defs.iter_mut() {
+        if let Expr::Closure(ClosureData {
+            captured_symbols, ..
+        }) = &def.loc_expr.value
+        {
+            total_capture_set.extend(captured_symbols.iter().copied());
+        }
+    }
+    total_capture_set.sort_by_key(|(sym, _)| *sym);
+    total_capture_set.dedup_by_key(|(sym, _)| *sym);
+    for def in defs.iter_mut() {
+        if let Expr::Closure(ClosureData {
+            captured_symbols, ..
+        }) = &mut def.loc_expr.value
+        {
+            *captured_symbols = total_capture_set.clone();
+        }
     }
 }
 
 fn fix_values_captured_in_closure_pattern(
     pattern: &mut crate::pattern::Pattern,
     no_capture_symbols: &mut VecSet<Symbol>,
+    closure_captures: &mut VecMap<Symbol, Vec<(Symbol, Variable)>>,
 ) {
     use crate::pattern::Pattern::*;
 
@@ -808,24 +872,35 @@ fn fix_values_captured_in_closure_pattern(
             ..
         } => {
             for (_, loc_arg) in loc_args.iter_mut() {
-                fix_values_captured_in_closure_pattern(&mut loc_arg.value, no_capture_symbols);
+                fix_values_captured_in_closure_pattern(
+                    &mut loc_arg.value,
+                    no_capture_symbols,
+                    closure_captures,
+                );
             }
         }
         UnwrappedOpaque { argument, .. } => {
             let (_, loc_arg) = &mut **argument;
-            fix_values_captured_in_closure_pattern(&mut loc_arg.value, no_capture_symbols);
+            fix_values_captured_in_closure_pattern(
+                &mut loc_arg.value,
+                no_capture_symbols,
+                closure_captures,
+            );
         }
         RecordDestructure { destructs, .. } => {
             for loc_destruct in destructs.iter_mut() {
                 use crate::pattern::DestructType::*;
                 match &mut loc_destruct.value.typ {
                     Required => {}
-                    Optional(_, loc_expr) => {
-                        fix_values_captured_in_closure_expr(&mut loc_expr.value, no_capture_symbols)
-                    }
+                    Optional(_, loc_expr) => fix_values_captured_in_closure_expr(
+                        &mut loc_expr.value,
+                        no_capture_symbols,
+                        closure_captures,
+                    ),
                     Guard(_, loc_pattern) => fix_values_captured_in_closure_pattern(
                         &mut loc_pattern.value,
                         no_capture_symbols,
+                        closure_captures,
                     ),
                 }
             }
@@ -848,19 +923,28 @@ fn fix_values_captured_in_closure_pattern(
 fn fix_values_captured_in_closure_expr(
     expr: &mut crate::expr::Expr,
     no_capture_symbols: &mut VecSet<Symbol>,
+    closure_captures: &mut VecMap<Symbol, Vec<(Symbol, Variable)>>,
 ) {
     use crate::expr::Expr::*;
 
     match expr {
         LetNonRec(def, loc_expr) => {
             // LetNonRec(Box<Def>, Box<Located<Expr>>, Variable, Aliases),
-            fix_values_captured_in_closure_def(def, no_capture_symbols);
-            fix_values_captured_in_closure_expr(&mut loc_expr.value, no_capture_symbols);
+            fix_values_captured_in_closure_def(def, no_capture_symbols, closure_captures);
+            fix_values_captured_in_closure_expr(
+                &mut loc_expr.value,
+                no_capture_symbols,
+                closure_captures,
+            );
         }
         LetRec(defs, loc_expr, _) => {
             // LetRec(Vec<Def>, Box<Located<Expr>>, Variable, Aliases),
-            fix_values_captured_in_closure_defs(defs, no_capture_symbols);
-            fix_values_captured_in_closure_expr(&mut loc_expr.value, no_capture_symbols);
+            fix_values_captured_in_closure_defs(defs, no_capture_symbols, closure_captures);
+            fix_values_captured_in_closure_expr(
+                &mut loc_expr.value,
+                no_capture_symbols,
+                closure_captures,
+            );
         }
 
         Expect {
@@ -868,8 +952,16 @@ fn fix_values_captured_in_closure_expr(
             loc_continuation,
             lookups_in_cond: _,
         } => {
-            fix_values_captured_in_closure_expr(&mut loc_condition.value, no_capture_symbols);
-            fix_values_captured_in_closure_expr(&mut loc_continuation.value, no_capture_symbols);
+            fix_values_captured_in_closure_expr(
+                &mut loc_condition.value,
+                no_capture_symbols,
+                closure_captures,
+            );
+            fix_values_captured_in_closure_expr(
+                &mut loc_continuation.value,
+                no_capture_symbols,
+                closure_captures,
+            );
         }
 
         ExpectFx {
@@ -891,16 +983,58 @@ fn fix_values_captured_in_closure_expr(
             captured_symbols.retain(|(s, _)| !no_capture_symbols.contains(s));
             captured_symbols.retain(|(s, _)| s != name);
 
+            let original_captures_len = captured_symbols.len();
+            let mut num_visited = 0;
+            let mut i = 0;
+            while num_visited < original_captures_len {
+                // If we've captured a capturing closure, replace the captured closure symbol with
+                // the symbols of its captures. That way, we can construct the closure with the
+                // captures it needs inside our body.
+                //
+                // E.g.
+                //   x = ""
+                //   inner = \{} -> x
+                //   outer = \{} -> inner {}
+                //
+                // initially `outer` captures [inner], but this is then replaced with just [x].
+                let (captured_symbol, _) = captured_symbols[i];
+                if let Some(captures) = closure_captures.get(&captured_symbol) {
+                    debug_assert!(!captures.is_empty());
+                    captured_symbols.swap_remove(i);
+                    captured_symbols.extend(captures);
+                    // Jump two, because the next element is now one of the newly-added captures,
+                    // which we don't need to check.
+                    i += 2;
+                } else {
+                    i += 1;
+                }
+                num_visited += 1;
+            }
+            if captured_symbols.len() > original_captures_len {
+                // Re-sort, since we've added new captures.
+                captured_symbols.sort_by_key(|(sym, _)| *sym);
+            }
+
             if captured_symbols.is_empty() {
                 no_capture_symbols.insert(*name);
+            } else {
+                closure_captures.insert(*name, captured_symbols.to_vec());
             }
 
             // patterns can contain default expressions, so much go over them too!
             for (_, _, loc_pat) in arguments.iter_mut() {
-                fix_values_captured_in_closure_pattern(&mut loc_pat.value, no_capture_symbols);
+                fix_values_captured_in_closure_pattern(
+                    &mut loc_pat.value,
+                    no_capture_symbols,
+                    closure_captures,
+                );
             }
 
-            fix_values_captured_in_closure_expr(&mut loc_body.value, no_capture_symbols);
+            fix_values_captured_in_closure_expr(
+                &mut loc_body.value,
+                no_capture_symbols,
+                closure_captures,
+            );
         }
 
         Num(..)
@@ -918,28 +1052,45 @@ fn fix_values_captured_in_closure_expr(
 
         List { loc_elems, .. } => {
             for elem in loc_elems.iter_mut() {
-                fix_values_captured_in_closure_expr(&mut elem.value, no_capture_symbols);
+                fix_values_captured_in_closure_expr(
+                    &mut elem.value,
+                    no_capture_symbols,
+                    closure_captures,
+                );
             }
         }
 
         When {
             loc_cond, branches, ..
         } => {
-            fix_values_captured_in_closure_expr(&mut loc_cond.value, no_capture_symbols);
+            fix_values_captured_in_closure_expr(
+                &mut loc_cond.value,
+                no_capture_symbols,
+                closure_captures,
+            );
 
             for branch in branches.iter_mut() {
-                fix_values_captured_in_closure_expr(&mut branch.value.value, no_capture_symbols);
+                fix_values_captured_in_closure_expr(
+                    &mut branch.value.value,
+                    no_capture_symbols,
+                    closure_captures,
+                );
 
                 // patterns can contain default expressions, so much go over them too!
                 for loc_pat in branch.patterns.iter_mut() {
                     fix_values_captured_in_closure_pattern(
                         &mut loc_pat.pattern.value,
                         no_capture_symbols,
+                        closure_captures,
                     );
                 }
 
                 if let Some(guard) = &mut branch.guard {
-                    fix_values_captured_in_closure_expr(&mut guard.value, no_capture_symbols);
+                    fix_values_captured_in_closure_expr(
+                        &mut guard.value,
+                        no_capture_symbols,
+                        closure_captures,
+                    );
                 }
             }
         }
@@ -950,23 +1101,43 @@ fn fix_values_captured_in_closure_expr(
             ..
         } => {
             for (loc_cond, loc_then) in branches.iter_mut() {
-                fix_values_captured_in_closure_expr(&mut loc_cond.value, no_capture_symbols);
-                fix_values_captured_in_closure_expr(&mut loc_then.value, no_capture_symbols);
+                fix_values_captured_in_closure_expr(
+                    &mut loc_cond.value,
+                    no_capture_symbols,
+                    closure_captures,
+                );
+                fix_values_captured_in_closure_expr(
+                    &mut loc_then.value,
+                    no_capture_symbols,
+                    closure_captures,
+                );
             }
 
-            fix_values_captured_in_closure_expr(&mut final_else.value, no_capture_symbols);
+            fix_values_captured_in_closure_expr(
+                &mut final_else.value,
+                no_capture_symbols,
+                closure_captures,
+            );
         }
 
         Call(function, arguments, _) => {
-            fix_values_captured_in_closure_expr(&mut function.1.value, no_capture_symbols);
+            fix_values_captured_in_closure_expr(
+                &mut function.1.value,
+                no_capture_symbols,
+                closure_captures,
+            );
 
             for (_, loc_arg) in arguments.iter_mut() {
-                fix_values_captured_in_closure_expr(&mut loc_arg.value, no_capture_symbols);
+                fix_values_captured_in_closure_expr(
+                    &mut loc_arg.value,
+                    no_capture_symbols,
+                    closure_captures,
+                );
             }
         }
         RunLowLevel { args, .. } | ForeignCall { args, .. } => {
             for (_, arg) in args.iter_mut() {
-                fix_values_captured_in_closure_expr(arg, no_capture_symbols);
+                fix_values_captured_in_closure_expr(arg, no_capture_symbols, closure_captures);
             }
         }
 
@@ -975,22 +1146,38 @@ fn fix_values_captured_in_closure_expr(
             updates: fields, ..
         } => {
             for (_, field) in fields.iter_mut() {
-                fix_values_captured_in_closure_expr(&mut field.loc_expr.value, no_capture_symbols);
+                fix_values_captured_in_closure_expr(
+                    &mut field.loc_expr.value,
+                    no_capture_symbols,
+                    closure_captures,
+                );
             }
         }
 
         Access { loc_expr, .. } => {
-            fix_values_captured_in_closure_expr(&mut loc_expr.value, no_capture_symbols);
+            fix_values_captured_in_closure_expr(
+                &mut loc_expr.value,
+                no_capture_symbols,
+                closure_captures,
+            );
         }
 
         Tag { arguments, .. } => {
             for (_, loc_arg) in arguments.iter_mut() {
-                fix_values_captured_in_closure_expr(&mut loc_arg.value, no_capture_symbols);
+                fix_values_captured_in_closure_expr(
+                    &mut loc_arg.value,
+                    no_capture_symbols,
+                    closure_captures,
+                );
             }
         }
         OpaqueRef { argument, .. } => {
             let (_, loc_arg) = &mut **argument;
-            fix_values_captured_in_closure_expr(&mut loc_arg.value, no_capture_symbols);
+            fix_values_captured_in_closure_expr(
+                &mut loc_arg.value,
+                no_capture_symbols,
+                closure_captures,
+            );
         }
         OpaqueWrapFunction(_) => {}
     }

crates/compiler/can/src/traverse.rs

@@ -244,7 +244,7 @@ pub fn walk_expr<V: Visitor>(visitor: &mut V, expr: &Expr, var: Variable) {
             walk_record_fields(visitor, updates.iter());
         }
         Expr::Tag {
-            variant_var: _,
+            tag_union_var: _,
             ext_var: _,
             name: _,
             arguments,