zinkc/
parser.rs

1//! Zink parser
2
3use 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/// WASM module parser
12#[derive(Default)]
13pub struct Parser<'p> {
14    /// Function environment
15    pub env: Env,
16    /// All functions
17    pub funcs: Functions<'p>,
18}
19
20impl<'p> Parser<'p> {
21    /// Parse WASM module.
22    pub fn parse(&mut self, wasm: &'p [u8]) -> Result<()> {
23        let mut validator = Validator::new();
24
25        // Compile functions.
26        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        // compute slots from functions
44        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            // process prarams for internal functions only
62            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    /// Drain selectors from parsed functions
71    pub fn drain_selectors(&mut self) -> Functions<'p> {
72        self.funcs.drain_selectors(&self.env.exports)
73    }
74
75    /// Parse data section.
76    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                // [i32.const offset call_indirect]
86                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    /// Parse export section
99    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    /// Parse import section.
115    pub fn imports(reader: &SectionLimited<Import>) -> Result<Imports> {
116        // TODO: use real index from WASM. (#122)
117        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}