cranelift: avoid per-variable value sets when tracking stack-map bindings

cranelift/frontend/src/frontend.rs

@@ -31,8 +31,6 @@ pub struct FunctionBuilderContext {
     ssa: SSABuilder,
     status: SecondaryMap<Block, BlockStatus>,
     variables: PrimaryMap<Variable, Type>,
-    stack_map_vars: EntitySet<Variable>,
-    stack_map_values: EntitySet<Value>,
     safepoints: safepoints::SafepointSpiller,
 }
 
@@ -72,15 +70,11 @@ impl FunctionBuilderContext {
             ssa,
             status,
             variables,
-            stack_map_vars,
-            stack_map_values,
             safepoints,
         } = self;
         ssa.clear();
         status.clear();
         variables.clear();
-        stack_map_values.clear();
-        stack_map_vars.clear();
         safepoints.clear();
     }
 
@@ -430,18 +424,19 @@ impl<'a> FunctionBuilder<'a> {
     /// the stack and it being reloading again, the stack can be updated to
     /// facilitate moving GCs.
     ///
-    /// This does not affect any pre-existing uses of the variable.
+    /// This must be called before any definition of the variable.
     ///
     /// # Panics
     ///
     /// Panics if the variable's type is larger than 16 bytes or if this
-    /// variable has not been declared yet.
+    /// variable has not been declared yet. In debug builds, also panics if
+    /// the variable has already been defined.
     pub fn declare_var_needs_stack_map(&mut self, var: Variable) {
         log::trace!("declare_var_needs_stack_map({var:?})");
         let ty = self.func_ctx.variables[var];
         assert!(ty != types::INVALID);
         assert!(ty.bytes() <= 16);
-        self.func_ctx.stack_map_vars.insert(var);
+        self.func_ctx.ssa.mark_var_needs_stack_map(var);
     }
 
     /// Returns the Cranelift IR necessary to use a previously defined user
@@ -559,7 +554,7 @@ impl<'a> FunctionBuilder<'a> {
         assert!(size <= 16);
         assert!(size.is_power_of_two());
 
-        self.func_ctx.stack_map_values.insert(val);
+        self.func_ctx.ssa.stack_map_values_mut().insert(val);
     }
 
     /// Creates a jump table in the function, to be used by [`br_table`](InstBuilder::br_table) instructions.
@@ -709,23 +704,13 @@ impl<'a> FunctionBuilder<'a> {
             }
         }
 
-        // Propagate the needs-stack-map bit from variables to each of their
-        // associated values.
-        for var in self.func_ctx.stack_map_vars.iter() {
-            for val in self.func_ctx.ssa.values_for_var(var) {
-                log::trace!("propagating needs-stack-map from {var:?} to {val:?}");
-                debug_assert_eq!(self.func.dfg.value_type(val), self.func_ctx.variables[var]);
-                self.func_ctx.stack_map_values.insert(val);
-            }
-        }
-
         // If we have any values that need inclusion in stack maps, then we need
         // to run our pass to spill those values to the stack at safepoints and
         // generate stack maps.
