Allocate less while constructing cranelift-fuzzgen tests (#4863)

* Improve panic message if typevar_operand is None

* cranelift-fuzzgen: Don't allocate for each choice

I don't think the performance of test-case generation is at all
important here. I'm actually doing this in preparation for a bigger
refactor where I want to be able to borrow the list of valid choices for
a given opcode without worrying about lifetimes.

* cranelift-fuzzgen: Remove next_func_index

It's only used locally within `generate_funcrefs`, so it doesn't need to
be in the FunctionBuilder struct.

Also there's already a local counter that I think is good enough for
this. As far as I know, the function indexes only need to be distinct,
not contiguous.

* cranelift-fuzzgen: Separate resources from config

The function-global variables, blocks, etc that are generated before
generating instructions are all owned collections without any lifetime
parameters. By contrast, the Unstructured and Config are both borrowed.
Separating them will make it easier to borrow from the owned resources.

cranelift/codegen/src/ir/dfg.rs

@@ -864,7 +864,12 @@ impl DataFlowGraph {
             self.value_type(
                 self[inst]
                     .typevar_operand(&self.value_lists)
-                    .expect("Instruction format doesn't have a designated operand, bad opcode."),
+                    .unwrap_or_else(|| {
+                        panic!(
+                            "Instruction format for {:?} doesn't have a designated operand",
+                            self[inst]
+                        )
+                    }),
             )
         } else {
             self.value_type(self.first_result(inst))

cranelift/fuzzgen/src/function_generator.rs

@@ -72,7 +72,7 @@ fn insert_call(
     _rets: &'static [Type],
 ) -> Result<()> {
     assert_eq!(opcode, Opcode::Call, "only call handled at the moment");
-    let (sig, func_ref) = fgen.u.choose(&fgen.func_refs)?.clone();
+    let (sig, func_ref) = fgen.u.choose(&fgen.resources.func_refs)?.clone();
 
     let actuals = fgen.generate_values_for_signature(
         builder,
@@ -92,7 +92,7 @@ fn insert_stack_load(
 ) -> Result<()> {
     let typevar = rets[0];
     let type_size = typevar.bytes();
-    let (slot, slot_size) = fgen.stack_slot_with_size(builder, type_size)?;
+    let (slot, slot_size) = fgen.stack_slot_with_size(type_size)?;
     let offset = fgen.u.int_in_range(0..=(slot_size - type_size))? as i32;
 
     let val = builder.ins().stack_load(typevar, slot, offset);
@@ -111,7 +111,7 @@ fn insert_stack_store(
 ) -> Result<()> {
     let typevar = args[0];
     let type_size = typevar.bytes();
-    let (slot, slot_size) = fgen.stack_slot_with_size(builder, type_size)?;
+    let (slot, slot_size) = fgen.stack_slot_with_size(type_size)?;
     let offset = fgen.u.int_in_range(0..=(slot_size - type_size))? as i32;
 
     let arg0 = fgen.get_variable_of_type(typevar)?;
@@ -558,12 +558,17 @@ where
 {
     u: &'r mut Unstructured<'data>,
     config: &'r Config,
-    vars: Vec<(Type, Variable)>,
+    resources: Resources,
+}
+
+#[derive(Default)]
+struct Resources {
+    vars: HashMap<Type, Vec<Variable>>,
     blocks: Vec<(Block, BlockSignature)>,
+    blocks_without_params: Vec<Block>,
     jump_tables: Vec<JumpTable>,
     func_refs: Vec<(Signature, FuncRef)>,
-    next_func_index: u32,
+    stack_slots: Vec<(StackSlot, StackSize)>,
-    static_stack_slots: Vec<StackSlot>,
 }
 
 impl<'r, 'data> FunctionGenerator<'r, 'data>
@@ -574,12 +579,7 @@ where
         Self {
             u,
             config,
-            vars: vec![],
+            resources: Resources::default(),
-            blocks: vec![],
-            jump_tables: vec![],
-            func_refs: vec![],
-            static_stack_slots: vec![],
-            next_func_index: 0,
         }
     }
 
@@ -646,19 +646,12 @@ where
     }
 
     /// Finds a stack slot with size of at least n bytes
-    fn stack_slot_with_size(
+    fn stack_slot_with_size(&mut self, n: u32) -> Result<(StackSlot, StackSize)> {
-        &mut self,
+        let first = self
-        builder: &mut FunctionBuilder,
+            .resources
-        n: u32,
+            .stack_slots
-    ) -> Result<(StackSlot, StackSize)> {
+            .partition_point(|&(_slot, size)| size < n);
-        let opts: Vec<_> = self
+        Ok(*self.u.choose(&self.resources.stack_slots[first..])?)
-            .static_stack_slots
-            .iter()
-            .map(|ss| (*ss, builder.func.sized_stack_slots[*ss].size))
-            .filter(|(_, size)| *size >= n)
-            .collect();
-
-        Ok(*self.u.choose(&opts[..])?)
     }
 
     /// Generates an address that should allow for a store or a load.
@@ -679,7 +672,7 @@ where
         // TODO: Currently our only source of addresses is stack_addr, but we should
         // add heap_addr, global_value, symbol_value eventually
         let (addr, available_size) = {
-            let (ss, slot_size) = self.stack_slot_with_size(builder, min_size)?;
+            let (ss, slot_size) = self.stack_slot_with_size(min_size)?;
             let max_offset = slot_size.saturating_sub(min_size);
             let offset = self.u.int_in_range(0..=max_offset)? as i32;
             let base_addr = builder.ins().stack_addr(I64, ss, offset);
@@ -697,27 +690,10 @@ where
         Ok((addr, offset.into()))
     }
 
-    /// Creates a new var
-    fn create_var(&mut self, builder: &mut FunctionBuilder, ty: Type) -> Result<Variable> {
-        let id = self.vars.len();
-        let var = Variable::new(id);
-        builder.declare_var(var, ty);
-        self.vars.push((ty, var));
-        Ok(var)
-    }
-
-    fn vars_of_type(&self, ty: Type) -> Vec<Variable> {
-        self.vars
-            .iter()
-            .filter(|(var_ty, _)| *var_ty == ty)
-            .map(|(_, v)| *v)
-            .collect()
-    }
-
     /// Get a variable of type `ty` from the current function
     fn get_variable_of_type(&mut self, ty: Type) -> Result<Variable> {
-        let opts = self.vars_of_type(ty);
+        let opts = self.resources.vars.get(&ty).map_or(&[][..], Vec::as_slice);
-        let var = self.u.choose(&opts[..])?;
+        let var = self.u.choose(opts)?;
         Ok(*var)
     }
 
@@ -761,22 +737,12 @@ where
         &mut self,
         builder: &mut FunctionBuilder,
     ) -> Result<(Block, Vec<Value>)> {
-        let block_targets = &self.blocks[1..];
+        let block_targets = &self.resources.blocks[1..];
         let (block, signature) = self.u.choose(block_targets)?.clone();
         let args = self.generate_values_for_signature(builder, signature.into_iter())?;
         Ok((block, args))
     }
 
-    /// Valid blocks for jump tables have to have no parameters in the signature, and must also
-    /// not be the first block.
-    fn generate_valid_jumptable_target_blocks(&mut self) -> Vec<Block> {
-        self.blocks[1..]
-            .iter()
-            .filter(|(_, sig)| sig.len() == 0)
-            .map(|(b, _)| *b)
-            .collect()
-    }
-
     fn generate_values_for_signature<I: Iterator<Item = Type>>(
         &mut self,
         builder: &mut FunctionBuilder,
@@ -813,10 +779,9 @@ where
         let var = self.get_variable_of_type(I32)?; // br_table only supports I32
         let val = builder.use_var(var);
 
-        let valid_blocks = self.generate_valid_jumptable_target_blocks();
+        let default_block = *self.u.choose(&self.resources.blocks_without_params)?;
-        let default_block = *self.u.choose(&valid_blocks[..])?;
 
-        let jt = *self.u.choose(&self.jump_tables[..])?;
+        let jt = *self.u.choose(&self.resources.jump_tables)?;
         builder.ins().br_table(val, default_block, jt);
         Ok(())
     }
@@ -872,8 +837,7 @@ where
         let switch_var = self.get_variable_of_type(_type)?;
         let switch_val = builder.use_var(switch_var);
 
-        let valid_blocks = self.generate_valid_jumptable_target_blocks();
+        let default_block = *self.u.choose(&self.resources.blocks_without_params)?;
-        let default_block = *self.u.choose(&valid_blocks[..])?;
 
         // Build this into a HashMap since we cannot have duplicate entries.
         let mut entries = HashMap::new();
@@ -894,7 +858,7 @@ where
             // Build the switch entries
             for i in 0..range_size {
                 let index = range_start.wrapping_add(i) % ty_max;
-                let block = *self.u.choose(&valid_blocks[..])?;
+                let block = *self.u.choose(&self.resources.blocks_without_params)?;
                 entries.insert(index, block);
             }
         }
@@ -936,26 +900,25 @@ where
     }
 
     fn generate_jumptables(&mut self, builder: &mut FunctionBuilder) -> Result<()> {
-        let valid_blocks = self.generate_valid_jumptable_target_blocks();
-
         for _ in 0..self.param(&self.config.jump_tables_per_function)? {
             let mut jt_data = JumpTableData::new();
 
             for _ in 0..self.param(&self.config.jump_table_entries)? {
-                let block = *self.u.choose(&valid_blocks[..])?;
+                let block = *self.u.choose(&self.resources.blocks_without_params)?;
                 jt_data.push_entry(block);
             }
 
-            self.jump_tables.push(builder.create_jump_table(jt_data));
+            self.resources
+                .jump_tables
+                .push(builder.create_jump_table(jt_data));
         }
         Ok(())
     }
 
     fn generate_funcrefs(&mut self, builder: &mut FunctionBuilder) -> Result<()> {
-        for _ in 0..self.param(&self.config.funcrefs_per_function)? {
+        let count = self.param(&self.config.funcrefs_per_function)?;
+        for func_index in 0..count.try_into().unwrap() {
             let (ext_name, sig) = if self.u.arbitrary::<bool>()? {
-                let func_index = self.next_func_index;
-                self.next_func_index = self.next_func_index.wrapping_add(1);
                 let user_func_ref = builder
                     .func
                     .declare_imported_user_function(UserExternalName {
@@ -980,7 +943,7 @@ where
                 colocated: self.u.arbitrary()?,
             });
 
-            self.func_refs.push((sig, func_ref));
+            self.resources.func_refs.push((sig, func_ref));
         }
 
         Ok(())
@@ -991,9 +954,13 @@ where
             let bytes = self.param(&self.config.static_stack_slot_size)? as u32;
             let ss_data = StackSlotData::new(StackSlotKind::ExplicitSlot, bytes);
             let slot = builder.create_sized_stack_slot(ss_data);
-
+            self.resources.stack_slots.push((slot, bytes));
-            self.static_stack_slots.push(slot);
         }
+
+        self.resources
+            .stack_slots
+            .sort_unstable_by_key(|&(_slot, bytes)| bytes);
+
         Ok(())
     }
 
@@ -1005,8 +972,7 @@ where
         let i16_zero = builder.ins().iconst(I16, 0);
         let i8_zero = builder.ins().iconst(I8, 0);
 
-        for &slot in self.static_stack_slots.iter() {
+        for &(slot, init_size) in self.resources.stack_slots.iter() {
-            let init_size = builder.func.sized_stack_slots[slot].size;
             let mut size = init_size;
 
             // Insert the largest available store for the remaining size.
@@ -1038,7 +1004,7 @@ where
         // the entry block.
         let block_count = 1 + extra_block_count;
 
-        let blocks = (0..block_count)
+        (0..block_count)
             .map(|i| {
                 let is_entry = i == 0;
                 let block = builder.create_block();
@@ -1063,9 +1029,7 @@ where
                     Ok((block, sig))
                 }
             })
-            .collect::<Result<Vec<_>>>()?;
+            .collect()
-
-        Ok(blocks)
     }
 
     fn generate_block_signature(&mut self) -> Result<BlockSignature> {
@@ -1080,21 +1044,33 @@ where
 
     fn build_variable_pool(&mut self, builder: &mut FunctionBuilder) -> Result<()> {
         let block = builder.current_block().unwrap();
-        let func_params = builder.func.signature.params.clone();
 
         // Define variables for the function signature
-        for (i, param) in func_params.iter().enumerate() {
+        let mut vars: Vec<_> = builder
-            let var = self.create_var(builder, param.value_type)?;
+            .func
-            let block_param = builder.block_params(block)[i];
+            .signature
-            builder.def_var(var, block_param);
+            .params
-        }
+            .iter()
+            .map(|param| param.value_type)
+            .zip(builder.block_params(block).iter().copied())
+            .collect();
 
         // Create a pool of vars that are going to be used in this function
         for _ in 0..self.param(&self.config.vars_per_function)? {
             let ty = self.generate_type()?;
-            let var = self.create_var(builder, ty)?;
             let value = self.generate_const(builder, ty)?;
+            vars.push((ty, value));
+        }
+
+        for (id, (ty, value)) in vars.into_iter().enumerate() {
+            let var = Variable::new(id);
+            builder.declare_var(var, ty);
             builder.def_var(var, value);
+            self.resources
+                .vars
+                .entry(ty)
+                .or_insert_with(Vec::new)
+                .push(var);
         }
 
         Ok(())
@@ -1117,7 +1093,15 @@ where
 
         let mut builder = FunctionBuilder::new(&mut func, &mut fn_builder_ctx);
 
-        self.blocks = self.generate_blocks(&mut builder, &sig)?;
+        self.resources.blocks = self.generate_blocks(&mut builder, &sig)?;
+
+        // Valid blocks for jump tables have to have no parameters in the signature, and must also
+        // not be the first block.
+        self.resources.blocks_without_params = self.resources.blocks[1..]
+            .iter()
+            .filter(|(_, sig)| sig.len() == 0)
+            .map(|(b, _)| *b)
+            .collect();
 
         // Function preamble
         self.generate_jumptables(&mut builder)?;
@@ -1125,7 +1109,7 @@ where
         self.generate_stack_slots(&mut builder)?;
 
         // Main instruction generation loop
-        for (i, (block, block_sig)) in self.blocks.clone().iter().enumerate() {
+        for (i, (block, block_sig)) in self.resources.blocks.clone().iter().enumerate() {
             let is_block0 = i == 0;
             builder.switch_to_block(*block);