Resolve aliases before inserting values into the live set (bytecodeal…

…liance#8945) * Refactor the internals of `FunctionBuilder::insert_safepoint_spills` into a few smaller methods * Initialize a logger for the `cranelift-fuzzgen` fuzz target * Resolve aliases before inserting values into the live set This fixes a fuzz bug found in the development of bytecodealliance#8941
ahaoboy · Jul 12, 2024 · 9f66134 · 9f66134
1 parent 68d3d83
commit 9f66134
Show file tree

Hide file tree

Showing 4 changed files with 63 additions and 11 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/cranelift/frontend/src/frontend.rs b/cranelift/frontend/src/frontend.rs
@@ -261,6 +261,8 @@ impl fmt::Display for DefVariableError {
     }
 }
 
+const LOG2_SIZE_CAPACITY: usize = (16u8.ilog2() as usize) + 1;
+
 /// This module allows you to create a function in Cranelift IR in a straightforward way, hiding
 /// all the complexity of its internal representation.
 ///
@@ -686,9 +688,40 @@ impl<'a> FunctionBuilder<'a> {
     /// (currently this is just non-tail calls).
     ///
     /// Second, take those results, add stack slots so we have a place to spill
-    /// to, and then finally spill the live needs-stack-map values at each
-    /// safepoint.
+    /// to.
+    ///
+    /// Third, and finally, we spill the live needs-stack-map values at each
+    /// safepoint into those stack slots.
     fn insert_safepoint_spills(&mut self) {
+        // Find all the GC values that are live across each safepoint.
+        let (safepoints, max_vals_in_stack_map_by_log2_size) =
+            self.find_live_stack_map_values_at_each_safepoint();
+
+        // Create the stack slots to spill them into.
+        let stack_slots = self.create_safepoint_slots(max_vals_in_stack_map_by_log2_size);
+
+        // Insert spills to our new stack slots before each safepoint
+        // instruction.
+        self.insert_safepoint_spills_to_stack_slots(safepoints, stack_slots);
+    }
+
+    /// Find the live GC references for each safepoint instruction in this
+    /// function.
+    ///
+    /// Returns a pair of
+    ///
+    /// 1. A map from each safepoint instruction to the set of GC references
+    ///    that are live across it
+    ///
+    /// 2. The maximum number of values we need to store in a stack map at the
+    ///    same time, bucketed by their type's size. This array is indexed by
+    ///    the log2 of the type's size.
+    fn find_live_stack_map_values_at_each_safepoint(
+        &mut self,
+    ) -> (
+        crate::HashMap<ir::Inst, SmallVec<[ir::Value; 4]>>,
+        [usize; LOG2_SIZE_CAPACITY],
+    ) {
         // A map from each safepoint to the set of GC references that are live
         // across it.
         let mut safepoints: crate::HashMap<ir::Inst, SmallVec<[ir::Value; 4]>> =
@@ -698,7 +731,6 @@ impl<'a> FunctionBuilder<'a> {
         // same time, bucketed by their type's size. This array is indexed by
         // the log2 of the type's size. We do not support recording values whose
         // size is greater than 16 in stack maps.
-        const LOG2_SIZE_CAPACITY: usize = (16u8.ilog2() as usize) + 1;
         let mut max_vals_in_stack_map_by_log2_size = [0; LOG2_SIZE_CAPACITY];
 
         // The set of needs-stack-maps values that are currently live in our
@@ -776,6 +808,7 @@ impl<'a> FunctionBuilder<'a> {
                 // instruction to the live set. This includes branch arguments,
                 // as mentioned above.
                 for val in self.func.dfg.inst_values(inst) {
+                    let val = self.func.dfg.resolve_aliases(val);
                     if self.func_ctx.stack_map_values.contains(val) {
                         live.insert(val);
                     }
@@ -791,17 +824,26 @@ impl<'a> FunctionBuilder<'a> {
             }
         }
 
-        // Create a stack slot for each size of needs-stack-map value. These
-        // slots are arrays capable of holding the maximum number of same-sized
-        // values that must appear in the same stack map at the same time.
-        //
-        // This is indexed by the log2 of the type size.
+        (safepoints, max_vals_in_stack_map_by_log2_size)
+    }
+
+    /// Create a stack slot for each size of needs-stack-map value.
+    ///
+    /// These slots are arrays capable of holding the maximum number of
+    /// same-sized values that must appear in the same stack map at the same
+    /// time.
+    ///
+    /// The resulting array of stack slots is indexed by the log2 of the type
+    /// size.
+    fn create_safepoint_slots(
+        &mut self,
+        max_vals_in_stack_map_by_log2_size: [usize; LOG2_SIZE_CAPACITY],
+    ) -> [PackedOption<ir::StackSlot>; LOG2_SIZE_CAPACITY] {
         let mut stack_slots = [PackedOption::<ir::StackSlot>::default(); LOG2_SIZE_CAPACITY];
         for (log2_size, capacity) in max_vals_in_stack_map_by_log2_size.into_iter().enumerate() {
             if capacity == 0 {
                 continue;
             }
-
             let size = 1usize << log2_size;
             let slot = self.func.create_sized_stack_slot(ir::StackSlotData::new(
                 ir::StackSlotKind::ExplicitSlot,
@@ -810,9 +852,16 @@ impl<'a> FunctionBuilder<'a> {
             ));
             stack_slots[log2_size] = Some(slot).into();
         }
+        stack_slots
+    }
 
-        // Insert spills to our new stack slots before each safepoint
-        // instruction.
+    /// Insert spills to the given stack slots before each safepoint
+    /// instruction.
+    fn insert_safepoint_spills_to_stack_slots(
+        &mut self,
+        safepoints: crate::HashMap<ir::Inst, SmallVec<[ir::Value; 4]>>,
+        stack_slots: [PackedOption<ir::StackSlot>; LOG2_SIZE_CAPACITY],
+    ) {
         let mut cursor = FuncCursor::new(self.func);
         for (inst, live_vals) in safepoints {
             cursor = cursor.at_inst(inst);

diff --git a/fuzz/Cargo.toml b/fuzz/Cargo.toml
@@ -12,6 +12,7 @@ cargo-fuzz = true
 
 [dependencies]
 anyhow = { workspace = true }
+env_logger = { workspace = true }
 once_cell = { workspace = true }
 cranelift-codegen = { workspace = true, features = ["incremental-cache", "x86", "arm64", "s390x", "riscv64"] }
 cranelift-reader = { workspace = true }

diff --git a/fuzz/fuzz_targets/cranelift-fuzzgen.rs b/fuzz/fuzz_targets/cranelift-fuzzgen.rs
@@ -171,6 +171,7 @@ impl fmt::Debug for TestCase {
 
 impl<'a> Arbitrary<'a> for TestCase {
     fn arbitrary(u: &mut Unstructured<'a>) -> arbitrary::Result<Self> {
+        let _ = env_logger::try_init();
         Self::generate(u).map_err(|_| {
             STATISTICS.invalid_inputs.fetch_add(1, Ordering::SeqCst);
             arbitrary::Error::IncorrectFormat