-        if !self.func_ctx.stack_map_values.is_empty() {
+        if !self.func_ctx.ssa.stack_map_values().is_empty() {
             self.func_ctx
                 .safepoints
-                .run(&mut self.func, &self.func_ctx.stack_map_values);
+                .run(&mut self.func, self.func_ctx.ssa.stack_map_values());
         }
 
         // Clear the state (but preserve the allocated buffers) in preparation

cranelift/frontend/src/ssa.rs

@@ -37,11 +37,11 @@ pub struct SSABuilder {
     /// the variable in the block.
     variables: SecondaryMap<Variable, SecondaryMap<Block, PackedOption<Value>>>,
 
-    /// Records every SSA `Value` ever associated with a variable, including
-    /// values that `variables` has since overwritten within the same block.
-    /// Used by `values_for_var` to feed `FunctionBuilder::finalize`'s
-    /// stack-map propagation.
-    all_var_values: SecondaryMap<Variable, EntitySet<Value>>,
+    /// Variables whose bindings must be tracked for stack maps.
+    stack_map_vars: EntitySet<Variable>,
+
+    /// SSA values that must be included in stack maps.
+    stack_map_values: EntitySet<Value>,
 
     /// Records the position of the basic blocks and the list of values used but not defined in the
     /// block.
@@ -116,7 +116,8 @@ impl SSABuilder {
     /// deallocating memory.
     pub fn clear(&mut self) {
         self.variables.clear();
-        self.all_var_values.clear();
+        self.stack_map_vars.clear();
+        self.stack_map_values.clear();
         self.ssa_blocks.clear();
         self.variable_pool.clear();
         self.inst_pool.clear();
@@ -128,7 +129,8 @@ impl SSABuilder {
     /// Tests whether an `SSABuilder` is in a cleared state.
     pub fn is_empty(&self) -> bool {
         self.variables.is_empty()
-            && self.all_var_values.is_empty()
+            && self.stack_map_vars.is_empty()
+            && self.stack_map_values.is_empty()
             && self.ssa_blocks.is_empty()
             && self.calls.is_empty()
             && self.results.is_empty()
@@ -209,18 +211,41 @@ fn emit_zero(ty: Type, mut cur: FuncCursor) -> Value {
 /// Phi functions.
 ///
 impl SSABuilder {
-    /// Get all of the values associated with the given variable that we have
-    /// inserted in the function thus far.
-    pub fn values_for_var(&self, var: Variable) -> impl Iterator<Item = Value> + '_ {
-        self.all_var_values[var].iter()
-    }
-
     /// Declares a new definition of a variable in a given basic block.
     /// The SSA value is passed as an argument because it should be created with
     /// `ir::DataFlowGraph::append_result`.
     pub fn def_var(&mut self, var: Variable, val: Value, block: Block) {
         self.variables[var][block] = PackedOption::from(val);
-        self.all_var_values[var].insert(val);
+        self.record_stack_map_binding(var, val);
+    }
+
+    /// Mark `var` as needing its bindings tracked for stack maps.
+    pub fn mark_var_needs_stack_map(&mut self, var: Variable) {
+        debug_assert!(
+            self.variables[var].values().all(|v| v.is_none()),
+            "declare_var_needs_stack_map({var:?}) must be called before any \
+             definition of the variable; otherwise earlier definitions are \
+             silently omitted from stack maps"
+        );
+        self.stack_map_vars.insert(var);
+    }
+
+    /// SSA values that must be included in stack maps.
+    pub fn stack_map_values(&self) -> &EntitySet<Value> {
+        &self.stack_map_values
+    }
+
+    /// Mutable view of [`Self::stack_map_values`].
+    pub fn stack_map_values_mut(&mut self) -> &mut EntitySet<Value> {
+        &mut self.stack_map_values
+    }
+
+    /// Propagate the needs-stack-map bit from `var` to `val`.
+    fn record_stack_map_binding(&mut self, var: Variable, val: Value) {
+        if self.stack_map_vars.contains(var) {
+            log::trace!("recording stack-map binding {var:?} -> {val:?}");
+            self.stack_map_values.insert(val);
+        }
     }
 
     /// Declares a use of a variable in a given basic block. Returns the SSA value corresponding
@@ -283,7 +308,7 @@ impl SSABuilder {
         // any of the blocks before `from`.
         //
         // So in either case there is no definition in these blocks yet and we can blindly set one.
-        debug_assert!(self.all_var_values[var].contains(val));
+        debug_assert!(!self.stack_map_vars.contains(var) || self.stack_map_values.contains(val));
         let var_defs = &mut self.variables[var];
         while block != from {
             debug_assert!(var_defs[block].is_none());
@@ -340,7 +365,7 @@ impl SSABuilder {
         // find a usable definition. So create one.
         let val = func.dfg.append_block_param(block, ty);
         var_defs[block] = PackedOption::from(val);
-        self.all_var_values[var].insert(val);
+        self.record_stack_map_binding(var, val);
 
         // Now every predecessor needs to pass its definition of this variable to the newly added
         // block parameter. To do that we have to "recursively" call `use_var`, but there are two

cranelift/fuzzgen/src/function_generator.rs

@@ -1920,7 +1920,6 @@ where
 
         for (ty, value) in vars.into_iter() {
             let var = builder.declare_var(ty);
-            builder.def_var(var, value);
 
             // Randomly declare variables as needing a stack map.
             // We limit these to only types that have fewer than 16 bytes
@@ -1929,6 +1928,8 @@ where
                 builder.declare_var_needs_stack_map(var);
             }
 
+            builder.def_var(var, value);
+
             self.resources
                 .vars
                 .entry(ty)