wasmparser/
arity.rs

1/* Copyright 2024 Mozilla Foundation
2 *
3 * Licensed under the Apache License, Version 2.0 (the "License");
4 * you may not use this file except in compliance with the License.
5 * You may obtain a copy of the License at
6 *
7 *     http://www.apache.org/licenses/LICENSE-2.0
8 *
9 * Unless required by applicable law or agreed to in writing, software
10 * distributed under the License is distributed on an "AS IS" BASIS,
11 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 * See the License for the specific language governing permissions and
13 * limitations under the License.
14 */
15
16use crate::prelude::*;
17use crate::{
18    BlockType, BrTable, CompositeInnerType, ContType, FrameKind, FuncType, Operator, RefType,
19    ResumeTable, SubType, TryTable, ValType,
20};
21
22/// To compute the arity (param and result counts) of "variable-arity"
23/// operators, the operator_arity macro needs information about the
24/// module's types and the current control stack. The ModuleArity
25/// trait exposes this information.
26pub trait ModuleArity {
27    /// Type with given index
28    fn sub_type_at(&self, type_idx: u32) -> Option<&SubType>;
29
30    /// Arity (param and result counts) of tag with given index
31    fn tag_type_arity(&self, at: u32) -> Option<(u32, u32)>;
32
33    /// Type index of function with given index
34    fn type_index_of_function(&self, function_idx: u32) -> Option<u32>;
35
36    /// Function type for a given continuation type
37    fn func_type_of_cont_type(&self, c: &ContType) -> Option<&FuncType>;
38
39    /// Sub type for a given reference type
40    fn sub_type_of_ref_type(&self, rt: &RefType) -> Option<&SubType>;
41
42    /// Current height of control stack
43    fn control_stack_height(&self) -> u32;
44
45    /// BlockType and FrameKind of label with given index
46    fn label_block(&self, depth: u32) -> Option<(BlockType, FrameKind)>;
47
48    /// Computes arity of given SubType
49    fn sub_type_arity(&self, t: &SubType) -> Option<(u32, u32)> {
50        match &t.composite_type.inner {
51            CompositeInnerType::Func(f) => {
52                Some((f.params().len() as u32, f.results().len() as u32))
53            }
54            CompositeInnerType::Struct(s) => Some((s.fields.len() as u32, s.fields.len() as u32)),
55            CompositeInnerType::Array(_) => None,
56            CompositeInnerType::Cont(c) => {
57                let f = self.func_type_of_cont_type(c)?;
58                Some((f.params().len() as u32, f.results().len() as u32))
59            }
60        }
61    }
62
63    /// Computes arity of given BlockType
64    fn block_type_arity(&self, ty: BlockType) -> Option<(u32, u32)> {
65        match ty {
66            BlockType::Empty => Some((0, 0)),
67            BlockType::Type(_) => Some((0, 1)),
68            BlockType::FuncType(t) => self.sub_type_arity(self.sub_type_at(t)?),
69        }
70    }
71}
72
73impl Operator<'_> {
74    /// Compute the arity (param and result counts) of the operator, given
75    /// an impl ModuleArity, which stores the necessary module state.
76    pub fn operator_arity(&self, module: &impl ModuleArity) -> Option<(u32, u32)> {
77        #[cfg_attr(not(feature = "simd"), allow(unused_macro_rules))]
78        macro_rules! operator_arity {
79            ($visit:ident $args:tt arity $a:tt -> $b:tt) => (Some(($a, $b)));
80            ($visit:ident { $($args:ident)* } arity custom) => ($visit(module, $($args),*));
81        }
82        macro_rules! define_arity {
83            ($(@$proposal:ident $op:ident $({ $($arg:ident: $argty:ty),* })? => $visit:ident ($($ann:tt)*) )*) => (
84                match self.clone() {
85                    $(
86                        Operator::$op $({ $($arg),* })? => {
87                            $(
88                                $(let _ = $arg;)*
89                            )?
90                            operator_arity!($visit {$( $($arg)* )?} $($ann)*)
91                        }
92                    )*
93                }
94            );
95        }
96        crate::for_each_operator!(define_arity)
97    }
98}
99
100fn visit_block(module: &dyn ModuleArity, block: BlockType) -> Option<(u32, u32)> {
101    let (params, _) = module.block_type_arity(block)?;
102    Some((params, params))
103}
104
105fn visit_loop(module: &dyn ModuleArity, block: BlockType) -> Option<(u32, u32)> {
106    visit_block(module, block)
107}
108
109fn visit_if(module: &dyn ModuleArity, block: BlockType) -> Option<(u32, u32)> {
110    let (params, results) = visit_block(module, block)?;
111    Some((params + 1, results))
112}
113
114fn visit_else(module: &dyn ModuleArity) -> Option<(u32, u32)> {
115    let (ty, _kind) = module.label_block(0)?;
116    let (params, results) = module.block_type_arity(ty)?;
117    Some((results, params))
118}
119
120fn visit_end(module: &dyn ModuleArity) -> Option<(u32, u32)> {
121    let (ty, _kind) = module.label_block(0)?;
122    let (_params, results) = module.block_type_arity(ty)?;
123    Some((results, results))
124}
125
126fn visit_br(module: &dyn ModuleArity, depth: u32) -> Option<(u32, u32)> {
127    let (ty, kind) = module.label_block(depth)?;
128    let (params, results) = module.block_type_arity(ty)?;
129    let n = match kind {
130        FrameKind::Loop => params,
131        _ => results,
132    };
133    Some((n, 0))
134}
135
136fn visit_br_if(module: &dyn ModuleArity, depth: u32) -> Option<(u32, u32)> {
137    let (params, _) = visit_br(module, depth)?;
138    Some((params + 1, params))
139}
140
141fn visit_br_table(module: &dyn ModuleArity, table: BrTable<'_>) -> Option<(u32, u32)> {
142    let (params, results) = visit_br(module, table.default())?;
143    Some((params + 1, results))
144}
145
146fn visit_return(module: &dyn ModuleArity) -> Option<(u32, u32)> {
147    let height = module.control_stack_height().checked_sub(1)?;
148    let (ty, _) = module.label_block(height)?;
149    let (_, results) = module.block_type_arity(ty)?;
150    Some((results, 0))
151}
152
153fn visit_call(module: &dyn ModuleArity, func: u32) -> Option<(u32, u32)> {
154    module.sub_type_arity(module.sub_type_at(module.type_index_of_function(func)?)?)
155}
156
157fn visit_call_indirect(module: &dyn ModuleArity, ty: u32, _table: u32) -> Option<(u32, u32)> {
158    let (params, results) = module.sub_type_arity(module.sub_type_at(ty)?)?;
159    Some((params + 1, results))
160}
161
162fn visit_struct_new(module: &dyn ModuleArity, ty: u32) -> Option<(u32, u32)> {
163    let (params, _results) = module.sub_type_arity(module.sub_type_at(ty)?)?;
164    Some((params, 1))
165}
166
167fn visit_array_new_fixed(_module: &dyn ModuleArity, _ty: u32, size: u32) -> Option<(u32, u32)> {
168    Some((size, 1))
169}
170
171fn visit_br_on_cast(
172    module: &dyn ModuleArity,
173    depth: u32,
174    _from: RefType,
175    _to: RefType,
176) -> Option<(u32, u32)> {
177    let (params, _) = visit_br(module, depth)?;
178    Some((params, params))
179}
180
181fn visit_br_on_cast_fail(
182    module: &dyn ModuleArity,
183    depth: u32,
184    _from: RefType,
185    _to: RefType,
186) -> Option<(u32, u32)> {
187    let (params, _) = visit_br(module, depth)?;
188    Some((params, params))
189}
190
191fn visit_typed_select_multi(_module: &dyn ModuleArity, tys: Vec<ValType>) -> Option<(u32, u32)> {
192    let len = u32::try_from(tys.len()).unwrap();
193    Some((1 + 2 * len, len))
194}
195
196fn visit_return_call(module: &dyn ModuleArity, func: u32) -> Option<(u32, u32)> {
197    let (params, _) = visit_call(module, func)?;
198    Some((params, 0))
199}
200
201fn visit_return_call_indirect(module: &dyn ModuleArity, ty: u32, table: u32) -> Option<(u32, u32)> {
202    let (params, _) = visit_call_indirect(module, ty, table)?;
203    Some((params, 0))
204}
205
206fn visit_try_table(module: &dyn ModuleArity, table: TryTable) -> Option<(u32, u32)> {
207    let (params, _) = module.block_type_arity(table.ty)?;
208    Some((params, params))
209}
210
211fn visit_throw(module: &dyn ModuleArity, tag: u32) -> Option<(u32, u32)> {
212    let (params, _) = module.tag_type_arity(tag)?;
213    Some((params, 0))
214}
215
216fn visit_try(module: &dyn ModuleArity, ty: BlockType) -> Option<(u32, u32)> {
217    visit_block(module, ty)
218}
219
220fn visit_catch(module: &dyn ModuleArity, tag: u32) -> Option<(u32, u32)> {
221    let (params, _) = visit_end(module)?;
222    let (tag_params, _) = module.tag_type_arity(tag)?;
223    Some((params, tag_params))
224}
225
226fn visit_delegate(module: &dyn ModuleArity, _depth: u32) -> Option<(u32, u32)> {
227    visit_end(module)
228}
229
230fn visit_catch_all(module: &dyn ModuleArity) -> Option<(u32, u32)> {
231    let (params, _) = visit_end(module)?;
232    Some((params, 0))
233}
234
235fn visit_call_ref(module: &dyn ModuleArity, ty: u32) -> Option<(u32, u32)> {
236    let (params, results) = module.sub_type_arity(module.sub_type_at(ty)?)?;
237    Some((params + 1, results))
238}
239
240fn visit_return_call_ref(module: &dyn ModuleArity, ty: u32) -> Option<(u32, u32)> {
241    let (params, _) = visit_call_ref(module, ty)?;
242    Some((params, 0))
243}
244
245fn visit_br_on_null(module: &dyn ModuleArity, depth: u32) -> Option<(u32, u32)> {
246    let (params, _results) = visit_br(module, depth)?;
247    Some((params + 1, params + 1))
248}
249
250fn visit_br_on_non_null(module: &dyn ModuleArity, depth: u32) -> Option<(u32, u32)> {
251    let (params, _results) = visit_br(module, depth)?;
252    Some((params, params.checked_sub(1)?))
253}
254
255fn visit_cont_bind(module: &dyn ModuleArity, arg: u32, result: u32) -> Option<(u32, u32)> {
256    let (arg_params, _) = module.sub_type_arity(module.sub_type_at(arg)?)?;
257    let (result_params, _) = module.sub_type_arity(module.sub_type_at(result)?)?;
258    Some((arg_params.checked_sub(result_params)? + 1, 1))
259}
260
261fn visit_suspend(module: &dyn ModuleArity, tag: u32) -> Option<(u32, u32)> {
262    module.tag_type_arity(tag)
263}
264
265fn visit_resume(module: &dyn ModuleArity, cont: u32, _table: ResumeTable) -> Option<(u32, u32)> {
266    let (params, results) = module.sub_type_arity(module.sub_type_at(cont)?)?;
267    Some((params + 1, results))
268}
269
270fn visit_resume_throw(
271    module: &dyn ModuleArity,
272    cont: u32,
273    tag: u32,
274    _table: ResumeTable,
275) -> Option<(u32, u32)> {
276    let (params, _) = module.tag_type_arity(tag)?;
277    let (_, results) = module.sub_type_arity(module.sub_type_at(cont)?)?;
278    Some((params + 1, results))
279}
280
281fn visit_switch(module: &dyn ModuleArity, cont: u32, _tag: u32) -> Option<(u32, u32)> {
282    let (params, _) = module.sub_type_arity(module.sub_type_at(cont)?)?;
283    let st = &module.sub_type_at(cont)?.composite_type.inner;
284    let CompositeInnerType::Cont(ct) = &st else {
285        return None;
286    };
287    let last_param = module.func_type_of_cont_type(ct)?.params().last()?;
288    let (cont_params, _) =
289        module.sub_type_arity(module.sub_type_of_ref_type(&last_param.as_reference_type()?)?)?;
290    Some((params, cont_params))
291}