fully canonicalize expect-fx

crates/compiler/can/src/def.rs

@@ -87,6 +87,7 @@ pub struct Annotation {
 pub(crate) struct CanDefs {
     defs: Vec<Option<Def>>,
     expects: Expects,
+    expects_fx: Expects,
     def_ordering: DefOrdering,
     aliases: VecMap<Symbol, Alias>,
 }
@@ -239,6 +240,7 @@ pub enum Declaration {
     DeclareRec(Vec<Def>, IllegalCycleMark),
     Builtin(Def),
     Expects(Expects),
+    ExpectsFx(Expects),
     /// If we know a cycle is illegal during canonicalization.
     /// Otherwise we will try to detect this during solving; see [`IllegalCycleMark`].
     InvalidCycle(Vec<CycleEntry>),
@@ -253,6 +255,7 @@ impl Declaration {
             InvalidCycle { .. } => 0,
             Builtin(_) => 0,
             Expects(_) => 0,
+            ExpectsFx(_) => 0,
         }
     }
 
@@ -268,7 +271,7 @@ impl Declaration {
                 &cycles.first().unwrap().expr_region,
                 &cycles.last().unwrap().expr_region,
             ),
-            Declaration::Expects(expects) => Region::span_across(
+            Declaration::Expects(expects) | Declaration::ExpectsFx(expects) => Region::span_across(
                 expects.regions.first().unwrap(),
                 expects.regions.last().unwrap(),
             ),
@@ -990,7 +993,7 @@ fn canonicalize_value_defs<'a>(
     }
 
     let mut expects = Expects::with_capacity(pending_expects.len());
-    let mut expect_fx = Expects::with_capacity(pending_expects.len());
+    let mut expects_fx = Expects::with_capacity(pending_expects.len());
 
     for pending in pending_expects {
         let (loc_can_condition, can_output) = canonicalize_expr(
@@ -1006,9 +1009,24 @@ fn canonicalize_value_defs<'a>(
         output.union(can_output);
     }
 
+    for pending in pending_expect_fx {
+        let (loc_can_condition, can_output) = canonicalize_expr(
+            env,
+            var_store,
+            scope,
+            pending.condition.region,
+            &pending.condition.value,
+        );
+
+        expects_fx.push(loc_can_condition, pending.preceding_comment);
+
+        output.union(can_output);
+    }
+
     let can_defs = CanDefs {
         defs,
         expects,
+        expects_fx,
         def_ordering,
         aliases,
     };
@@ -1392,6 +1410,7 @@ pub(crate) fn sort_can_defs_new(
     let CanDefs {
         defs,
         expects,
+        expects_fx,
         def_ordering,
         aliases,
     } = defs;
@@ -1422,6 +1441,19 @@ pub(crate) fn sort_can_defs_new(
         declarations.push_expect(preceding_comment, name, Loc::at(region, condition));
     }
 
+    let it = expects_fx
+        .conditions
+        .into_iter()
+        .zip(expects_fx.regions)
+        .zip(expects_fx.preceding_comment);
+
+    for ((condition, region), preceding_comment) in it {
+        // an `expect` does not have a user-defined name, but we'll need a name to call the expectation
+        let name = scope.gen_unique_symbol();
+
+        declarations.push_expect_fx(preceding_comment, name, Loc::at(region, condition));
+    }
+
     for (symbol, alias) in aliases.into_iter() {
         output.aliases.insert(symbol, alias);
     }
@@ -1599,6 +1631,7 @@ pub(crate) fn sort_can_defs(
     let CanDefs {
         mut defs,
         expects,
+        expects_fx,
         def_ordering,
         aliases,
     } = defs;
@@ -1713,6 +1746,10 @@ pub(crate) fn sort_can_defs(
         declarations.push(Declaration::Expects(expects));
     }
 
+    if !expects_fx.conditions.is_empty() {
+        declarations.push(Declaration::ExpectsFx(expects_fx));
+    }
+
     (declarations, output)
 }
 
@@ -2205,6 +2242,10 @@ fn decl_to_let(decl: Declaration, loc_ret: Loc<Expr>) -> Loc<Expr> {
             // Expects should only be added to top-level decls, not to let-exprs!
             unreachable!("{:?}", &expects)
         }
+        Declaration::ExpectsFx(expects) => {
+            // Expects should only be added to top-level decls, not to let-exprs!
+            unreachable!("{:?}", &expects)
+        }
     }
 }
 
@@ -2518,7 +2559,7 @@ fn to_pending_value_def<'a>(
         ExpectFx {
             condition,
             preceding_comment,
-        } => PendingValue::Expect(PendingExpect {
+        } => PendingValue::ExpectFx(PendingExpect {
             condition,
             preceding_comment: *preceding_comment,
         }),

crates/compiler/can/src/expr.rs

@@ -2311,6 +2311,24 @@ impl Declarations {
         index
     }
 
+    pub fn push_expect_fx(
+        &mut self,
+        preceding_comment: Region,
+        name: Symbol,
+        loc_expr: Loc<Expr>,
+    ) -> usize {
+        let index = self.declarations.len();
+
+        self.declarations.push(DeclarationTag::Expectation);
+        self.variables.push(Variable::BOOL);
+        self.symbols.push(Loc::at(preceding_comment, name));
+        self.annotations.push(None);
+
+        self.expressions.push(loc_expr);
+
+        index
+    }
+
     pub fn push_value_def(
         &mut self,
         symbol: Loc<Symbol>,
@@ -2491,6 +2509,12 @@ impl Declarations {
 
                     collector.visit_expr(&loc_expr.value, loc_expr.region, var);
                 }
+                ExpectationFx => {
+                    let loc_expr =
+                        toplevel_expect_to_inline_expect(self.expressions[index].clone());
+
+                    collector.visit_expr(&loc_expr.value, loc_expr.region, var);
+                }
             }
         }
 
@@ -2504,6 +2528,7 @@ roc_error_macros::assert_sizeof_default!(DeclarationTag, 8);
 pub enum DeclarationTag {
     Value,
     Expectation,
+    ExpectationFx,
     Function(Index<Loc<FunctionDef>>),
     Recursive(Index<Loc<FunctionDef>>),
     TailRecursive(Index<Loc<FunctionDef>>),
@@ -2516,14 +2541,14 @@ pub enum DeclarationTag {
 
 impl DeclarationTag {
     fn len(self) -> usize {
+        use DeclarationTag::*;
+
         match self {
-            DeclarationTag::Function(_) => 1,
-            DeclarationTag::Recursive(_) => 1,
-            DeclarationTag::TailRecursive(_) => 1,
-            DeclarationTag::Value => 1,
-            DeclarationTag::Expectation => 1,
-            DeclarationTag::Destructure(_) => 1,
-            DeclarationTag::MutualRecursion { length, .. } => length as usize + 1,
+            Function(_) | Recursive(_) | TailRecursive(_) => 1,
+            Value => 1,
+            Expectation | ExpectationFx => 1,
+            Destructure(_) => 1,
+            MutualRecursion { length, .. } => length as usize + 1,
         }
     }
 }

crates/compiler/can/src/module.rs

@@ -620,6 +620,7 @@ pub fn canonicalize_module_defs<'a>(
                 // the declarations of this group will be treaded individually by later iterations
             }
             Expectation => { /* ignore */ }
+            ExpectationFx => { /* ignore */ }
         }
     }
 
