Resolve ability specializations looked up in expects

Cargo.lock

@@ -3959,6 +3959,7 @@ dependencies = [
  "pretty_assertions",
  "roc_build",
  "roc_builtins",
+ "roc_can",
  "roc_collections",
  "roc_gen_llvm",
  "roc_intern",

crates/compiler/can/src/expr.rs

@@ -230,14 +230,14 @@ pub enum Expr {
     Expect {
         loc_condition: Box<Loc<Expr>>,
         loc_continuation: Box<Loc<Expr>>,
-        lookups_in_cond: Vec<(Symbol, Variable)>,
+        lookups_in_cond: Vec<ExpectLookup>,
     },
 
     // not parsed, but is generated when lowering toplevel effectful expects
     ExpectFx {
         loc_condition: Box<Loc<Expr>>,
         loc_continuation: Box<Loc<Expr>>,
-        lookups_in_cond: Vec<(Symbol, Variable)>,
+        lookups_in_cond: Vec<ExpectLookup>,
     },
 
     /// Rendered as empty box in editor
@@ -247,6 +247,13 @@ pub enum Expr {
     RuntimeError(RuntimeError),
 }
 
+#[derive(Clone, Copy, Debug)]
+pub struct ExpectLookup {
+    pub symbol: Symbol,
+    pub var: Variable,
+    pub ability_info: Option<SpecializationId>,
+}
+
 impl Expr {
     pub fn category(&self) -> Category {
         match self {
@@ -2541,7 +2548,7 @@ impl Declarations {
             })
     }
 
-    pub fn expects(&self) -> VecMap<Region, Vec<(Symbol, Variable)>> {
+    pub fn expects(&self) -> VecMap<Region, Vec<ExpectLookup>> {
         let mut collector = ExpectCollector {
             expects: VecMap::default(),
         };
@@ -2625,9 +2632,9 @@ pub struct DestructureDef {
     pub pattern_vars: VecMap<Symbol, Variable>,
 }
 
-fn get_lookup_symbols(expr: &Expr) -> Vec<(Symbol, Variable)> {
+fn get_lookup_symbols(expr: &Expr) -> Vec<ExpectLookup> {
     let mut stack: Vec<&Expr> = vec![expr];
-    let mut symbols = Vec::new();
+    let mut lookups: Vec<ExpectLookup> = Vec::new();
 
     while let Some(expr) = stack.pop() {
         match expr {
@@ -2636,11 +2643,23 @@ fn get_lookup_symbols(expr: &Expr) -> Vec<(Symbol, Variable)> {
                 symbol,
                 record_var: var,
                 ..
-            }
+            } => {
-            | Expr::AbilityMember(symbol, _, var) => {
                 // Don't introduce duplicates, or make unused variables
-                if !symbols.iter().any(|(sym, _)| sym == symbol) {
+                if !lookups.iter().any(|l| l.symbol == *symbol) {
-                    symbols.push((*symbol, *var));
+                    lookups.push(ExpectLookup {
+                        symbol: *symbol,
+                        var: *var,
+                        ability_info: None,
+                    });
+                }
+            }
+            Expr::AbilityMember(symbol, spec_id, var) => {
+                if !lookups.iter().any(|l| l.symbol == *symbol) {
+                    lookups.push(ExpectLookup {
+                        symbol: *symbol,
+                        var: *var,
+                        ability_info: *spec_id,
+                    });
                 }
             }
             Expr::List { loc_elems, .. } => {
@@ -2737,7 +2756,7 @@ fn get_lookup_symbols(expr: &Expr) -> Vec<(Symbol, Variable)> {
         }
     }
 
-    symbols
+    lookups
 }
 
 /// Here we transform
@@ -2784,14 +2803,22 @@ fn toplevel_expect_to_inline_expect_help(mut loc_expr: Loc<Expr>, has_effects: b
     loop {
         match loc_expr.value {
             Expr::LetNonRec(boxed_def, remainder) => {
-                lookups_in_cond.extend(boxed_def.pattern_vars.iter().map(|(a, b)| (*a, *b)));
+                lookups_in_cond.extend(boxed_def.pattern_vars.iter().map(|(a, b)| ExpectLookup {
+                    symbol: *a,
+                    var: *b,
+                    ability_info: None,
+                }));
 
                 stack.push(StoredDef::NonRecursive(loc_expr.region, boxed_def));
                 loc_expr = *remainder;
             }
             Expr::LetRec(defs, remainder, mark) => {
                 for def in &defs {
-                    lookups_in_cond.extend(def.pattern_vars.iter().map(|(a, b)| (*a, *b)));
+                    lookups_in_cond.extend(def.pattern_vars.iter().map(|(a, b)| ExpectLookup {
+                        symbol: *a,
+                        var: *b,
+                        ability_info: None,
+                    }));
                 }
 
                 stack.push(StoredDef::Recursive(loc_expr.region, defs, mark));
@@ -2834,7 +2861,7 @@ fn toplevel_expect_to_inline_expect_help(mut loc_expr: Loc<Expr>, has_effects: b
 }
 
 struct ExpectCollector {
-    expects: VecMap<Region, Vec<(Symbol, Variable)>>,
+    expects: VecMap<Region, Vec<ExpectLookup>>,
 }
 
 impl crate::traverse::Visitor for ExpectCollector {

crates/compiler/can/src/module.rs

@@ -3,7 +3,7 @@ use crate::annotation::canonicalize_annotation;
 use crate::def::{canonicalize_defs, Def};
 use crate::effect_module::HostedGeneratedFunctions;
 use crate::env::Env;
-use crate::expr::{ClosureData, Declarations, Expr, Output, PendingDerives};
+use crate::expr::{ClosureData, Declarations, ExpectLookup, Expr, Output, PendingDerives};
 use crate::pattern::{BindingsFromPattern, Pattern};
 use crate::scope::Scope;
 use bumpalo::Bump;
@@ -130,7 +130,7 @@ pub struct Module {
     pub aliases: MutMap<Symbol, (bool, Alias)>,
     pub rigid_variables: RigidVariables,
     pub abilities_store: PendingAbilitiesStore,
-    pub loc_expects: VecMap<Region, Vec<(Symbol, Variable)>>,
+    pub loc_expects: VecMap<Region, Vec<ExpectLookup>>,
 }
 
 #[derive(Debug, Default)]
@@ -152,7 +152,7 @@ pub struct ModuleOutput {
     pub symbols_from_requires: Vec<(Loc<Symbol>, Loc<Type>)>,
     pub pending_derives: PendingDerives,
     pub scope: Scope,
-    pub loc_expects: VecMap<Region, Vec<(Symbol, Variable)>>,
+    pub loc_expects: VecMap<Region, Vec<ExpectLookup>>,
 }
 
 fn validate_generate_with<'a>(

crates/compiler/constrain/src/expr.rs

@@ -13,8 +13,8 @@ use roc_can::expected::Expected::{self, *};
 use roc_can::expected::PExpected;
 use roc_can::expr::Expr::{self, *};
 use roc_can::expr::{
-    AccessorData, AnnotatedMark, ClosureData, DeclarationTag, Declarations, DestructureDef, Field,
+    AccessorData, AnnotatedMark, ClosureData, DeclarationTag, Declarations, DestructureDef,
-    FunctionDef, OpaqueWrapFunctionData, WhenBranch,
+    ExpectLookup, Field, FunctionDef, OpaqueWrapFunctionData, WhenBranch,
 };
 use roc_can::pattern::Pattern;
 use roc_can::traverse::symbols_introduced_from_pattern;
@@ -519,7 +519,12 @@ pub fn constrain_expr(
 
             let mut vars = Vec::with_capacity(lookups_in_cond.len());
 
-            for (symbol, var) in lookups_in_cond.iter() {
+            for ExpectLookup {
+                symbol,
+                var,
+                ability_info: _,
+            } in lookups_in_cond.iter()
+            {
                 vars.push(*var);
 
                 all_constraints.push(constraints.lookup(
@@ -566,7 +571,12 @@ pub fn constrain_expr(
 
             let mut vars = Vec::with_capacity(lookups_in_cond.len());
 
-            for (symbol, var) in lookups_in_cond.iter() {
+            for ExpectLookup {
+                symbol,
+                var,
+                ability_info: _,
+            } in lookups_in_cond.iter()
+            {
                 vars.push(*var);
 
                 all_constraints.push(constraints.lookup(

crates/compiler/load_internal/src/file.rs

@@ -7,8 +7,8 @@ use parking_lot::Mutex;
 use roc_builtins::roc::module_source;
 use roc_can::abilities::{AbilitiesStore, PendingAbilitiesStore, ResolvedImpl};
 use roc_can::constraint::{Constraint as ConstraintSoa, Constraints};
-use roc_can::expr::Declarations;
 use roc_can::expr::PendingDerives;
+use roc_can::expr::{Declarations, ExpectLookup};
 use roc_can::module::{
     canonicalize_module_defs, ExposedByModule, ExposedForModule, ExposedModuleTypes, Module,
     ResolvedImplementations, TypeState,
@@ -719,7 +719,7 @@ pub enum EntryPoint<'a> {
 pub struct Expectations {
     pub subs: roc_types::subs::Subs,
     pub path: PathBuf,
-    pub expectations: VecMap<Region, Vec<(Symbol, Variable)>>,
+    pub expectations: VecMap<Region, Vec<ExpectLookup>>,
     pub ident_ids: IdentIds,
 }
 
@@ -763,7 +763,7 @@ struct ParsedModule<'a> {
     header_for: HeaderFor<'a>,
 }
 
-type LocExpects = VecMap<Region, Vec<(Symbol, Variable)>>;
+type LocExpects = VecMap<Region, Vec<ExpectLookup>>;
 
 /// A message sent out _from_ a worker thread,
 /// representing a result of work done, or a request for further work

crates/compiler/mono/src/ir.rs

@@ -9,7 +9,7 @@ use bumpalo::collections::{CollectIn, Vec};
 use bumpalo::Bump;
 use roc_builtins::bitcode::{FloatWidth, IntWidth};
 use roc_can::abilities::SpecializationId;
-use roc_can::expr::{AnnotatedMark, ClosureData, IntValue};
+use roc_can::expr::{AnnotatedMark, ClosureData, ExpectLookup, IntValue};
 use roc_can::module::ExposedByModule;
 use roc_collections::all::{default_hasher, BumpMap, BumpMapDefault, MutMap};
 use roc_collections::VecMap;
@@ -6350,11 +6350,25 @@ pub fn from_can<'a>(
             let rest = from_can(env, variable, loc_continuation.value, procs, layout_cache);
             let cond_symbol = env.unique_symbol();
 
-            let lookups = Vec::from_iter_in(lookups_in_cond.iter().map(|t| t.0), env.arena);
+            let mut lookups = Vec::with_capacity_in(lookups_in_cond.len(), env.arena);
-
             let mut layouts = Vec::with_capacity_in(lookups_in_cond.len(), env.arena);
 
-            for (_, var) in lookups_in_cond {
+            for ExpectLookup {
+                symbol,
+                var,
+                ability_info,
+            } in lookups_in_cond
+            {
+                let symbol = match ability_info {
+                    Some(specialization_id) => late_resolve_ability_specialization(
+                        env,
+                        symbol,
+                        Some(specialization_id),
+                        var,
+                    ),
+                    None => symbol,
+                };
+                lookups.push(symbol);
                 let res_layout = layout_cache.from_var(env.arena, var, env.subs);
                 let layout = return_on_layout_error!(env, res_layout, "Expect");
                 layouts.push(layout);
@@ -6389,11 +6403,25 @@ pub fn from_can<'a>(
             let rest = from_can(env, variable, loc_continuation.value, procs, layout_cache);
             let cond_symbol = env.unique_symbol();
 
-            let lookups = Vec::from_iter_in(lookups_in_cond.iter().map(|t| t.0), env.arena);
+            let mut lookups = Vec::with_capacity_in(lookups_in_cond.len(), env.arena);
-
             let mut layouts = Vec::with_capacity_in(lookups_in_cond.len(), env.arena);
 
-            for (_, var) in lookups_in_cond {
+            for ExpectLookup {
+                symbol,
+                var,
+                ability_info,
+            } in lookups_in_cond
+            {
+                let symbol = match ability_info {
+                    Some(specialization_id) => late_resolve_ability_specialization(
+                        env,
+                        symbol,
+                        Some(specialization_id),
+                        var,
+                    ),
+                    None => symbol,
+                };
+                lookups.push(symbol);
                 let res_layout = layout_cache.from_var(env.arena, var, env.subs);
                 let layout = return_on_layout_error!(env, res_layout, "Expect");
                 layouts.push(layout);

crates/repl_expect/Cargo.toml

@@ -10,6 +10,7 @@ bumpalo = {version = "3.11.0", features = ["collections"]}
 target-lexicon = "0.12.2"
 
 roc_builtins = {path = "../compiler/builtins"}
+roc_can = {path = "../compiler/can"}
 roc_collections = {path = "../compiler/collections"}
 roc_intern = {path = "../compiler/intern"}
 roc_load = {path = "../compiler/load"}

crates/repl_expect/src/run.rs

@@ -4,6 +4,7 @@ use bumpalo::collections::Vec as BumpVec;
 use bumpalo::Bump;
 use inkwell::context::Context;
 use roc_build::link::llvm_module_to_dylib;
+use roc_can::expr::ExpectLookup;
 use roc_collections::{MutSet, VecMap};
 use roc_gen_llvm::{
     llvm::{build::LlvmBackendMode, externs::add_default_roc_externs},
@@ -389,7 +390,16 @@ fn render_expect_failure<'a>(
     };
     let subs = arena.alloc(&mut data.subs);
 
-    let (symbols, variables): (Vec<_>, Vec<_>) = current.iter().map(|(a, b)| (*a, *b)).unzip();
+    let (symbols, variables): (Vec<_>, Vec<_>) = current
+        .iter()
+        .map(
+            |ExpectLookup {
+                 symbol,
+                 var,
+                 ability_info: _,
+             }| (*symbol, *var),
+        )
+        .unzip();
 
     let (offset, expressions) = crate::get_values(
         target_info,