1use crate::{Error, Result};
4use std::iter::IntoIterator;
5use wasmparser::{
6 Data, DataKind, Export, ExternalKind, Import, Operator, Payload, SectionLimited, TypeRef,
7 ValidPayload, Validator,
8};
9use zingen::wasm::{Data as DataSet, Env, Exports, Functions, HostFunc, Imports};
10
11#[derive(Default)]
13pub struct Parser<'p> {
14 pub env: Env,
16 pub funcs: Functions<'p>,
18}
19
20impl<'p> Parser<'p> {
21 pub fn parse(&mut self, wasm: &'p [u8]) -> Result<()> {
23 let mut validator = Validator::new();
24
25 for payload in wasmparser::Parser::new(0).parse_all(wasm) {
27 let payload = payload?;
28 let valid_payload = validator.payload(&payload)?;
29
30 match &payload {
31 Payload::ImportSection(reader) => self.env.imports = Self::imports(reader)?,
32 Payload::DataSection(reader) => self.env.data = Self::data(reader)?,
33 Payload::ExportSection(reader) => self.env.exports = Self::exports(reader)?,
34 _ => {}
35 }
36
37 if let ValidPayload::Func(to_validator, body) = valid_payload {
38 self.funcs
39 .add(to_validator.into_validator(Default::default()), body);
40 }
41 }
42
43 let mut slots = self.env.imports.reserved();
45 for (idx, fun) in self.funcs.iter() {
46 let sig = fun.sig()?;
47 let locals = fun.body.get_locals_reader()?.get_count();
48 let params = sig.params().len();
49 tracing::trace!(
50 "computing slots for function {idx}, locals: {locals}, params: {params}, reserved: {slots}, external: {}",
51 self.env.is_external(fun.index())
52 );
53
54 self.env.slots.insert(fun.index(), slots);
55 self.env
56 .funcs
57 .insert(fun.index(), (params as u32, sig.results().len() as u32));
58
59 slots += locals;
60
61 if !self.env.is_external(fun.index()) && !self.env.is_main(fun.index()) {
63 slots += params as u32;
64 }
65 }
66
67 Ok(())
68 }
69
70 pub fn drain_selectors(&mut self) -> Functions<'p> {
72 self.funcs.drain_selectors(&self.env.exports)
73 }
74
75 fn data(reader: &SectionLimited<Data>) -> Result<DataSet> {
77 let mut dataset = DataSet::default();
78 let mut iter = reader.clone().into_iter();
79 while let Some(Ok(data)) = iter.next() {
80 if let DataKind::Active {
81 memory_index: _,
82 offset_expr,
83 } = data.kind
84 {
85 let mut reader = offset_expr.get_binary_reader();
87 let Operator::I32Const { value: offset } = reader.read_operator()? else {
88 return Err(Error::InvalidDataOffset);
89 };
90
91 dataset.insert(offset, data.data.into());
92 }
93 }
94
95 Ok(dataset)
96 }
97
98 pub fn exports(reader: &SectionLimited<Export>) -> Result<Exports> {
100 let mut exports = Exports::default();
101 let mut iter = reader.clone().into_iter();
102 while let Some(Ok(Export {
103 name,
104 kind: ExternalKind::Func,
105 index,
106 })) = iter.next()
107 {
108 exports.insert(index, name.into());
109 }
110
111 Ok(exports)
112 }
113
114 pub fn imports(reader: &SectionLimited<Import>) -> Result<Imports> {
116 let mut index = 0;
118
119 let mut imports = Imports::default();
120 let mut iter = reader.clone().into_iter();
121 while let Some(Ok(Import {
122 module,
123 name,
124 ty: TypeRef::Func(_),
125 })) = iter.next()
126 {
127 let func = HostFunc::try_from((module, name))?;
128 tracing::trace!("imported function: {}::{} at {index}", module, name);
129 imports.insert(index, func);
130 index += 1;
131 }
132
133 Ok(imports)
134 }
135}
136
137impl<'p> TryFrom<&'p [u8]> for Parser<'p> {
138 type Error = Error;
139
140 fn try_from(wasm: &'p [u8]) -> Result<Self> {
141 let mut parser = Self::default();
142 parser.parse(wasm)?;
143 Ok(parser)
144 }
145}