@@ -742,6 +743,10 @@ pub fn canonicalize_module_defs<'a>(
                 let loc_expr = &mut declarations.expressions[index];
                 fix_values_captured_in_closure_expr(&mut loc_expr.value, &mut VecSet::default());
             }
+            ExpectationFx => {
+                let loc_expr = &mut declarations.expressions[index];
+                fix_values_captured_in_closure_expr(&mut loc_expr.value, &mut VecSet::default());
+            }
         }
     }
 

crates/compiler/can/src/traverse.rs

@@ -52,6 +52,11 @@ pub fn walk_decls<V: Visitor>(visitor: &mut V, decls: &Declarations) {
 
                 visitor.visit_expr(&loc_condition.value, loc_condition.region, Variable::BOOL);
             }
+            ExpectationFx => {
+                let loc_condition = &decls.expressions[index];
+
+                visitor.visit_expr(&loc_condition.value, loc_condition.region, Variable::BOOL);
+            }
             Function(function_index)
             | Recursive(function_index)
             | TailRecursive(function_index) => {
@@ -119,6 +124,12 @@ fn walk_decl<V: Visitor>(visitor: &mut V, decl: &Declaration) {
                 visitor.visit_expr(condition, *region, Variable::BOOL);
             }
         }
+        Declaration::ExpectsFx(expects) => {
+            let it = expects.regions.iter().zip(expects.conditions.iter());
+            for (region, condition) in it {
+                visitor.visit_expr(condition, *region, Variable::BOOL);
+            }
+        }
         Declaration::Builtin(def) => visitor.visit_def(def),
         Declaration::InvalidCycle(_cycles) => {
             // ignore