1 use crate::config::Config;
2 use anyhow::Result;
3 use arbitrary::{Arbitrary, Unstructured};
4 use cranelift::codegen::ir::types::*;
5 use cranelift::codegen::ir::{
6     AbiParam, Block, ExternalName, Function, JumpTable, Opcode, Signature, Type, Value,
7 };
8 use cranelift::codegen::isa::CallConv;
9 use cranelift::frontend::{FunctionBuilder, FunctionBuilderContext, Variable};
10 use cranelift::prelude::{EntityRef, InstBuilder, IntCC, JumpTableData};
11 use std::ops::RangeInclusive;
12 
13 type BlockSignature = Vec<Type>;
14 
15 fn insert_opcode_arity_0(
16     _fgen: &mut FunctionGenerator,
17     builder: &mut FunctionBuilder,
18     opcode: Opcode,
19     _args: &'static [Type],
20     _rets: &'static [Type],
21 ) -> Result<()> {
22     builder.ins().NullAry(opcode, INVALID);
23     Ok(())
24 }
25 
26 fn insert_opcode_arity_2(
27     fgen: &mut FunctionGenerator,
28     builder: &mut FunctionBuilder,
29     opcode: Opcode,
30     args: &'static [Type],
31     rets: &'static [Type],
32 ) -> Result<()> {
33     let arg0 = fgen.get_variable_of_type(args[0])?;
34     let arg0 = builder.use_var(arg0);
35 
36     let arg1 = fgen.get_variable_of_type(args[1])?;
37     let arg1 = builder.use_var(arg1);
38 
39     let typevar = rets[0];
40     let (inst, dfg) = builder.ins().Binary(opcode, typevar, arg0, arg1);
41     let results = dfg.inst_results(inst).to_vec();
42 
43     for (val, ty) in results.into_iter().zip(rets) {
44         let var = fgen.get_variable_of_type(*ty)?;
45         builder.def_var(var, val);
46     }
47     Ok(())
48 }
49 
50 type OpcodeInserter = fn(
51     fgen: &mut FunctionGenerator,
52     builder: &mut FunctionBuilder,
53     Opcode,
54     &'static [Type],
55     &'static [Type],
56 ) -> Result<()>;
57 
58 // TODO: Derive this from the `cranelift-meta` generator.
59 const OPCODE_SIGNATURES: &'static [(
60     Opcode,
61     &'static [Type], // Args
62     &'static [Type], // Rets
63     OpcodeInserter,
64 )] = &[
65     (Opcode::Nop, &[], &[], insert_opcode_arity_0),
66     // Iadd
67     (Opcode::Iadd, &[I8, I8], &[I8], insert_opcode_arity_2),
68     (Opcode::Iadd, &[I16, I16], &[I16], insert_opcode_arity_2),
69     (Opcode::Iadd, &[I32, I32], &[I32], insert_opcode_arity_2),
70     (Opcode::Iadd, &[I64, I64], &[I64], insert_opcode_arity_2),
71     // Isub
72     (Opcode::Isub, &[I8, I8], &[I8], insert_opcode_arity_2),
73     (Opcode::Isub, &[I16, I16], &[I16], insert_opcode_arity_2),
74     (Opcode::Isub, &[I32, I32], &[I32], insert_opcode_arity_2),
75     (Opcode::Isub, &[I64, I64], &[I64], insert_opcode_arity_2),
76     // Imul
77     (Opcode::Imul, &[I8, I8], &[I8], insert_opcode_arity_2),
78     (Opcode::Imul, &[I16, I16], &[I16], insert_opcode_arity_2),
79     (Opcode::Imul, &[I32, I32], &[I32], insert_opcode_arity_2),
80     (Opcode::Imul, &[I64, I64], &[I64], insert_opcode_arity_2),
81     // Udiv
82     (Opcode::Udiv, &[I8, I8], &[I8], insert_opcode_arity_2),
83     (Opcode::Udiv, &[I16, I16], &[I16], insert_opcode_arity_2),
84     (Opcode::Udiv, &[I32, I32], &[I32], insert_opcode_arity_2),
85     (Opcode::Udiv, &[I64, I64], &[I64], insert_opcode_arity_2),
86     // Sdiv
87     (Opcode::Sdiv, &[I8, I8], &[I8], insert_opcode_arity_2),
88     (Opcode::Sdiv, &[I16, I16], &[I16], insert_opcode_arity_2),
89     (Opcode::Sdiv, &[I32, I32], &[I32], insert_opcode_arity_2),
90     (Opcode::Sdiv, &[I64, I64], &[I64], insert_opcode_arity_2),
91 ];
92 
93 pub struct FunctionGenerator<'r, 'data>
94 where
95     'data: 'r,
96 {
97     u: &'r mut Unstructured<'data>,
98     config: &'r Config,
99     vars: Vec<(Type, Variable)>,
100     blocks: Vec<(Block, BlockSignature)>,
101     jump_tables: Vec<JumpTable>,
102 }
103 
104 impl<'r, 'data> FunctionGenerator<'r, 'data>
105 where
106     'data: 'r,
107 {
108     pub fn new(u: &'r mut Unstructured<'data>, config: &'r Config) -> Self {
109         Self {
110             u,
111             config,
112             vars: vec![],
113             blocks: vec![],
114             jump_tables: vec![],
115         }
116     }
117 
118     /// Generates a random value for config `param`
119     fn param(&mut self, param: &RangeInclusive<usize>) -> Result<usize> {
120         Ok(self.u.int_in_range(param.clone())?)
121     }
122 
123     fn generate_callconv(&mut self) -> Result<CallConv> {
124         // TODO: Generate random CallConvs per target
125         Ok(CallConv::SystemV)
126     }
127 
128     fn generate_intcc(&mut self) -> Result<IntCC> {
129         Ok(*self.u.choose(
130             &[
131                 IntCC::Equal,
132                 IntCC::NotEqual,
133                 IntCC::SignedLessThan,
134                 IntCC::SignedGreaterThanOrEqual,
135                 IntCC::SignedGreaterThan,
136                 IntCC::SignedLessThanOrEqual,
137                 IntCC::UnsignedLessThan,
138                 IntCC::UnsignedGreaterThanOrEqual,
139                 IntCC::UnsignedGreaterThan,
140                 IntCC::UnsignedLessThanOrEqual,
141                 IntCC::Overflow,
142                 IntCC::NotOverflow,
143             ][..],
144         )?)
145     }
146 
147     fn generate_type(&mut self) -> Result<Type> {
148         // TODO: It would be nice if we could get these directly from cranelift
149         let scalars = [
150             // IFLAGS, FFLAGS,
151             B1, // B8, B16, B32, B64, B128,
152             I8, I16, I32, I64,
153             // I128,
154             // F32, F64,
155             // R32, R64,
156         ];
157         // TODO: vector types
158 
159         let ty = self.u.choose(&scalars[..])?;
160         Ok(*ty)
161     }
162 
163     fn generate_abi_param(&mut self) -> Result<AbiParam> {
164         // TODO: Generate more advanced abi params (structs/purposes/extensions/etc...)
165         let ty = self.generate_type()?;
166         Ok(AbiParam::new(ty))
167     }
168 
169     fn generate_signature(&mut self) -> Result<Signature> {
170         let callconv = self.generate_callconv()?;
171         let mut sig = Signature::new(callconv);
172 
173         for _ in 0..self.param(&self.config.signature_params)? {
174             sig.params.push(self.generate_abi_param()?);
175         }
176 
177         for _ in 0..self.param(&self.config.signature_rets)? {
178             sig.returns.push(self.generate_abi_param()?);
179         }
180 
181         Ok(sig)
182     }
183 
184     /// Creates a new var
185     fn create_var(&mut self, builder: &mut FunctionBuilder, ty: Type) -> Result<Variable> {
186         let id = self.vars.len();
187         let var = Variable::new(id);
188         builder.declare_var(var, ty);
189         self.vars.push((ty, var));
190         Ok(var)
191     }
192 
193     fn vars_of_type(&self, ty: Type) -> Vec<Variable> {
194         self.vars
195             .iter()
196             .filter(|(var_ty, _)| *var_ty == ty)
197             .map(|(_, v)| *v)
198             .collect()
199     }
200 
201     /// Get a variable of type `ty` from the current function
202     fn get_variable_of_type(&mut self, ty: Type) -> Result<Variable> {
203         let opts = self.vars_of_type(ty);
204         let var = self.u.choose(&opts[..])?;
205         Ok(*var)
206     }
207 
208     /// Generates an instruction(`iconst`/`fconst`/etc...) to introduce a constant value
209     fn generate_const(&mut self, builder: &mut FunctionBuilder, ty: Type) -> Result<Value> {
210         Ok(match ty {
211             ty if ty.is_int() => {
212                 let imm64 = match ty {
213                     I8 => self.u.arbitrary::<i8>()? as i64,
214                     I16 => self.u.arbitrary::<i16>()? as i64,
215                     I32 => self.u.arbitrary::<i32>()? as i64,
216                     I64 => self.u.arbitrary::<i64>()?,
217                     _ => unreachable!(),
218                 };
219                 builder.ins().iconst(ty, imm64)
220             }
221             ty if ty.is_bool() => builder.ins().bconst(B1, bool::arbitrary(self.u)?),
222             _ => unimplemented!(),
223         })
224     }
225 
226     /// Chooses a random block which can be targeted by a jump / branch.
227     /// This means any block that is not the first block.
228     ///
229     /// For convenience we also generate values that match the block's signature
230     fn generate_target_block(
231         &mut self,
232         builder: &mut FunctionBuilder,
233     ) -> Result<(Block, Vec<Value>)> {
234         let block_targets = &self.blocks[1..];
235         let (block, signature) = self.u.choose(block_targets)?.clone();
236         let args = self.generate_values_for_signature(builder, signature.into_iter())?;
237         Ok((block, args))
238     }
239 
240     /// Valid blocks for jump tables have to have no parameters in the signature, and must also
241     /// not be the first block.
242     fn generate_valid_jumptable_target_blocks(&mut self) -> Vec<Block> {
243         self.blocks[1..]
244             .iter()
245             .filter(|(_, sig)| sig.len() == 0)
246             .map(|(b, _)| *b)
247             .collect()
248     }
249 
250     fn generate_values_for_signature<I: Iterator<Item = Type>>(
251         &mut self,
252         builder: &mut FunctionBuilder,
253         signature: I,
254     ) -> Result<Vec<Value>> {
255         signature
256             .map(|ty| {
257                 let var = self.get_variable_of_type(ty)?;
258                 let val = builder.use_var(var);
259                 Ok(val)
260             })
261             .collect()
262     }
263 
264     fn generate_return(&mut self, builder: &mut FunctionBuilder) -> Result<()> {
265         let types: Vec<Type> = {
266             let rets = &builder.func.signature.returns;
267             rets.iter().map(|p| p.value_type).collect()
268         };
269         let vals = self.generate_values_for_signature(builder, types.into_iter())?;
270 
271         builder.ins().return_(&vals[..]);
272         Ok(())
273     }
274 
275     fn generate_jump(&mut self, builder: &mut FunctionBuilder) -> Result<()> {
276         let (block, args) = self.generate_target_block(builder)?;
277         builder.ins().jump(block, &args[..]);
278         Ok(())
279     }
280 
281     /// Generates a br_table into a random block
282     fn generate_br_table(&mut self, builder: &mut FunctionBuilder) -> Result<()> {
283         let _type = *self.u.choose(&[I8, I16, I32, I64][..])?;
284         let var = self.get_variable_of_type(_type)?;
285         let val = builder.use_var(var);
286 
287         let valid_blocks = self.generate_valid_jumptable_target_blocks();
288         let default_block = *self.u.choose(&valid_blocks[..])?;
289 
290         let jt = *self.u.choose(&self.jump_tables[..])?;
291         builder.ins().br_table(val, default_block, jt);
292         Ok(())
293     }
294 
295     /// Generates a brz/brnz into a random block
296     fn generate_br(&mut self, builder: &mut FunctionBuilder) -> Result<()> {
297         let (block, args) = self.generate_target_block(builder)?;
298 
299         let condbr_types = [
300             I8, I16, I32, I64, // TODO: I128
301             B1,
302         ];
303         let _type = *self.u.choose(&condbr_types[..])?;
304         let var = self.get_variable_of_type(_type)?;
305         let val = builder.use_var(var);
306 
307         if bool::arbitrary(self.u)? {
308             builder.ins().brz(val, block, &args[..]);
309         } else {
310             builder.ins().brnz(val, block, &args[..]);
311         }
312 
313         // After brz/brnz we must generate a jump
314         self.generate_jump(builder)?;
315         Ok(())
316     }
317 
318     fn generate_bricmp(&mut self, builder: &mut FunctionBuilder) -> Result<()> {
319         let (block, args) = self.generate_target_block(builder)?;
320         let cond = self.generate_intcc()?;
321 
322         let bricmp_types = [
323             I8, I16, I32, I64, // TODO: I128
324         ];
325         let _type = *self.u.choose(&bricmp_types[..])?;
326 
327         let lhs_var = self.get_variable_of_type(_type)?;
328         let lhs_val = builder.use_var(lhs_var);
329 
330         let rhs_var = self.get_variable_of_type(_type)?;
331         let rhs_val = builder.use_var(rhs_var);
332 
333         builder
334             .ins()
335             .br_icmp(cond, lhs_val, rhs_val, block, &args[..]);
336 
337         // After bricmp's we must generate a jump
338         self.generate_jump(builder)?;
339         Ok(())
340     }
341 
342     /// We always need to exit safely out of a block.
343     /// This either means a jump into another block or a return.
344     fn finalize_block(&mut self, builder: &mut FunctionBuilder) -> Result<()> {
345         let gen = self.u.choose(
346             &[
347                 Self::generate_bricmp,
348                 Self::generate_br,
349                 Self::generate_br_table,
350                 Self::generate_jump,
351                 Self::generate_return,
352             ][..],
353         )?;
354 
355         gen(self, builder)
356     }
357 
358     /// Fills the current block with random instructions
359     fn generate_instructions(&mut self, builder: &mut FunctionBuilder) -> Result<()> {
360         for _ in 0..self.param(&self.config.instructions_per_block)? {
361             let (op, args, rets, inserter) = *self.u.choose(OPCODE_SIGNATURES)?;
362             inserter(self, builder, op, args, rets)?;
363         }
364 
365         Ok(())
366     }
367 
368     fn generate_jumptables(&mut self, builder: &mut FunctionBuilder) -> Result<()> {
369         let valid_blocks = self.generate_valid_jumptable_target_blocks();
370 
371         for _ in 0..self.param(&self.config.jump_tables_per_function)? {
372             let mut jt_data = JumpTableData::new();
373 
374             for _ in 0..self.param(&self.config.jump_table_entries)? {
375                 let block = *self.u.choose(&valid_blocks[..])?;
376                 jt_data.push_entry(block);
377             }
378 
379             self.jump_tables.push(builder.create_jump_table(jt_data));
380         }
381         Ok(())
382     }
383 
384     /// Creates a random amount of blocks in this function
385     fn generate_blocks(
386         &mut self,
387         builder: &mut FunctionBuilder,
388         sig: &Signature,
389     ) -> Result<Vec<(Block, BlockSignature)>> {
390         let extra_block_count = self.param(&self.config.blocks_per_function)?;
391 
392         // We must always have at least one block, so we generate the "extra" blocks and add 1 for
393         // the entry block.
394         let block_count = 1 + extra_block_count;
395 
396         let blocks = (0..block_count)
397             .map(|i| {
398                 let block = builder.create_block();
399 
400                 // The first block has to have the function signature, but for the rest of them we generate
401                 // a random signature;
402                 if i == 0 {
403                     builder.append_block_params_for_function_params(block);
404                     Ok((block, sig.params.iter().map(|a| a.value_type).collect()))
405                 } else {
406                     let sig = self.generate_block_signature()?;
407                     sig.iter().for_each(|ty| {
408                         builder.append_block_param(block, *ty);
409                     });
410                     Ok((block, sig))
411                 }
412             })
413             .collect::<Result<Vec<_>>>()?;
414 
415         Ok(blocks)
416     }
417 
418     fn generate_block_signature(&mut self) -> Result<BlockSignature> {
419         let param_count = self.param(&self.config.block_signature_params)?;
420 
421         let mut params = Vec::with_capacity(param_count);
422         for _ in 0..param_count {
423             params.push(self.generate_type()?);
424         }
425         Ok(params)
426     }
427 
428     fn build_variable_pool(&mut self, builder: &mut FunctionBuilder) -> Result<()> {
429         let block = builder.current_block().unwrap();
430         let func_params = builder.func.signature.params.clone();
431 
432         // Define variables for the function signature
433         for (i, param) in func_params.iter().enumerate() {
434             let var = self.create_var(builder, param.value_type)?;
435             let block_param = builder.block_params(block)[i];
436             builder.def_var(var, block_param);
437         }
438 
439         // Create a pool of vars that are going to be used in this function
440         for _ in 0..self.param(&self.config.vars_per_function)? {
441             let ty = self.generate_type()?;
442             let var = self.create_var(builder, ty)?;
443             let value = self.generate_const(builder, ty)?;
444             builder.def_var(var, value);
445         }
446 
447         Ok(())
448     }
449 
450     /// We generate a function in multiple stages:
451     ///
452     /// * First we generate a random number of empty blocks
453     /// * Then we generate a random pool of variables to be used throughout the function
454     /// * We then visit each block and generate random instructions
455     ///
456     /// Because we generate all blocks and variables up front we already know everything that
457     /// we need when generating instructions (i.e. jump targets / variables)
458     pub fn generate(mut self) -> Result<Function> {
459         let sig = self.generate_signature()?;
460 
461         let mut fn_builder_ctx = FunctionBuilderContext::new();
462         let mut func = Function::with_name_signature(ExternalName::user(0, 0), sig.clone());
463 
464         let mut builder = FunctionBuilder::new(&mut func, &mut fn_builder_ctx);
465 
466         self.blocks = self.generate_blocks(&mut builder, &sig)?;
467 
468         // Function preamble
469         self.generate_jumptables(&mut builder)?;
470 
471         // Main instruction generation loop
472         for (i, (block, block_sig)) in self.blocks.clone().iter().enumerate() {
473             let is_block0 = i == 0;
474             builder.switch_to_block(*block);
475 
476             if is_block0 {
477                 // The first block is special because we must create variables both for the
478                 // block signature and for the variable pool. Additionally, we must also define
479                 // initial values for all variables that are not the function signature.
480                 self.build_variable_pool(&mut builder)?;
481             } else {
482                 // Define variables for the block params
483                 for (i, ty) in block_sig.iter().enumerate() {
484                     let var = self.get_variable_of_type(*ty)?;
485                     let block_param = builder.block_params(*block)[i];
486                     builder.def_var(var, block_param);
487                 }
488             }
489 
490             // Generate block instructions
491             self.generate_instructions(&mut builder)?;
492 
493             self.finalize_block(&mut builder)?;
494         }
495 
496         builder.seal_all_blocks();
497         builder.finalize();
498 
499         Ok(func)
500     }
501 }
502