Merge pull request #3733 from rtfeldman/expect-fx

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>,
 }
@@ -106,6 +107,12 @@ impl Expects {
             preceding_comment: Vec::with_capacity(capacity),
         }
     }
+
+    fn push(&mut self, loc_can_condition: Loc<Expr>, preceding_comment: Region) {
+        self.conditions.push(loc_can_condition.value);
+        self.regions.push(loc_can_condition.region);
+        self.preceding_comment.push(preceding_comment);
+    }
 }
 
 /// A Def that has had patterns and type annnotations canonicalized,
@@ -233,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>),
@@ -247,6 +255,7 @@ impl Declaration {
             InvalidCycle { .. } => 0,
             Builtin(_) => 0,
             Expects(_) => 0,
+            ExpectsFx(_) => 0,
         }
     }
 
@@ -262,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(),
             ),
@@ -907,6 +916,7 @@ fn canonicalize_value_defs<'a>(
     // once we've finished assembling the entire scope.
     let mut pending_value_defs = 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());
 
     for loc_pending_def in value_defs {
         match loc_pending_def.value {
@@ -921,6 +931,10 @@ fn canonicalize_value_defs<'a>(
             PendingValue::Expect(pending_expect) => {
                 pending_expects.push(pending_expect);
             }
+
+            PendingValue::ExpectFx(pending_expect) => {
+                pending_expect_fx.push(pending_expect);
+            }
         }
     }
 
@@ -979,6 +993,7 @@ fn canonicalize_value_defs<'a>(
     }
 
     let mut expects = 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(
@@ -989,9 +1004,21 @@ fn canonicalize_value_defs<'a>(
             &pending.condition.value,
         );
 
-        expects.conditions.push(loc_can_condition.value);
-        expects.regions.push(loc_can_condition.region);
-        expects.preceding_comment.push(pending.preceding_comment);
+        expects.push(loc_can_condition, pending.preceding_comment);
+
+        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);
     }
@@ -999,6 +1026,7 @@ fn canonicalize_value_defs<'a>(
     let can_defs = CanDefs {
         defs,
         expects,
+        expects_fx,
         def_ordering,
         aliases,
     };
@@ -1382,6 +1410,7 @@ pub(crate) fn sort_can_defs_new(
     let CanDefs {
         defs,
         expects,
+        expects_fx,
         def_ordering,
         aliases,
     } = defs;
@@ -1412,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);
     }
@@ -1589,6 +1631,7 @@ pub(crate) fn sort_can_defs(
     let CanDefs {
         mut defs,
         expects,
+        expects_fx,
         def_ordering,
         aliases,
     } = defs;
@@ -1703,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)
 }
 
@@ -2195,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)
+        }
     }
 }
 
@@ -2391,6 +2442,7 @@ fn to_pending_type_def<'a>(
 enum PendingValue<'a> {
     Def(PendingValueDef<'a>),
     Expect(PendingExpect<'a>),
+    ExpectFx(PendingExpect<'a>),
     SignatureDefMismatch,
 }
 
@@ -2503,6 +2555,14 @@ fn to_pending_value_def<'a>(
             condition,
             preceding_comment: *preceding_comment,
         }),
+
+        ExpectFx {
+            condition,
+            preceding_comment,
+        } => 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/operator.rs

@@ -92,6 +92,16 @@ fn desugar_value_def<'a>(arena: &'a Bump, def: &'a ValueDef<'a>) -> ValueDef<'a>
                 preceding_comment: *preceding_comment,
             }
         }
+        ExpectFx {
+            condition,
+            preceding_comment,
+        } => {
+            let desugared_condition = &*arena.alloc(desugar_expr(arena, condition));
+            ExpectFx {
+                condition: desugared_condition,
+                preceding_comment: *preceding_comment,
+            }
+        }
     }
 }
 

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

crates/compiler/constrain/src/expr.rs

