orca_wasm/ir/
function.rs

1//! Function Builder
2
3use crate::ir::id::{FunctionID, ImportsID, LocalID, ModuleID, TypeID};
4use crate::ir::module::module_functions::{add_local, add_locals, LocalFunction};
5use crate::ir::module::{Module, ReIndexable};
6use crate::ir::types::DataType;
7use crate::ir::types::InstrumentationMode;
8use crate::ir::types::{Body, FuncInstrFlag, FuncInstrMode};
9use crate::module_builder::AddLocal;
10use crate::opcode::{Inject, InjectAt, Instrumenter, MacroOpcode, Opcode};
11use crate::{Component, Location};
12use wasmparser::{Operator, TypeRef};
13
14// TODO: probably need better reasoning with lifetime here
15/// Build a function from scratch
16/// See an example [here].
17///
18/// [here]: https://github.com/thesuhas/orca/blob/314af2df01203e7715aa728e7388cf39c564e9d7/fac_orca/src/main.rs#L16
19pub struct FunctionBuilder<'a> {
20    // pub(crate) id: u32, // function index
21    pub(crate) params: Vec<DataType>,
22    pub(crate) results: Vec<DataType>,
23    #[allow(dead_code)]
24    pub(crate) name: Option<String>,
25    pub body: Body<'a>,
26}
27
28impl<'a> FunctionBuilder<'a> {
29    pub fn new(params: &[DataType], results: &[DataType]) -> Self {
30        Self {
31            params: params.to_vec(),
32            results: results.to_vec(),
33            name: None,
34            body: Body::default(),
35        }
36    }
37
38    /// Finish building a function (have side effect on module IR),
39    /// return function index
40    pub fn finish_module(mut self, module: &mut Module<'a>) -> FunctionID {
41        // add End as last instruction
42        self.end();
43        let id = module.add_local_func(self.name, &self.params, &self.results, self.body.clone());
44
45        assert_eq!(
46            module.functions.len() as u32,
47            module.num_local_functions + module.imports.num_funcs
48        );
49
50        id
51    }
52
53    pub fn replace_import_in_module(mut self, module: &mut Module<'a>, import_id: ImportsID) {
54        // add End as last instruction
55        self.end();
56
57        let err_msg = "Could not replace the specified import with this function,";
58        let imp = module.imports.get(import_id);
59        if let TypeRef::Func(imp_ty_id) = imp.ty {
60            if let Some(ty) = module.types.get(TypeID(imp_ty_id)) {
61                if *ty.params() == self.params && *ty.results() == self.results {
62                    let mut local_func = LocalFunction::new(
63                        TypeID(imp_ty_id),
64                        FunctionID(*import_id),
65                        self.body.clone(),
66                        self.params.len(),
67                    );
68                    local_func.body.name = Some(imp.name.to_string());
69                    module.convert_import_fn_to_local(import_id, local_func);
70                } else {
71                    panic!("{err_msg} types are not equivalent.")
72                }
73            } else {
74                panic!(
75                    "{} could not find an associated type for the specified import ID: {:?}.",
76                    err_msg, import_id
77                )
78            }
79        } else {
80            panic!("{err_msg} the specified import ID does not point to a function!")
81        }
82    }
83
84    /// Finish building a function (have side effect on component IR),
85    /// return function index
86    pub fn finish_component(mut self, comp: &mut Component<'a>, mod_idx: ModuleID) -> FunctionID {
87        // add End as last instruction
88        self.end();
89
90        let id = comp.modules[*mod_idx as usize].add_local_func(
91            self.name,
92            &self.params,
93            &self.results,
94            self.body.clone(),
95        );
96
97        assert_eq!(
98            comp.modules[*mod_idx as usize].functions.len() as u32,
99            comp.modules[*mod_idx as usize].num_local_functions
100                + comp.modules[*mod_idx as usize].imports.num_funcs
101                + comp.modules[*mod_idx as usize].imports.num_funcs_added
102        );
103        id
104    }
105
106    pub fn set_name(&mut self, name: String) {
107        self.name = Some(name)
108    }
109}
110
111impl<'a> Inject<'a> for FunctionBuilder<'a> {
112    /// Inject an operator at the end of the function
113    // here the location of the injection is always at the end of the function
114    fn inject(&mut self, op: Operator<'a>) {
115        self.body.push_op(op)
116    }
117}
118impl<'a> Opcode<'a> for FunctionBuilder<'a> {}
119impl<'a> MacroOpcode<'a> for FunctionBuilder<'a> {}
120
121impl AddLocal for FunctionBuilder<'_> {
122    /// add a local and return local index
123    /// (note that local indices start after)
124    fn add_local(&mut self, ty: DataType) -> LocalID {
125        add_local(
126            ty,
127            self.params.len(),
128            &mut self.body.num_locals,
129            &mut self.body.locals,
130        )
131    }
132}
133
134/// Modify a function
135/// Uses same injection logic as Iterator, which is different from
136/// FunctionBuilder since FunctionModifier does side effect to operators at encoding
137/// (it only modifies the Instrument type)
138pub struct FunctionModifier<'a, 'b> {
139    pub instr_flag: &'a mut FuncInstrFlag<'b>,
140    pub body: &'a mut Body<'b>,
141    pub args: &'a mut Vec<LocalID>,
142    pub(crate) instr_idx: Option<usize>,
143}
144
145impl<'a, 'b> FunctionModifier<'a, 'b> {
146    // by default, the instr_idx the last instruction (always Operator::End indicating end of the function)
147    // and the Instrument type is set to before
148    pub fn init(
149        instr_flag: &'a mut FuncInstrFlag<'b>,
150        body: &'a mut Body<'b>,
151        args: &'a mut Vec<LocalID>,
152    ) -> Self {
153        let instr_idx = body.instructions.len() - 1;
154        let mut func_modifier = FunctionModifier {
155            instr_flag,
156            body,
157            args,
158            instr_idx: None,
159        };
160        func_modifier.before_at(Location::Module {
161            func_idx: FunctionID(0), // not used
162            instr_idx,
163        });
164        func_modifier
165    }
166
167    pub fn add_locals(&mut self, types: &[DataType]) {
168        add_locals(
169            types,
170            self.args.len(),
171            &mut self.body.num_locals,
172            &mut self.body.locals,
173        );
174    }
175}
176impl AddLocal for FunctionModifier<'_, '_> {
177    /// add a local and return local index
178    fn add_local(&mut self, ty: DataType) -> LocalID {
179        add_local(
180            ty,
181            self.args.len(),
182            &mut self.body.num_locals,
183            &mut self.body.locals,
184        )
185    }
186}
187
188impl<'b> Inject<'b> for FunctionModifier<'_, 'b> {
189    // TODO: refactor the inject the function to return a Result rather than panicking?
190    fn inject(&mut self, instr: Operator<'b>) {
191        if self.instr_flag.current_mode.is_some() {
192            // inject at the function level
193            self.instr_flag.add_instr(instr);
194        } else {
195            // inject at instruction level
196            if let Some(idx) = self.instr_idx {
197                let is_special = self.body.instructions[idx].add_instr(instr);
198                // remember if we injected a special instrumentation (to be resolved before encoding)
199                self.instr_flag.has_special_instr |= is_special;
200            } else {
201                panic!("Instruction index not set");
202            }
203        }
204    }
205}
206impl<'b> InjectAt<'b> for FunctionModifier<'_, 'b> {
207    fn inject_at(&mut self, idx: usize, mode: InstrumentationMode, instr: Operator<'b>) {
208        let loc = Location::Module {
209            func_idx: FunctionID(0), // not used
210            instr_idx: idx,
211        };
212        self.set_instrument_mode_at(mode, loc);
213        self.add_instr_at(loc, instr);
214    }
215}
216impl<'b> Opcode<'b> for FunctionModifier<'_, 'b> {}
217impl<'b> MacroOpcode<'b> for FunctionModifier<'_, 'b> {}
218
219impl<'b> Instrumenter<'b> for FunctionModifier<'_, 'b> {
220    ///Can be called after finishing some instrumentation to reset the mode.
221    fn finish_instr(&mut self) {
222        self.instr_flag.finish_instr();
223    }
224    fn curr_instrument_mode(&self) -> &Option<InstrumentationMode> {
225        if let Some(idx) = self.instr_idx {
226            &self.body.instructions[idx].instr_flag.current_mode
227        } else {
228            panic!("Instruction index not set");
229        }
230    }
231
232    fn set_instrument_mode_at(&mut self, mode: InstrumentationMode, loc: Location) {
233        if let Location::Module { instr_idx, .. } = loc {
234            self.instr_idx = Some(instr_idx);
235            self.body.instructions[instr_idx].instr_flag.current_mode = Some(mode);
236        } else {
237            panic!("Should have gotten module location");
238        }
239    }
240
241    fn curr_func_instrument_mode(&self) -> &Option<FuncInstrMode> {
242        &self.instr_flag.current_mode
243    }
244
245    fn set_func_instrument_mode(&mut self, mode: FuncInstrMode) {
246        self.instr_flag.current_mode = Some(mode);
247    }
248
249    fn clear_instr_at(&mut self, loc: Location, mode: InstrumentationMode) {
250        if let Location::Module { instr_idx, .. } = loc {
251            self.body.clear_instr(instr_idx, mode);
252        } else {
253            panic!("Should have gotten module location");
254        }
255    }
256
257    fn add_instr_at(&mut self, loc: Location, instr: Operator<'b>) {
258        if let Location::Module { instr_idx, .. } = loc {
259            self.body.instructions[instr_idx].add_instr(instr);
260        } else {
261            panic!("Should have gotten module location");
262        }
263    }
264
265    fn empty_alternate_at(&mut self, loc: Location) -> &mut Self {
266        if let Location::Module { instr_idx, .. } = loc {
267            self.body.instructions[instr_idx].instr_flag.alternate = Some(vec![]);
268        } else {
269            panic!("Should have gotten Component Location and not Module Location!")
270        }
271
272        self
273    }
274
275    fn empty_block_alt_at(&mut self, loc: Location) -> &mut Self {
276        if let Location::Module { instr_idx, .. } = loc {
277            self.body.instructions[instr_idx].instr_flag.block_alt = Some(vec![]);
278            self.instr_flag.has_special_instr |= true;
279        } else {
280            panic!("Should have gotten Component Location and not Module Location!")
281        }
282
283        self
284    }
285
286    fn get_injected_val(&self, idx: usize) -> &Operator {
287        self.body.instructions[idx].instr_flag.get_instr(idx)
288    }
289}
orca_wasm/ir/function.rs

orca_wasm/ir/
function.rs
