use crate::{Error, Result};
use std::iter::IntoIterator;
use wasmparser::{
Data, DataKind, Export, ExternalKind, Import, Operator, Payload, SectionLimited, TypeRef,
ValidPayload, Validator,
};
use zingen::wasm::{Data as DataSet, Env, Exports, Functions, HostFunc, Imports};
#[derive(Default)]
pub struct Parser<'p> {
pub env: Env,
pub funcs: Functions<'p>,
}
impl<'p> Parser<'p> {
pub fn parse(&mut self, wasm: &'p [u8]) -> Result<()> {
let mut validator = Validator::new();
for payload in wasmparser::Parser::new(0).parse_all(wasm) {
let payload = payload?;
let valid_payload = validator.payload(&payload)?;
match &payload {
Payload::ImportSection(reader) => self.env.imports = Self::imports(reader)?,
Payload::DataSection(reader) => self.env.data = Self::data(reader)?,
Payload::ExportSection(reader) => self.env.exports = Self::exports(reader)?,
_ => {}
}
if let ValidPayload::Func(to_validator, body) = valid_payload {
self.funcs
.add(to_validator.into_validator(Default::default()), body);
}
}
let mut slots = self.env.imports.reserved();
for (idx, fun) in self.funcs.iter() {
let sig = fun.sig()?;
let locals = fun.body.get_locals_reader()?.get_count();
let params = sig.params().len();
tracing::trace!(
"computing slots for function {idx}, locals: {locals}, params: {params}, reserved: {slots}, external: {}",
self.env.is_external(fun.index())
);
self.env.slots.insert(fun.index(), slots);
self.env
.funcs
.insert(fun.index(), (params as u32, sig.results().len() as u32));
slots += locals;
if !self.env.is_external(fun.index()) && !self.env.is_main(fun.index()) {
slots += params as u32;
}
}
Ok(())
}
pub fn drain_selectors(&mut self) -> Functions<'p> {
self.funcs.drain_selectors(&self.env.exports)
}
fn data(reader: &SectionLimited<Data>) -> Result<DataSet> {
let mut dataset = DataSet::default();
let mut iter = reader.clone().into_iter();
while let Some(Ok(data)) = iter.next() {
if let DataKind::Active {
memory_index: _,
offset_expr,
} = data.kind
{
let mut reader = offset_expr.get_binary_reader();
let Operator::I32Const { value: offset } = reader.read_operator()? else {
return Err(Error::InvalidDataOffset);
};
dataset.insert(offset, data.data.into());
}
}
Ok(dataset)
}
pub fn exports(reader: &SectionLimited<Export>) -> Result<Exports> {
let mut exports = Exports::default();
let mut iter = reader.clone().into_iter();
while let Some(Ok(Export {
name,
kind: ExternalKind::Func,
index,
})) = iter.next()
{
exports.insert(index, name.into());
}
Ok(exports)
}
pub fn imports(reader: &SectionLimited<Import>) -> Result<Imports> {
let mut index = 0;
let mut imports = Imports::default();
let mut iter = reader.clone().into_iter();
while let Some(Ok(Import {
module,
name,
ty: TypeRef::Func(_),
})) = iter.next()
{
let func = HostFunc::try_from((module, name))?;
tracing::trace!("imported function: {}::{} at {index}", module, name);
imports.insert(index, func);
index += 1;
}
Ok(imports)
}
}
impl<'p> TryFrom<&'p [u8]> for Parser<'p> {
type Error = Error;
fn try_from(wasm: &'p [u8]) -> Result<Self> {
let mut parser = Self::default();
parser.parse(wasm)?;
Ok(parser)
}
}