@@ -1997,6 +1997,23 @@ pub fn constrain_decls(
 
                 constraint = constraints.let_constraint([], [], [], expect_constraint, constraint)
             }
+            ExpectationFx => {
+                let loc_expr = &declarations.expressions[index];
+
+                let bool_type = Type::Variable(Variable::BOOL);
+                let expected =
+                    Expected::ForReason(Reason::ExpectCondition, bool_type, loc_expr.region);
+
+                let expect_constraint = constrain_expr(
+                    constraints,
+                    &mut env,
+                    loc_expr.region,
+                    &loc_expr.value,
+                    expected,
+                );
+
+                constraint = constraints.let_constraint([], [], [], expect_constraint, constraint)
+            }
             Function(function_def_index) => {
                 constraint = constrain_function_def(
                     constraints,

crates/compiler/fmt/src/def.rs

@@ -159,6 +159,7 @@ impl<'a> Formattable for ValueDef<'a> {
             Body(loc_pattern, loc_expr) => loc_pattern.is_multiline() || loc_expr.is_multiline(),
             AnnotatedBody { .. } => true,
             Expect { condition, .. } => condition.is_multiline(),
+            ExpectFx { condition, .. } => condition.is_multiline(),
         }
     }
 
@@ -233,6 +234,9 @@ impl<'a> Formattable for ValueDef<'a> {
                 fmt_body(buf, &loc_pattern.value, &loc_expr.value, indent);
             }
             Expect { condition, .. } => fmt_expect(buf, condition, self.is_multiline(), indent),
+            ExpectFx { condition, .. } => {
+                fmt_expect_fx(buf, condition, self.is_multiline(), indent)
+            }
             AnnotatedBody {
                 ann_pattern,
                 ann_type,
@@ -303,6 +307,27 @@ fn fmt_expect<'a, 'buf>(
     condition.format(buf, return_indent);
 }
 
+fn fmt_expect_fx<'a, 'buf>(
+    buf: &mut Buf<'buf>,
+    condition: &'a Loc<Expr<'a>>,
+    is_multiline: bool,
+    indent: u16,
+) {
+    buf.ensure_ends_with_newline();
+    buf.indent(indent);
+    buf.push_str("expect-fx");
+
+    let return_indent = if is_multiline {
+        buf.newline();
+        indent + INDENT
+    } else {
+        buf.spaces(1);
+        indent
+    };
+
+    condition.format(buf, return_indent);
+}
+
 pub fn fmt_value_def<'a, 'buf>(
     buf: &mut Buf<'buf>,
     def: &roc_parse::ast::ValueDef<'a>,

crates/compiler/fmt/src/spaces.rs

@@ -547,6 +547,13 @@ impl<'a> RemoveSpaces<'a> for ValueDef<'a> {
                 condition: arena.alloc(condition.remove_spaces(arena)),
                 preceding_comment,
             },
+            ExpectFx {
+                condition,
+                preceding_comment,
+            } => ExpectFx {
+                condition: arena.alloc(condition.remove_spaces(arena)),
+                preceding_comment,
+            },
         }
     }
 }

crates/compiler/load_internal/src/docs.rs

@@ -203,6 +203,10 @@ fn generate_entry_docs<'a>(
                 ValueDef::Expect { .. } => {
                     // Don't generate docs for `expect`s
                 }
