move record index to start of update

crates/compiler/mono/src/ir.rs

@@ -13,7 +13,7 @@ use roc_can::abilities::SpecializationId;
 use roc_can::expr::{AnnotatedMark, ClosureData, ExpectLookup};
 use roc_can::module::ExposedByModule;
 use roc_collections::all::{default_hasher, BumpMap, BumpMapDefault, MutMap};
-use roc_collections::VecMap;
+use roc_collections::{MutSet, VecMap};
 use roc_debug_flags::dbg_do;
 #[cfg(debug_assertions)]
 use roc_debug_flags::{
@@ -1386,6 +1386,7 @@ pub struct Env<'a, 'i> {
     pub abilities: AbilitiesView<'i>,
     pub exposed_by_module: &'i ExposedByModule,
     pub derived_module: &'i SharedDerivedModule,
+    pub struct_indexing: UsageTrackingMap<(Symbol, u64), Symbol>,
 }
 
 impl<'a, 'i> Env<'a, 'i> {
@@ -4225,6 +4226,7 @@ pub fn with_hole<'a>(
             // If this symbol is a raw value, find the real name we gave to its specialized usage.
             if let ReuseSymbol::Value(_symbol) = can_reuse_symbol(
                 env,
+                layout_cache,
                 procs,
                 &roc_can::expr::Expr::Var(symbol, variable),
                 variable,
@@ -4324,7 +4326,7 @@ pub fn with_hole<'a>(
         OpaqueRef { argument, .. } => {
             let (arg_var, loc_arg_expr) = *argument;
 
-            match can_reuse_symbol(env, procs, &loc_arg_expr.value, arg_var) {
+            match can_reuse_symbol(env, layout_cache, procs, &loc_arg_expr.value, arg_var) {
                 // Opaques decay to their argument.
                 ReuseSymbol::Value(symbol) => {
                     let real_name = procs.get_or_insert_symbol_specialization(
@@ -4909,13 +4911,13 @@ pub fn with_hole<'a>(
         RecordUpdate {
             record_var,
             symbol: structure,
-            updates,
+            ref updates,
             ..
         } => {
             use FieldType::*;
 
             enum FieldType<'a> {
-                CopyExisting(u64),
+                CopyExisting,
                 UpdateExisting(&'a roc_can::expr::Field),
             }
 
@@ -4938,43 +4940,55 @@ pub fn with_hole<'a>(
                 Err(_) => return runtime_error(env, "Can't update record with improper layout"),
             };
 
-            let mut field_layouts = Vec::with_capacity_in(sorted_fields.len(), env.arena);
-
-            let mut symbols = Vec::with_capacity_in(sorted_fields.len(), env.arena);
+            // The struct indexing generated by the current context
+            let mut current_struct_indexing = MutSet::default();
+            // The symbols that are used to create the new struct
+            let mut new_struct_symbols = Vec::with_capacity_in(sorted_fields.len(), env.arena);
+            // Information about the fields that are being updated
             let mut fields = Vec::with_capacity_in(sorted_fields.len(), env.arena);
+            let mut index = 0;
+            for (label, _, opt_field_layout) in sorted_fields.iter() {
+                let record_index = (structure, index);
 
-            let mut current = 0;
-            for (label, _, opt_field_layout) in sorted_fields.into_iter() {
                 match opt_field_layout {
                     Err(_) => {
                         debug_assert!(!updates.contains_key(&label));
                         // this was an optional field, and now does not exist!
-                        // do not increment `current`!
+                        // do not increment `index`!
                     }
-                    Ok(field_layout) => {
-                        field_layouts.push(field_layout);
+                    Ok(_field_layout) => {
+                        current_struct_indexing.insert(record_index);
 
+                        let original_struct_symbol = env.unique_symbol();
                         if let Some(field) = updates.get(&label) {
-                            let field_symbol = possible_reuse_symbol_or_specialize(
-                                env,
-                                procs,
-                                layout_cache,
-                                &field.loc_expr.value,
-                                field.var,
-                            );
+                            // TODO this should probably used around here.
+                            // let field_symbol = possible_reuse_symbol_or_specialize(
+                            //     env,
+                            //     procs,
+                            //     layout_cache,
+                            //     &field.loc_expr.value,
+                            //     field.var,
+                            // );
 
+                            env.struct_indexing
+                                .insert(record_index, original_struct_symbol);
+                            let new_struct_symbol = env.unique_symbol();
+                            new_struct_symbols.push(new_struct_symbol);
                             fields.push(UpdateExisting(field));
-                            symbols.push(field_symbol);
                         } else {
-                            fields.push(CopyExisting(current));
-                            symbols.push(env.unique_symbol());
+                            env.struct_indexing
+                                .insert(record_index, original_struct_symbol);
+                            new_struct_symbols
+                                .push(*env.struct_indexing.get(record_index).unwrap());
+                            fields.push(CopyExisting);
                         }
 
-                        current += 1;
+                        index += 1;
                     }
                 }
             }
-            let symbols = symbols.into_bump_slice();
+
+            let new_struct_symbols = new_struct_symbols.into_bump_slice();
 
             let record_layout = layout_cache
                 .from_var(env.arena, record_var, env.subs)
@@ -4985,8 +4999,8 @@ pub fn with_hole<'a>(
                 _ => arena.alloc([record_layout]),
             };
 
-            if symbols.len() == 1 {
-                // TODO we can probably special-case this more, skippiing the generation of
+            if new_struct_symbols.len() == 1 {
+                // TODO we can probably special-case this more, skipping the generation of
                 // UpdateExisting
                 let mut stmt = hole.clone();
 
@@ -4994,7 +5008,7 @@ pub fn with_hole<'a>(
 
                 match what_to_do {
                     UpdateExisting(field) => {
-                        substitute_in_exprs(env.arena, &mut stmt, assigned, symbols[0]);
+                        substitute_in_exprs(env.arena, &mut stmt, assigned, new_struct_symbols[0]);
 
                         stmt = assign_to_symbol(
                             env,
@@ -5002,11 +5016,11 @@ pub fn with_hole<'a>(
                             layout_cache,
                             field.var,
                             *field.loc_expr.clone(),
-                            symbols[0],
+                            new_struct_symbols[0],
                             stmt,
                         );
                     }
-                    CopyExisting(_) => {
+                    CopyExisting => {
                         unreachable!(
                             r"when a record has just one field and is updated, it must update that one field"
                         );
@@ -5015,12 +5029,12 @@ pub fn with_hole<'a>(
 
                 stmt
             } else {
-                let expr = Expr::Struct(symbols);
+                let expr = Expr::Struct(new_struct_symbols);
                 let mut stmt = Stmt::Let(assigned, expr, record_layout, hole);
 
-                let it = field_layouts.iter().zip(symbols.iter()).zip(fields);
+                let it = new_struct_symbols.iter().zip(fields);
 
-                for ((field_layout, symbol), what_to_do) in it {
+                for (new_struct_symbol, what_to_do) in it {
                     match what_to_do {
                         UpdateExisting(field) => {
                             stmt = assign_to_symbol(
@@ -5029,47 +5043,55 @@ pub fn with_hole<'a>(
                                 layout_cache,
                                 field.var,
                                 *field.loc_expr.clone(),
-                                *symbol,
+                                *new_struct_symbol,
                                 stmt,
                             );
                         }
-                        CopyExisting(index) => {
-                            let structure_needs_specialization =
-                                procs.ability_member_aliases.get(structure).is_some()
-                                    || procs.is_module_thunk(structure)
-                                    || procs.is_imported_module_thunk(structure);
-
-                            let specialized_structure_sym = if structure_needs_specialization {
-                                // We need to specialize the record now; create a new one for it.
-                                // TODO: reuse this symbol for all updates
-                                env.unique_symbol()
-                            } else {
-                                // The record is already good.
-                                structure
-                            };
-
-                            let access_expr = Expr::StructAtIndex {
-                                structure: specialized_structure_sym,
-                                index,
-                                field_layouts,
-                            };
-                            stmt =
-                                Stmt::Let(*symbol, access_expr, *field_layout, arena.alloc(stmt));
-
-                            if structure_needs_specialization {
-                                stmt = specialize_symbol(
-                                    env,
-                                    procs,
-                                    layout_cache,
-                                    Some(record_var),
-                                    specialized_structure_sym,
-                                    env.arena.alloc(stmt),
-                                    structure,
-                                );
-                            }
-                        }
+                        CopyExisting => {}
                     }
                 }
+
+                let structure_needs_specialization =
+                    procs.ability_member_aliases.get(structure).is_some()
+                        || procs.is_module_thunk(structure)
+                        || procs.is_imported_module_thunk(structure);
+
+                let specialized_structure_sym = if structure_needs_specialization {
+                    // We need to specialize the record now; create a new one for it.
+                    env.unique_symbol()
+                } else {
+                    // The record is already good.
+                    structure
+                };
+
+                for record_index in current_struct_indexing.iter() {
+                    let (symbol, usage) = env.struct_indexing.pop(record_index).unwrap();
+                    match usage {
+                        Usage::Used => {
+                            let layout = field_layouts[record_index.1 as usize];
+                            let access_expr = Expr::StructAtIndex {
+                                structure: specialized_structure_sym,
+                                index: record_index.1,
+                                field_layouts,
+                            };
+                            stmt = Stmt::Let(symbol, access_expr, layout, arena.alloc(stmt));
+                        }
+                        Usage::Unused => {}
+                    }
+                }
+
+                if structure_needs_specialization {
+                    stmt = specialize_symbol(
+                        env,
+                        procs,
+                        layout_cache,
+                        Some(record_var),
+                        specialized_structure_sym,
+                        env.arena.alloc(stmt),
+                        structure,
+                    );
+                }
+
                 stmt
             }
         }
@@ -5227,7 +5249,7 @@ pub fn with_hole<'a>(
                     // re-use that symbol, and don't define its value again
                     let mut result;
                     use ReuseSymbol::*;
-                    match can_reuse_symbol(env, procs, &loc_expr.value, fn_var) {
+                    match can_reuse_symbol(env, layout_cache, procs, &loc_expr.value, fn_var) {
                         LocalFunction(_) => {
                             unreachable!("if this was known to be a function, we would not be here")
                         }
@@ -5685,22 +5707,28 @@ fn compile_struct_like<'a, L, UnusedLayout>(
         // TODO how should function pointers be handled here?
         use ReuseSymbol::*;
         match take_elem_expr(index) {
-            Some((var, loc_expr)) => match can_reuse_symbol(env, procs, &loc_expr.value, var) {
-                Imported(symbol) | LocalFunction(symbol) | UnspecializedExpr(symbol) => {
-                    elem_symbols.push(symbol);
-                    can_elems.push(Field::FunctionOrUnspecialized(symbol, variable));
+            Some((var, loc_expr)) => {
+                match can_reuse_symbol(env, layout_cache, procs, &loc_expr.value, var) {
+                    Imported(symbol) | LocalFunction(symbol) | UnspecializedExpr(symbol) => {
+                        elem_symbols.push(symbol);
+                        can_elems.push(Field::FunctionOrUnspecialized(symbol, variable));
+                    }
+                    Value(symbol) => {
+                        let reusable = procs.get_or_insert_symbol_specialization(
+                            env,
+                            layout_cache,
+                            symbol,
+                            var,
+                        );
+                        elem_symbols.push(reusable);
+                        can_elems.push(Field::ValueSymbol);
+                    }
+                    NotASymbol => {
+                        elem_symbols.push(env.unique_symbol());
+                        can_elems.push(Field::Field(var, *loc_expr));
+                    }
                 }
-                Value(symbol) => {
-                    let reusable =
-                        procs.get_or_insert_symbol_specialization(env, layout_cache, symbol, var);
-                    elem_symbols.push(reusable);
-                    can_elems.push(Field::ValueSymbol);
-                }
-                NotASymbol => {
-                    elem_symbols.push(env.unique_symbol());
-                    can_elems.push(Field::Field(var, *loc_expr));
-                }
-            },
+            }
             None => {
                 // this field was optional, but not given
                 continue;
@@ -6816,7 +6844,7 @@ pub fn from_can<'a>(
             store_specialized_expectation_lookups(env, [variable], &[spec_var]);
 
             let symbol_is_reused = matches!(
-                can_reuse_symbol(env, procs, &loc_condition.value, variable),
+                can_reuse_symbol(env, layout_cache, procs, &loc_condition.value, variable),
                 ReuseSymbol::Value(_)
             );
 
@@ -7615,7 +7643,8 @@ enum ReuseSymbol {
 
 fn can_reuse_symbol<'a>(
     env: &mut Env<'a, '_>,
-    procs: &Procs<'a>,
+    layout_cache: &mut LayoutCache<'a>,
+    procs: &mut Procs<'a>,
     expr: &roc_can::expr::Expr,
     expr_var: Variable,
 ) -> ReuseSymbol {
@@ -7627,6 +7656,61 @@ fn can_reuse_symbol<'a>(
             late_resolve_ability_specialization(env, *member, *specialization_id, expr_var)
         }
         Var(symbol, _) => *symbol,
+        RecordAccess {
+            record_var,
+            field,
+            loc_expr,
+            ..
+        } => {
+            let sorted_fields_result = {
+                let mut layout_env = layout::Env::from_components(
+                    layout_cache,
+                    env.subs,
+                    env.arena,
+                    env.target_info,
+                );
+                layout::sort_record_fields(&mut layout_env, *record_var)
+            };
+
+            let sorted_fields = match sorted_fields_result {
+                Ok(fields) => fields,
+                Err(_) => unreachable!("Can't access record with improper layout"),
+            };
+
+            let index =
+                sorted_fields
+                    .into_iter()
+                    .enumerate()
+                    .find_map(
+                        |(current, (label, _, _))| {
+                            if label == *field {
+                                Some(current)
+                            } else {
+                                None
+                            }
+                        },
+                    );
+
+            let struct_index = index.expect("field not in its own type");
+
+            let struct_symbol = possible_reuse_symbol_or_specialize(
+                env,
+                procs,
+                layout_cache,
+                &loc_expr.value,
+                *record_var,
+            );
+
+            match env
+                .struct_indexing
+                .get((struct_symbol, struct_index as u64))
+            {
+                Some(symbol) => *symbol,
+                None => {
+                    return NotASymbol;
+                }
+            }
+        }
         _ => return NotASymbol,
     };
 
@@ -7660,7 +7744,7 @@ fn possible_reuse_symbol_or_specialize<'a>(
     expr: &roc_can::expr::Expr,
     var: Variable,
 ) -> Symbol {
-    match can_reuse_symbol(env, procs, expr, var) {
+    match can_reuse_symbol(env, layout_cache, procs, expr, var) {
         ReuseSymbol::Value(symbol) => {
             procs.get_or_insert_symbol_specialization(env, layout_cache, symbol, var)
         }
@@ -7999,7 +8083,7 @@ fn assign_to_symbol<'a>(
     result: Stmt<'a>,
 ) -> Stmt<'a> {
     use ReuseSymbol::*;
-    match can_reuse_symbol(env, procs, &loc_arg.value, arg_var) {
+    match can_reuse_symbol(env, layout_cache, procs, &loc_arg.value, arg_var) {
         Imported(original) | LocalFunction(original) | UnspecializedExpr(original) => {
             // for functions we must make sure they are specialized correctly
             specialize_symbol(
@@ -9983,3 +10067,40 @@ where
 
     answer
 }
+
+enum Usage {
+    Used,
+    Unused,
+}
+
+pub struct UsageTrackingMap<K, V>
+where
+    K: std::cmp::Eq + std::hash::Hash,
+{
+    map: MutMap<K, (V, Usage)>,
+}
+
+impl<K, V> UsageTrackingMap<K, V>
+where
+    K: std::cmp::Eq + std::hash::Hash,
+{
+    pub fn new() -> Self {
+        Self {
+            map: MutMap::default(),
+        }
+    }
+
+    pub fn insert(&mut self, key: K, value: V) {
+        self.map.insert(key, (value, Usage::Unused));
+    }
+
+    pub fn get(&mut self, key: K) -> Option<&V> {
+        let (value, usage) = self.map.get_mut(&key)?;
+        *usage = Usage::Used;
+        Some(value)
+    }
+
+    fn pop(&mut self, key: &K) -> Option<(V, Usage)> {
+        self.map.remove(key)
+    }
+}