Simplify implementation.

Author: cjgillot

compiler/rustc_mir_transform/src/ssa.rs

@@ -101,14 +101,15 @@ impl SsaLocals {
             .retain(|&local| matches!(visitor.assignments[local], Set1::One(_)));
         debug!(?visitor.assignment_order);
 
-        let copy_classes = compute_copy_classes(&mut visitor, body);
-
-        SsaLocals {
+        let mut ssa = SsaLocals {
             assignments: visitor.assignments,
             assignment_order: visitor.assignment_order,
             direct_uses: visitor.direct_uses,
-            copy_classes,
-        }
+            // This is filled by `compute_copy_classes`.
+            copy_classes: IndexVec::default(),
+        };
+        compute_copy_classes(&mut ssa, body);
+        ssa
     }
 
     pub fn num_locals(&self) -> usize {
@@ -261,27 +262,19 @@ impl<'tcx> Visitor<'tcx> for SsaVisitor {
 }
 
 #[instrument(level = "trace", skip(ssa, body))]
-fn compute_copy_classes(ssa: &mut SsaVisitor, body: &Body<'_>) -> IndexVec<Local, Local> {
+fn compute_copy_classes(ssa: &mut SsaLocals, body: &Body<'_>) {
+    let mut direct_uses = std::mem::take(&mut ssa.direct_uses);
     let mut copies = IndexVec::from_fn_n(|l| l, body.local_decls.len());
 
-    for &local in &ssa.assignment_order {
-        debug!(?local);
-
-        // This is not SSA: mark that we don't know the value.
-        debug!(assignments = ?ssa.assignments[local]);
-        let Set1::One(LocationExtended::Plain(loc)) = ssa.assignments[local] else { continue };
-
-        // `loc` must point to a direct assignment to `local`.
-        let Either::Left(stmt) = body.stmt_at(loc) else { bug!() };
-        let Some((_target, rvalue)) = stmt.kind.as_assign() else { bug!() };
-        assert_eq!(_target.as_local(), Some(local));
-
+    for (local, rvalue, _) in ssa.assignments(body) {
         let (Rvalue::Use(Operand::Copy(place) | Operand::Move(place)) | Rvalue::CopyForDeref(place))
             = rvalue
         else { continue };
 
         let Some(rhs) = place.as_local() else { continue };
-        let Set1::One(_) = ssa.assignments[rhs] else { continue };
+        if !ssa.is_ssa(rhs) {
+            continue;
+        }
 
         // We visit in `assignment_order`, ie. reverse post-order, so `rhs` has been
         // visited before `local`, and we just have to copy the representing local.
@@ -302,11 +295,11 @@ fn compute_copy_classes(ssa: &mut SsaVisitor, body: &Body<'_>) -> IndexVec<Local
         } else {
             copies[local] = head;
         }
-        ssa.direct_uses[rhs] -= 1;
+        direct_uses[rhs] -= 1;
     }
 
     debug!(?copies);
-    debug!(?ssa.direct_uses);
+    debug!(?direct_uses);
 
     // Invariant: `copies` must point to the head of an equivalence class.
     #[cfg(debug_assertions)]
@@ -315,7 +308,8 @@ fn compute_copy_classes(ssa: &mut SsaVisitor, body: &Body<'_>) -> IndexVec<Local
     }
     debug_assert_eq!(copies[RETURN_PLACE], RETURN_PLACE);
 
-    copies
+    ssa.direct_uses = direct_uses;
+    ssa.copy_classes = copies;
 }
 
 #[derive(Debug)]