Store specialized variables of expect lookups in expect frames

crates/compiler/gen_llvm/src/llvm/build.rs

@@ -2586,7 +2586,7 @@ pub fn build_exp_stmt<'a, 'ctx, 'env>(
             condition: cond_symbol,
             region,
             lookups,
-            variables: _, // TODO
+            variables,
             remainder,
         } => {
             let bd = env.builder;
@@ -2621,6 +2621,7 @@ pub fn build_exp_stmt<'a, 'ctx, 'env>(
                             *cond_symbol,
                             *region,
                             lookups,
+                            variables,
                         );
 
                         if let LlvmBackendMode::BinaryDev = env.mode {
@@ -2655,7 +2656,7 @@ pub fn build_exp_stmt<'a, 'ctx, 'env>(
             condition: cond_symbol,
             region,
             lookups,
-            variables: _, // TODO
+            variables,
             remainder,
         } => {
             let bd = env.builder;
@@ -2690,6 +2691,7 @@ pub fn build_exp_stmt<'a, 'ctx, 'env>(
                             *cond_symbol,
                             *region,
                             lookups,
+                            variables,
                         );
 
                         bd.build_unconditional_branch(then_block);

crates/compiler/gen_llvm/src/llvm/expect.rs

@@ -10,6 +10,7 @@ use inkwell::values::{BasicValueEnum, FunctionValue, IntValue, PointerValue};
 use inkwell::AddressSpace;
 use roc_builtins::bitcode;
 use roc_module::symbol::Symbol;
+use roc_mono::ir::LookupType;
 use roc_mono::layout::{Builtin, Layout, LayoutIds, UnionLayout};
 use roc_region::all::Region;
 
@@ -143,6 +144,20 @@ pub(crate) fn notify_parent_dbg(env: &Env, shared_memory: &SharedMemoryPointer)
     );
 }
 
+// Shape of expect frame:
+//
+//     ===
+//     Fixed-size header
+//     ===
+// /-- ptr_lookup_1  (ptr_size)
+// |   var_lookup_1  (u32)
+// |   ..
+// |   ptr_lookup_n  (ptr_size)
+// |   var_lookup_n  (u32)
+// \-> lookup_val_1  (varsize)
+//     ..
+//     lookup_val_n  (varsize)
+//
 pub(crate) fn clone_to_shared_memory<'a, 'ctx, 'env>(
     env: &Env<'a, 'ctx, 'env>,
     scope: &Scope<'a, 'ctx>,
@@ -151,6 +166,7 @@ pub(crate) fn clone_to_shared_memory<'a, 'ctx, 'env>(
     condition: Symbol,
     region: Region,
     lookups: &[Symbol],
