Allow field use after a conditional field move

compiler/src/mir/passes.rs

@@ -1420,17 +1420,14 @@ pub(crate) struct LowerMethod<'a> {
     /// a register. We map fields to registers here so that for field A we
     /// always write to register R, removing the need for tracking field states
     /// separately.
-    field_mapping: HashMap<types::FieldId, RegisterId>,
+    field_mapping: IndexMap<types::FieldId, RegisterId>,
 
     /// The registers to write to when a register is moved.
     drop_flags: HashMap<RegisterId, RegisterId>,
 
     /// Variables to remap to field reads, and the types to expose the fields
     /// as.
     variable_fields: HashMap<types::VariableId, types::FieldId>,
-
-    /// The number of fields that are moved.
-    moved_fields: usize,
 }
 
 impl<'a> LowerMethod<'a> {
@@ -1450,15 +1447,14 @@ impl<'a> LowerMethod<'a> {
             current_block,
             scope: Scope::root_scope(),
             variable_mapping: IndexMap::new(),
-            field_mapping: HashMap::new(),
+            field_mapping: IndexMap::new(),
             drop_flags: HashMap::new(),
             register_states: RegisterStates::new(),
             register_kinds: Vec::new(),
             surrounding_type_register: RegisterId(SELF_ID),
             self_register: RegisterId(SELF_ID),
             variable_fields: HashMap::new(),
             used_variables: HashSet::new(),
-            moved_fields: 0,
         }
     }
 