+
+                ValueDef::ExpectFx { .. } => {
+                    // Don't generate docs for `expect-fx`s
+                }
             },
             Ok(type_index) => match &defs.type_defs[type_index.index()] {
                 TypeDef::Alias {

crates/compiler/load_internal/src/file.rs

@@ -5067,6 +5067,80 @@ fn build_pending_specializations<'a>(
                 let expr_region = declarations.expressions[index].region;
                 let region = Region::span_across(&name_region, &expr_region);
 
+                toplevel_expects.insert(symbol, region);
+                procs_base.partial_procs.insert(symbol, proc);
+            }
+            ExpectationFx => {
+                // skip expectations if we're not going to run them
+                match execution_mode {
+                    ExecutionMode::Test => { /* fall through */ }
+                    ExecutionMode::Check | ExecutionMode::Executable => continue,
+                }
+
+                // mark this symbol as a top-level thunk before any other work on the procs
+                module_thunks.push(symbol);
+
+                let expr_var = Variable::EMPTY_RECORD;
+
+                let is_host_exposed = true;
+
+                // If this is an exposed symbol, we need to
+                // register it as such. Otherwise, since it
+                // never gets called by Roc code, it will never
+                // get specialized!
+                if is_host_exposed {
+                    let layout_result =
+                        layout_cache.raw_from_var(mono_env.arena, expr_var, mono_env.subs);
+
+                    // cannot specialize when e.g. main's type contains type variables
+                    if let Err(e) = layout_result {
+                        match e {
+                            LayoutProblem::Erroneous => {
+                                let message = "top level function has erroneous type";
+                                procs_base.runtime_errors.insert(symbol, message);
+                                continue;
+                            }
+                            LayoutProblem::UnresolvedTypeVar(v) => {
+                                let message = format!(
+                                    "top level function has unresolved type variable {:?}",
+                                    v
+                                );
+                                procs_base
+                                    .runtime_errors
+                                    .insert(symbol, mono_env.arena.alloc(message));
+                                continue;
+                            }
+                        }
+                    }
+
+                    procs_base.host_specializations.insert_host_exposed(
+                        mono_env.subs,
+                        LambdaName::no_niche(symbol),
+                        annotation,
+                        expr_var,
+                    );
+                }
+
+                let body = roc_can::expr::toplevel_expect_to_inline_expect(body);
+
+                let proc = PartialProc {
+                    annotation: expr_var,
+                    // This is a 0-arity thunk, so it has no arguments.
+                    pattern_symbols: &[],
+                    // This is a top-level definition, so it cannot capture anything
+                    captured_symbols: CapturedSymbols::None,
+                    body: body.value,
+                    body_var: expr_var,
+                    // This is a 0-arity thunk, so it cannot be recursive
+                    is_self_recursive: false,
+                };
+
+                // extend the region of the expect expression with the region of the preceding
+                // comment, so it is shown in failure/panic messages
+                let name_region = declarations.symbols[index].region;
+                let expr_region = declarations.expressions[index].region;
+                let region = Region::span_across(&name_region, &expr_region);
+
                 toplevel_expects.insert(symbol, region);
                 procs_base.partial_procs.insert(symbol, proc);
             }

crates/compiler/load_internal/tests/test_load.rs

@@ -278,6 +278,10 @@ fn expect_types(mut loaded_module: LoadedModule, mut expected_types: HashMap<&st
                 // at least at the moment this does not happen
                 panic!("Unexpected expectation in module declarations");
             }
+            ExpectationFx => {
+                // at least at the moment this does not happen
+                panic!("Unexpected expectation in module declarations");
+            }
         };
     }
 
@@ -363,7 +367,7 @@ fn interface_with_deps() {
                 def_count += 1;
             }
             MutualRecursion { .. } => { /* do nothing, not a def */ }
-            Expectation => { /* do nothing, not a def */ }
+            Expectation | ExpectationFx => { /* do nothing, not a def */ }
         }
     }
 

crates/compiler/parse/src/ast.rs

@@ -335,6 +335,11 @@ pub enum ValueDef<'a> {
         condition: &'a Loc<Expr<'a>>,
         preceding_comment: Region,
     },
+
+    ExpectFx {
+        condition: &'a Loc<Expr<'a>>,
+        preceding_comment: Region,
+    },
 }
 
 #[derive(Debug, Clone, PartialEq, Default)]

crates/compiler/parse/src/expr.rs

@@ -633,6 +633,11 @@ fn parse_defs_end<'a>(
 
         let start = state.pos();
 