+    lookup_variables: &[LookupType],
 ) {
     let original_ptr = shared_memory.0;
 
@@ -160,9 +176,11 @@ pub(crate) fn clone_to_shared_memory<'a, 'ctx, 'env>(
 
     let after_header = offset;
 
-    let space_for_offsets = env
+    let space_for_offsets = env.ptr_int().const_int(
-        .ptr_int()
+        (lookups.len() * env.target_info.ptr_size() + lookups.len() * std::mem::size_of::<u32>())
-        .const_int((lookups.len() * env.target_info.ptr_size()) as _, false);
+            as _,
+        false,
+    );
 
     let mut lookup_starts = bumpalo::collections::Vec::with_capacity_in(lookups.len(), env.arena);
 
@@ -203,14 +221,43 @@ pub(crate) fn clone_to_shared_memory<'a, 'ctx, 'env>(
     {
         let mut offset = after_header;
 
-        for lookup_start in lookup_starts {
+        for (lookup_start, lookup_var) in lookup_starts.into_iter().zip(lookup_variables) {
-            build_copy(env, original_ptr, offset, lookup_start.into());
+            // Store the pointer to the value
+            {
+                build_copy(env, original_ptr, offset, lookup_start.into());
 
-            let ptr_width = env
+                let ptr_width = env
-                .ptr_int()
+                    .ptr_int()
-                .const_int(env.target_info.ptr_size() as _, false);
+                    .const_int(env.target_info.ptr_size() as _, false);
 
-            offset = env.builder.build_int_add(offset, ptr_width, "offset")
+                offset = env.builder.build_int_add(offset, ptr_width, "offset");
+            }
+
+            // Store the specialized variable of the value
+            {
+                let ptr = unsafe {
+                    env.builder
+                        .build_in_bounds_gep(original_ptr, &[offset], "at_current_offset")
+                };
+
+                let u32_ptr = env.context.i32_type().ptr_type(AddressSpace::Generic);
+                let ptr = env
+                    .builder
+                    .build_pointer_cast(ptr, u32_ptr, "cast_ptr_type");
+
+                let var_value = env
+                    .context
+                    .i32_type()
+                    .const_int(lookup_var.index() as _, false);
+
+                env.builder.build_store(ptr, var_value);
+
+                let var_size = env
+                    .ptr_int()
+                    .const_int(std::mem::size_of::<u32>() as _, false);
+
+                offset = env.builder.build_int_add(offset, var_size, "offset");
+            }
         }
     }
 

crates/compiler/gen_llvm/src/llvm/lowlevel.rs

@@ -1136,6 +1136,7 @@ pub(crate) fn run_low_level<'a, 'ctx, 'env>(
                     args[0],
                     region,
                     &[args[0]],
+                    &[roc_mono::ir::LookupType::NULL], // TODO
                 );
 
                 crate::llvm::expect::notify_parent_dbg(env, &shared_memory);

crates/compiler/mono/src/ir.rs

@@ -1604,6 +1604,9 @@ pub fn cond<'a>(
 
 pub type Stores<'a> = &'a [(Symbol, Layout<'a>, Expr<'a>)];
 
+/// The specialized type of a lookup. Represented as a type-variable.
+pub type LookupType = Variable;
+
 #[derive(Clone, Debug, PartialEq)]
 pub enum Stmt<'a> {
     Let(Symbol, Expr<'a>, Layout<'a>, &'a Stmt<'a>),
@@ -1625,7 +1628,7 @@ pub enum Stmt<'a> {
         condition: Symbol,
         region: Region,
         lookups: &'a [Symbol],
-        variables: &'a [Variable],
+        variables: &'a [LookupType],
         /// what happens after the expect
         remainder: &'a Stmt<'a>,
     },
@@ -1633,7 +1636,7 @@ pub enum Stmt<'a> {
         condition: Symbol,
         region: Region,
         lookups: &'a [Symbol],
-        variables: &'a [Variable],
+        variables: &'a [LookupType],
         /// what happens after the expect
         remainder: &'a Stmt<'a>,
     },
@@ -6573,7 +6576,8 @@ pub fn from_can<'a>(
             let cond_symbol = env.unique_symbol();
 
             let mut lookups = Vec::with_capacity_in(lookups_in_cond.len(), env.arena);
-            let mut variables = Vec::with_capacity_in(lookups_in_cond.len(), env.arena);
+            let mut lookup_variables = Vec::with_capacity_in(lookups_in_cond.len(), env.arena);
+            let mut specialized_variables = Vec::with_capacity_in(lookups_in_cond.len(), env.arena);
 
             for ExpectLookup {
                 symbol,
@@ -6595,22 +6599,23 @@ pub fn from_can<'a>(
                     .expectation_subs
                     .as_deref_mut()
                     .expect("if expects are compiled, their subs should be available");
-                let variable = expectation_subs.fresh_unnamed_flex_var();
+                let spec_var = expectation_subs.fresh_unnamed_flex_var();
 
                 if !env.subs.is_function(var) {
                     // Exclude functions from lookups
                     lookups.push(symbol);
-                    variables.push(variable);
+                    lookup_variables.push(var);
+                    specialized_variables.push(spec_var);
                 }
             }
 
-            let variables = variables.into_bump_slice();
+            let specialized_variables = specialized_variables.into_bump_slice();
 
             let mut stmt = Stmt::Expect {
                 condition: cond_symbol,
                 region: loc_condition.region,
                 lookups: lookups.into_bump_slice(),
-                variables,
+                variables: specialized_variables,
                 remainder: env.arena.alloc(rest),
             };
 
@@ -6624,7 +6629,9 @@ pub fn from_can<'a>(
                 env.arena.alloc(stmt),
             );
 
-            store_specialized_expectation_lookups(env, lookups_in_cond, variables);
+            // Now that the condition has been specialized, export the specialized types of our
+            // lookups into the expectation subs.
+            store_specialized_expectation_lookups(env, lookup_variables, specialized_variables);
 
             stmt
         }
@@ -6638,7 +6645,8 @@ pub fn from_can<'a>(
             let cond_symbol = env.unique_symbol();
 
             let mut lookups = Vec::with_capacity_in(lookups_in_cond.len(), env.arena);