@@ -1553,7 +1549,6 @@ impl<'a> LowerMethod<'a> {
         location: InstructionLocation,
     ) {
         self.mark_register_as_moved(register);
-        self.partially_move_self_if_field(register);
         self.drop_all_registers(location);
         self.check_for_unused_variables();
         self.return_register(register, location);
@@ -2099,7 +2094,9 @@ impl<'a> LowerMethod<'a> {
             types::Receiver::Implicit => {
                 let reg = self.self_register;
 
-                if !self.register_is_available(self.self_register) {
+                if !self.register_is_available(reg)
+                    || self.self_has_moved_field()
+                {
                     let name = info.id.name(self.db()).clone();
 
                     self.state.diagnostics.implicit_receiver_moved(
@@ -2474,16 +2471,6 @@ impl<'a> LowerMethod<'a> {
                 self.drop_register(old, loc);
             }
 
-            // We allow the use of `self` again once all moved fields are
-            // assigned a new value.
-            if is_moved {
-                self.moved_fields -= 1;
-
-                if self.moved_fields == 0 {
-                    self.mark_register_as_available(self.self_register);
-                }
-            }
-
             self.update_register_state(reg, RegisterState::Available);
             self.current_block_mut().set_field(rec, tid, id, new_val, loc);
         } else {
@@ -3778,18 +3765,7 @@ impl<'a> LowerMethod<'a> {
         let name = &node.name;
         let check_loc = node.location;
 
-        match self.register_state(rec) {
-            RegisterState::Available | RegisterState::PartiallyMoved => {
-                self.check_if_moved(reg, name, check_loc);
-            }
-            _ => {
-                self.state.diagnostics.implicit_receiver_moved(
-                    name,
-                    self.file(),
-                    node.location,
-                );
-            }
-        }
+        self.check_if_moved(reg, name, check_loc);
 
         if info.as_pointer {
             self.current_block_mut().field_pointer(reg, rec, tid, id, loc);
@@ -3827,8 +3803,14 @@ impl<'a> LowerMethod<'a> {
         self.verify_type(node.resolved_type, node.location);
 
         let reg = self.self_register;
+        if !self.register_is_available(reg) || self.self_has_moved_field() {
+            self.state.diagnostics.moved_variable(
+                SELF_NAME,
+                self.file(),
+                node.location,
+            );
+        }
 
-        self.check_if_moved(reg, SELF_NAME, node.location);
         reg
     }
 
@@ -4124,7 +4106,6 @@ impl<'a> LowerMethod<'a> {
 
         if typ.is_owned_or_uni(self.db()) {
             self.mark_register_as_moved(reg);
-            self.partially_move_self_if_field(reg);
 
             if let Some(flag) = self.drop_flags.get(&reg).cloned() {
                 self.current_block_mut().bool_literal(flag, false, ins_loc);
@@ -4229,13 +4210,8 @@ impl<'a> LowerMethod<'a> {
 
         self.check_field_move(register, location);
         self.mark_register_as_moved(register);
-        self.partially_move_self_if_field(register);
         self.record_loop_move(register, location);
 
-        if self.register_kind(register).is_field() {
-            self.mark_register_as_partially_moved(self.self_register);
-        }
-
         if let Some(flag) = self.drop_flags.get(&register).cloned() {
             self.current_block_mut().bool_literal(flag, false, ins_loc);
         }
@@ -4300,7 +4276,6 @@ impl<'a> LowerMethod<'a> {
             }
 
             self.record_loop_move(register, location);
-            self.partially_move_self_if_field(register);
             self.mark_register_as_moved(register);
 
             if let Some(flag) = self.drop_flags.get(&register).cloned() {
@@ -4328,15 +4303,6 @@ impl<'a> LowerMethod<'a> {
         register
     }
 
-    fn partially_move_self_if_field(&mut self, register: RegisterId) {
-        if !self.register_kind(register).is_field() {
-            return;
-        }
-
-        self.moved_fields += 1;
-        self.mark_register_as_partially_moved(self.self_register);
-    }
-
     fn clone_value_type(
         &mut self,
         source: RegisterId,
@@ -4492,22 +4458,23 @@ impl<'a> LowerMethod<'a> {
         });
 
         match self.register_state(self_reg) {
+            RegisterState::Moved => {}
+            RegisterState::MaybeMoved => {
+                self.drop_register(self_reg, location);
+            }
             RegisterState::PartiallyMoved => {
                 self.drop_remaining_fields(&fields, location);
                 self.current_block_mut()
                     .drop_without_dropper(self_reg, location);
             }
-            // Parts of self may be moved conditionally, such as for code like
-            // `if foo { drop(field) }`. In this case we need to determine at
-            // runtime if `self` is to be dropped as a whole or only its memory
-            // should be deallocated.
-            RegisterState::MaybeMoved if partially_moved => {
-                self.drop_partial_self(&fields, location);
+            RegisterState::Available if partially_moved => {
+                self.drop_remaining_fields(&fields, location);
+                self.current_block_mut()
+                    .drop_without_dropper(self_reg, location);
             }
-            RegisterState::Available | RegisterState::MaybeMoved => {
+            RegisterState::Available => {
                 self.drop_register(self_reg, location);
             }
-            RegisterState::Moved => {}
         }
     }
 
@@ -4528,37 +4495,6 @@ impl<'a> LowerMethod<'a> {
         }
     }
 
-    fn drop_partial_self(
-        &mut self,
-        fields: &[(FieldId, TypeRef)],
-        location: InstructionLocation,
-    ) {
-        let self_reg = self.surrounding_type_register;
-        let before_block = self.current_block;
-        let drop_block = self.add_block();
-        let after_block = self.add_block();
-        let drop_flag = self.drop_flags.get(&self_reg).cloned().unwrap();
-
-        self.current_block_mut().branch(
-            drop_flag,
-            drop_block,
-            after_block,
-            location,
-        );
-
-        self.add_edge(before_block, drop_block);
-        self.add_edge(before_block, after_block);
-        self.add_edge(drop_block, after_block);
-
-        self.current_block = drop_block;
-        self.drop_remaining_fields(fields, location);
-        self.current_block_mut().drop_without_dropper(self_reg, location);
-        self.current_block_mut().goto(after_block, location);
-
-        self.current_block = after_block;
-        self.mark_register_as_moved(self_reg);
-    }
-
     fn drop_loop_registers(&mut self, location: InstructionLocation) {
         let mut registers = Vec::new();
         let mut scope = Some(&self.scope);
@@ -4937,6 +4873,23 @@ impl<'a> LowerMethod<'a> {
 
     fn mark_register_as_moved(&mut self, register: RegisterId) {
         self.update_register_state(register, RegisterState::Moved);
+
+        // Moving `self` as a whole also moves all of its fields.
+        if register == self.self_register {
+            for &field in self.field_mapping.values() {
+                self.register_states.set(
+                    self.current_block,
+                    field,
+                    RegisterState::Moved,
+                );
+            }
+        }
+    }
+
+    /// Returns `true` if any field of `self` has been (possibly) moved.
+    fn self_has_moved_field(&mut self) -> bool {
+        (0..self.field_mapping.len())
+            .any(|i| !self.register_is_available(self.field_mapping[i]))
     }
 
     fn mark_register_as_available(&mut self, register: RegisterId) {

std/fixtures/diagnostics/closure_explicitly_captures_partially_moved_self.inko

@@ -0,0 +1,21 @@
+import std.drop (drop)
+
+type Inner {}
+
+type inline Outer {
+  let @a: Inner
+  let @b: Inner
+
+  fn move submit {
+    drop(@a)
+    fn { self.helper }.call
+  }
+
+  fn helper {}
+}
+
+type async Main {
+  fn async main {
+    Outer(a: Inner(), b: Inner()).submit
+  }
+}

std/fixtures/diagnostics/closure_implicitly_captures_partially_moved_self.inko

@@ -0,0 +1,21 @@
+import std.drop (drop)
+
+type Inner {}
+
+type inline Outer {
+  let @a: Inner
+  let @b: Inner
+
+  fn move submit {
+    drop(@a)
+    fn { helper }.call
+  }
+
+  fn helper {}
+}
+
+type async Main {
+  fn async main {
+    Outer(a: Inner(), b: Inner()).submit
+  }
+}

std/fixtures/diagnostics/field_used_after_conditional_field_move.inko

@@ -0,0 +1,22 @@
+import std.drop (drop)
+
+type Inner {}
+
+type Foo {}
+
+type inline Outer {
+  let @inner: Inner
+  let @foo: Foo
+
+  fn move submit {
+    if true { drop(@inner) }
+
+    drop(@foo)
+  }
+}
+
+type async Main {
+  fn async main {
+    Outer(inner: Inner(), foo: Foo()).submit
+  }
+}

std/fixtures/diagnostics/field_used_after_conditional_self_move.inko

@@ -0,0 +1,26 @@
+import std.drop (drop)
+
+type Inner {}
+
+type Foo {}
+
+type inline Outer {
+  let @inner: Inner
+  let @foo: Foo
+
+  fn move submit {
+    if true {
+      drop(self)
+    }
+
+    drop(@foo)
+  }
+}
+
+type async Main {
+  fn async main {
+    Outer(inner: Inner(), foo: Foo()).submit
+  }
+}
+
+# field_used_after_conditional_self_move.inko:16:10 error(moved): 'foo' can't be used as it has been moved

std/fixtures/diagnostics/implicit_method_call_after_field_move.inko

@@ -0,0 +1,23 @@
+import std.drop (drop)
+
+type Inner {}
+
+type inline Outer {
+  let @a: Inner
+  let @b: Inner
+
+  fn move submit {
+    drop(@a)
+    helper
+  }
+
+  fn move helper {}
+}
+
+type async Main {
+  fn async main {
+    Outer(a: Inner(), b: Inner()).submit
+  }
+}
+
+# implicit_method_call_after_field_move.inko:11:5 error(moved): 'helper' can't be used, as 'self' has been moved

std/fixtures/diagnostics/self_used_after_field_move.inko

@@ -0,0 +1,21 @@
+import std.drop (drop)
+
+type Inner {}
+
+type inline Outer {
+  let @a: Inner
+  let @b: Inner
+
+  fn move submit {
+    drop(@a)
+    drop(self)
+  }
+}
+
+type async Main {
+  fn async main {
+    Outer(a: Inner(), b: Inner()).submit
+  }
+}
+
+# self_used_after_field_move.inko:11:10 error(moved): 'self' can't be used as it has been moved