+        let parse_expect_vanilla =
+            crate::parser::keyword_e(crate::keyword::EXPECT, EExpect::Expect);
+        let parse_expect_fx = crate::parser::keyword_e(crate::keyword::EXPECT_FX, EExpect::Expect);
+        let parse_expect = either!(parse_expect_vanilla, parse_expect_fx);
+
         match space0_after_e(
             crate::pattern::loc_pattern_help(min_indent),
             min_indent,
@@ -641,14 +646,13 @@ fn parse_defs_end<'a>(
         .parse(arena, state.clone())
         {
             Err((NoProgress, _, _)) => {
-                match crate::parser::keyword_e(crate::keyword::EXPECT, EExpect::Expect)
-                    .parse(arena, state)
-                {
+                match parse_expect.parse(arena, state) {
                     Err((_, _, _)) => {
                         // a hacky way to get expression-based error messages. TODO fix this
                         return Ok((NoProgress, defs, initial));
                     }
-                    Ok((_, _, state)) => {
+
+                    Ok((_, expect_flavor, state)) => {
                         let parse_def_expr = space0_before_e(
                             move |a, s| parse_loc_expr(min_indent + 1, a, s),
                             min_indent,
@@ -677,10 +681,17 @@ fn parse_defs_end<'a>(
 
                         let preceding_comment = Region::new(spaces_before_current_start, start);
 
-                        let value_def = ValueDef::Expect {
-                            condition: arena.alloc(loc_def_expr),
-                            preceding_comment,
+                        let value_def = match expect_flavor {
+                            Either::First(_) => ValueDef::Expect {
+                                condition: arena.alloc(loc_def_expr),
+                                preceding_comment,
+                            },
+                            Either::Second(_) => ValueDef::ExpectFx {
+                                condition: arena.alloc(loc_def_expr),
+                                preceding_comment,
+                            },
                         };
+
                         defs.push_value_def(value_def, region, spaces_before_current, &[]);
 
                         global_state = state;

crates/compiler/parse/src/keyword.rs

@@ -5,5 +5,6 @@ pub const WHEN: &str = "when";
 pub const AS: &str = "as";
 pub const IS: &str = "is";
 pub const EXPECT: &str = "expect";
+pub const EXPECT_FX: &str = "expect-fx";
 
-pub const KEYWORDS: [&str; 7] = [IF, THEN, ELSE, WHEN, AS, IS, EXPECT];
+pub const KEYWORDS: [&str; 8] = [IF, THEN, ELSE, WHEN, AS, IS, EXPECT, EXPECT_FX];

crates/compiler/parse/tests/snapshots/pass/expect_fx.module.result-ast

@@ -0,0 +1,39 @@
+Defs {
+    tags: [
+        Index(2147483648),
+    ],
+    regions: [
+        @25-41,
+    ],
+    space_before: [
+        Slice(start = 0, length = 1),
+    ],
+    space_after: [
+        Slice(start = 1, length = 1),
+    ],
+    spaces: [
+        LineComment(
+            " expecting some effects",
+        ),
+        Newline,
+    ],
+    type_defs: [],
+    value_defs: [
+        ExpectFx {
+            condition: @35-41 BinOps(
+                [
+                    (
+                        @35-36 Num(
+                            "5",
+                        ),
+                        @37-39 Equals,
+                    ),
+                ],
+                @40-41 Num(
+                    "2",
+                ),
+            ),
+            preceding_comment: @25-25,
+        },
+    ],
+}

crates/compiler/parse/tests/snapshots/pass/expect_fx.module.roc

@@ -0,0 +1,2 @@
+# expecting some effects
+expect-fx 5 == 2

crates/compiler/parse/tests/test_parse.rs

@@ -251,6 +251,7 @@ mod test_parse {
         pass/spaced_singleton_list.expr,
         pass/spaces_inside_empty_list.expr,
         pass/standalone_module_defs.module,
+        pass/expect_fx.module,
         pass/string_without_escape.expr,
         pass/sub_var_with_spaces.expr,
         pass/sub_with_spaces.expr,