From c2a2f008208da646be44bec9c2494188f2c027f3 Mon Sep 17 00:00:00 2001
From: Jamey Sharp <jsharp@fastly.com>
Date: Sun, 24 Mar 2024 16:13:54 -0700
Subject: [PATCH] cranelift: New InstructionData::map_values helper

This is like the DataFlowGraph::map_inst_values method, but that method
mutably borrows `self`, so it goes to a fair bit of trouble to drop its
borrows around callbacks so the mutable borrow can be used in the
callback too.

This new helper only actually needs to borrow two public fields from the
DFG: `value_lists` and `jump_tables`. Therefore it's possible to use
other fields in the callback as long as the compiler can see all the
fields being used in the same scope. That's good enough for everywhere
we were using this pattern, so we can simplify this way.

The case which motivated this change isn't shown here: It isn't possible
to call `dfg.map_inst_values` on every instruction in `dfg.insts`,
because the borrow on `dfg.insts` prevents taking a mutable borrow on
`dfg`. But we can call this new helper in that case.
---
 cranelift/codegen/src/egraph.rs          |  2 +-
 cranelift/codegen/src/ir/dfg.rs          | 36 ++++++------------------
 cranelift/codegen/src/ir/instructions.rs | 19 +++++++++++++
 3 files changed, 29 insertions(+), 28 deletions(-)

diff --git a/cranelift/codegen/src/egraph.rs b/cranelift/codegen/src/egraph.rs
index b1be475ace3d..73f98a9b949c 100644
--- a/cranelift/codegen/src/egraph.rs
+++ b/cranelift/codegen/src/egraph.rs
@@ -662,7 +662,7 @@ impl<'a> EgraphPass<'a> {
                         // Rewrite args of *all* instructions using the
                         // value-to-opt-value map.
                         cursor.func.dfg.resolve_aliases_in_arguments(inst);
-                        cursor.func.dfg.map_inst_values(inst, |_, arg| {
+                        cursor.func.dfg.map_inst_values(inst, |arg| {
                             let new_value = value_to_opt_value[arg];
                             trace!("rewriting arg {} of inst {} to {}", arg, inst, new_value);
                             debug_assert_ne!(new_value, Value::reserved_value());
diff --git a/cranelift/codegen/src/ir/dfg.rs b/cranelift/codegen/src/ir/dfg.rs
index 3204f312277e..b7c8e161a602 100644
--- a/cranelift/codegen/src/ir/dfg.rs
+++ b/cranelift/codegen/src/ir/dfg.rs
@@ -383,7 +383,9 @@ impl DataFlowGraph {
     /// For each argument of inst which is defined by an alias, replace the
     /// alias with the aliased value.
     pub fn resolve_aliases_in_arguments(&mut self, inst: Inst) {
-        self.map_inst_values(inst, |dfg, arg| resolve_aliases(&dfg.values, arg));
+        self.insts[inst].map_values(&mut self.value_lists, &mut self.jump_tables, |arg| {
+            resolve_aliases(&self.values, arg)
+        });
     }
 
     /// Turn a value into an alias of another.
@@ -732,24 +734,11 @@ impl DataFlowGraph {
     }
 
     /// Map a function over the values of the instruction.
-    pub fn map_inst_values<F>(&mut self, inst: Inst, mut body: F)
+    pub fn map_inst_values<F>(&mut self, inst: Inst, body: F)
     where
-        F: FnMut(&mut DataFlowGraph, Value) -> Value,
+        F: FnMut(Value) -> Value,
     {
-        for i in 0..self.inst_args(inst).len() {
-            let arg = self.inst_args(inst)[i];
-            self.inst_args_mut(inst)[i] = body(self, arg);
-        }
-
-        for block_ix in 0..self.insts[inst].branch_destination(&self.jump_tables).len() {
-            // We aren't changing the size of the args list, so we won't need to write the branch
-            // back to the instruction.
-            let mut block = self.insts[inst].branch_destination(&self.jump_tables)[block_ix];
-            for i in 0..block.args_slice(&self.value_lists).len() {
-                let arg = block.args_slice(&self.value_lists)[i];
-                block.args_slice_mut(&mut self.value_lists)[i] = body(self, arg);
-            }
-        }
+        self.insts[inst].map_values(&mut self.value_lists, &mut self.jump_tables, body);
     }
 
     /// Overwrite the instruction's value references with values from the iterator.
@@ -759,16 +748,9 @@ impl DataFlowGraph {
     where
         I: Iterator<Item = Value>,
     {
-        for arg in self.inst_args_mut(inst) {
-            *arg = values.next().unwrap();
-        }
-
-        for block_ix in 0..self.insts[inst].branch_destination(&self.jump_tables).len() {
-            let mut block = self.insts[inst].branch_destination(&self.jump_tables)[block_ix];
-            for arg in block.args_slice_mut(&mut self.value_lists) {
-                *arg = values.next().unwrap();
-            }
-        }
+        self.insts[inst].map_values(&mut self.value_lists, &mut self.jump_tables, |_| {
+            values.next().unwrap()
+        });
     }
 
     /// Get all value arguments on `inst` as a slice.
diff --git a/cranelift/codegen/src/ir/instructions.rs b/cranelift/codegen/src/ir/instructions.rs
index 35607e73ea1c..eebe2c0e3591 100644
--- a/cranelift/codegen/src/ir/instructions.rs
+++ b/cranelift/codegen/src/ir/instructions.rs
@@ -341,6 +341,25 @@ impl InstructionData {
         }
     }
 
+    /// Replace the values used in this instruction according to the given
+    /// function.
+    pub fn map_values(
+        &mut self,
+        pool: &mut ValueListPool,
+        jump_tables: &mut ir::JumpTables,
+        mut f: impl FnMut(Value) -> Value,
+    ) {
+        for arg in self.arguments_mut(pool) {
+            *arg = f(*arg);
+        }
+
+        for block in self.branch_destination_mut(jump_tables) {
+            for arg in block.args_slice_mut(pool) {
+                *arg = f(*arg);
+            }
+        }
+    }
+
     /// If this is a trapping instruction, get its trap code. Otherwise, return
     /// `None`.
     pub fn trap_code(&self) -> Option<TrapCode> {