-            let mut variables = Vec::with_capacity_in(lookups_in_cond.len(), env.arena);
+            let mut lookup_variables = Vec::with_capacity_in(lookups_in_cond.len(), env.arena);
+            let mut specialized_variables = Vec::with_capacity_in(lookups_in_cond.len(), env.arena);
 
             for ExpectLookup {
                 symbol,
@@ -6660,22 +6668,23 @@ pub fn from_can<'a>(
                     .expectation_subs
                     .as_deref_mut()
                     .expect("if expects are compiled, their subs should be available");
-                let variable = expectation_subs.fresh_unnamed_flex_var();
+                let spec_var = expectation_subs.fresh_unnamed_flex_var();
 
                 if !env.subs.is_function(var) {
                     // Exclude functions from lookups
                     lookups.push(symbol);
-                    variables.push(variable);
+                    lookup_variables.push(var);
+                    specialized_variables.push(spec_var);
                 }
             }
 
-            let variables = variables.into_bump_slice();
+            let specialized_variables = specialized_variables.into_bump_slice();
 
             let mut stmt = Stmt::ExpectFx {
                 condition: cond_symbol,
                 region: loc_condition.region,
                 lookups: lookups.into_bump_slice(),
-                variables,
+                variables: specialized_variables,
                 remainder: env.arena.alloc(rest),
             };
 
@@ -6689,7 +6698,7 @@ pub fn from_can<'a>(
                 env.arena.alloc(stmt),
             );
 
-            store_specialized_expectation_lookups(env, lookups_in_cond, variables);
+            store_specialized_expectation_lookups(env, lookup_variables, specialized_variables);
 
             stmt
         }
@@ -6776,14 +6785,14 @@ pub fn from_can<'a>(
 
 fn store_specialized_expectation_lookups(
     env: &mut Env,
-    lookups_in_cond: impl IntoIterator<Item = ExpectLookup>,
+    lookup_variables: impl IntoIterator<Item = Variable>,
-    variables: &[Variable],
+    specialized_variables: &[Variable],
 ) {
     let subs = &env.subs;
     let expectation_subs = env.expectation_subs.as_deref_mut().unwrap();
-    for (ExpectLookup { var, .. }, stored_var) in lookups_in_cond.into_iter().zip(variables) {
+    for (lookup_var, stored_var) in lookup_variables.into_iter().zip(specialized_variables) {
         let stored_specialized_var =
-            storage_copy_var_to(&mut Default::default(), subs, expectation_subs, var);
+            storage_copy_var_to(&mut Default::default(), subs, expectation_subs, lookup_var);
         let stored_specialized_desc = expectation_subs.get(stored_specialized_var);
         expectation_subs.union(*stored_var, stored_specialized_var, stored_specialized_desc);
     }

crates/repl_expect/src/lib.rs

@@ -46,12 +46,19 @@ pub fn get_values<'a>(
     let app = arena.alloc(app);
 
     for (i, variable) in variables.iter().enumerate() {
-        let start = app.memory.deref_usize(start_offset + i * 8);
+        let size_of_lookup_header = 8 /* pointer to value */ + 4 /* type variable */;
+
+        let start = app
+            .memory
+            .deref_usize(start_offset + i * size_of_lookup_header);
+        let variable = app.memory.deref_u32(
+            start_offset + i * size_of_lookup_header + 8, /* skip the pointer */
+        );
+        let variable = unsafe { Variable::from_index(variable) };
+
         app.offset = start;
 
         let expr = {
-            let variable = *variable;
-
             // TODO: pass layout_cache to jit_to_ast directly
             let mut layout_cache = LayoutCache::new(layout_interner.fork(), target_info);
             let layout = layout_cache.from_var(arena, variable, subs).unwrap();

crates/repl_expect/src/run.rs

@@ -574,14 +574,13 @@ fn render_expect_failure<'a>(
         None => panic!("region {failure_region:?} not in list of expects"),
         Some(current) => current,
     };
-    let subs = arena.alloc(&mut data.subs);
 
-    let (symbols, variables) = split_expect_lookups(subs, current);
+    let (symbols, variables) = split_expect_lookups(&data.subs, current);
 
     let (offset, expressions) = crate::get_values(
         target_info,
         arena,
-        subs,
+        &data.subs,
         interns,
         layout_interner,
         start,
@@ -591,7 +590,7 @@ fn render_expect_failure<'a>(
 
     renderer.render_failure(
         writer,
-        subs,
+        &mut data.subs,
         &symbols,
         &variables,
         &expressions,