diff --git a/src/frontend.rs b/src/frontend.rs index c250642..4fd9c26 100644 --- a/src/frontend.rs +++ b/src/frontend.rs @@ -27,6 +27,24 @@ pub fn wasm_to_ir(bytes: &[u8]) -> Result> { Ok(module) } +fn parse_init_expr<'a>(init_expr: &wasmparser::InitExpr<'a>) -> Result { + let operators = init_expr + .get_operators_reader() + .into_iter() + .collect::, _>>()?; + if operators.len() != 2 || !matches!(&operators[1], &wasmparser::Operator::End) { + anyhow::bail!( + "Unsupported operator seq in base-address expr: {:?}", + operators + ); + } + Ok(match &operators[0] { + &wasmparser::Operator::I32Const { value } => value as usize, + &wasmparser::Operator::I64Const { value } => value as usize, + op => anyhow::bail!("Unsupported data segment base-address operator: {:?}", op), + }) +} + fn handle_payload<'a>( module: &mut Module<'a>, payload: Payload<'a>, @@ -138,26 +156,7 @@ fn handle_payload<'a>( } => { let data = segment.data.to_vec(); let memory = Memory::from(*memory_index); - let operators = init_expr - .get_operators_reader() - .into_iter() - .collect::, _>>()?; - if operators.len() != 2 - || !matches!(&operators[1], &wasmparser::Operator::End) - { - anyhow::bail!( - "Unsupported operator seq in base-address expr: {:?}", - operators - ); - } - let offset = match &operators[0] { - &wasmparser::Operator::I32Const { value } => value as usize, - &wasmparser::Operator::I64Const { value } => value as usize, - op => anyhow::bail!( - "Unsupported data segment base-address operator: {:?}", - op - ), - }; + let offset = parse_init_expr(init_expr)?; module .frontend_memory_mut(memory) .segments @@ -166,7 +165,54 @@ fn handle_payload<'a>( } } } - _ => {} + Payload::CustomSection { .. } => {} + Payload::CodeSectionStart { .. } => {} + Payload::Version { .. } => {} + Payload::ElementSection(reader) => { + for element in reader { + let element = element?; + if element.ty != wasmparser::Type::FuncRef { + anyhow::bail!("Unsupported table type: {:?}", element.ty); + } + match &element.kind { + wasmparser::ElementKind::Passive => {} + wasmparser::ElementKind::Declared => {} + wasmparser::ElementKind::Active { + table_index, + init_expr, + } => { + let table = Table::from(*table_index); + let offset = parse_init_expr(&init_expr)?; + let items = element + .items + .get_items_reader()? + .into_iter() + .collect::, _>>()?; + let mut funcs = vec![]; + for item in items { + let func = match item { + wasmparser::ElementItem::Func(func_idx) => Func::from(func_idx), + _ => anyhow::bail!("Unsupported element item: {:?}", item), + }; + funcs.push(func); + } + + let table_items = module + .frontend_table_mut(table) + .func_elements + .as_mut() + .unwrap(); + if (offset + funcs.len()) > table_items.len() { + table_items.resize(offset + funcs.len(), Func::invalid()); + } + table_items[offset..(offset + funcs.len())].copy_from_slice(&funcs[..]); + } + } + } + } + payload => { + log::warn!("Skipping section: {:?}", payload); + } } Ok(()) diff --git a/src/ir/display.rs b/src/ir/display.rs index 4ede1f9..ab13671 100644 --- a/src/ir/display.rs +++ b/src/ir/display.rs @@ -155,8 +155,13 @@ impl<'a> Display for ModuleDisplay<'a> { for (global, global_ty) in self.0.globals() { writeln!(f, " {}: {}", global, global_ty)?; } - for (table, table_ty) in self.0.tables() { - writeln!(f, " {}: {}", table, table_ty)?; + for (table, table_data) in self.0.tables() { + writeln!(f, " {}: {}", table, table_data.ty)?; + if let Some(funcs) = &table_data.func_elements { + for (i, &func) in funcs.iter().enumerate() { + writeln!(f, " {}[{}]: {}", table, i, func)?; + } + } } for (memory, memory_data) in self.0.memories() { writeln!( diff --git a/src/ir/module.rs b/src/ir/module.rs index 3b80b37..77814f9 100644 --- a/src/ir/module.rs +++ b/src/ir/module.rs @@ -10,7 +10,7 @@ pub struct Module<'a> { funcs: EntityVec, signatures: EntityVec, globals: EntityVec, - tables: EntityVec, + tables: EntityVec, imports: Vec, exports: Vec, memories: EntityVec, @@ -37,6 +37,12 @@ pub struct MemorySegment { pub data: Vec, } +#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct TableData { + pub ty: Type, + pub func_elements: Option>, +} + impl From<&wasmparser::FuncType> for SignatureData { fn from(fty: &wasmparser::FuncType) -> Self { Self { @@ -149,11 +155,11 @@ impl<'a> Module<'a> { pub fn globals<'b>(&'b self) -> impl Iterator + 'b { self.globals.entries().map(|(id, ty)| (id, *ty)) } - pub fn table_ty(&self, id: Table) -> Type { - self.tables[id] + pub fn table<'b>(&'b self, id: Table) -> &'b TableData { + &self.tables[id] } - pub fn tables<'b>(&'b self) -> impl Iterator + 'b { - self.tables.entries().map(|(id, ty)| (id, *ty)) + pub fn tables<'b>(&'b self) -> impl Iterator + 'b { + self.tables.entries() } pub fn memories<'b>(&'b self) -> impl Iterator + 'b { self.memories.entries() @@ -172,7 +178,15 @@ impl<'a> Module<'a> { self.funcs.push(body) } pub(crate) fn frontend_add_table(&mut self, ty: Type) -> Table { - self.tables.push(ty) + let func_elements = if ty == Type::FuncRef { + Some(vec![]) + } else { + None + }; + self.tables.push(TableData { ty, func_elements }) + } + pub(crate) fn frontend_table_mut<'b>(&'b mut self, table: Table) -> &'b mut TableData { + &mut self.tables[table] } pub(crate) fn frontend_add_global(&mut self, ty: Type) -> Global { self.globals.push(ty) diff --git a/src/op_traits.rs b/src/op_traits.rs index 15f4894..27cd0d3 100644 --- a/src/op_traits.rs +++ b/src/op_traits.rs @@ -1,7 +1,7 @@ //! Metadata on operators. use crate::entity::EntityVec; -use crate::ir::{Global, Local, Module, Signature, Table, Value, Type}; +use crate::ir::{Global, Local, Module, Signature, Table, Type, Value}; use crate::Operator; use anyhow::Result; @@ -216,7 +216,7 @@ pub fn op_inputs( Operator::I32ReinterpretF32 => Ok(vec![Type::F32]), Operator::I64ReinterpretF64 => Ok(vec![Type::F64]), Operator::TableGet { .. } => Ok(vec![Type::I32]), - Operator::TableSet { table_index } => Ok(vec![Type::I32, module.table_ty(*table_index)]), + Operator::TableSet { table_index } => Ok(vec![Type::I32, module.table(*table_index).ty]), Operator::TableGrow { .. } => Ok(vec![Type::I32]), Operator::TableSize { .. } => Ok(vec![]), Operator::MemorySize { .. } => Ok(vec![]), @@ -425,7 +425,7 @@ pub fn op_outputs( Operator::F64ReinterpretI64 => Ok(vec![Type::F64]), Operator::I32ReinterpretF32 => Ok(vec![Type::I32]), Operator::I64ReinterpretF64 => Ok(vec![Type::I64]), - Operator::TableGet { table_index } => Ok(vec![module.table_ty(*table_index)]), + Operator::TableGet { table_index } => Ok(vec![module.table(*table_index).ty]), Operator::TableSet { .. } => Ok(vec![]), Operator::TableGrow { .. } => Ok(vec![]), Operator::TableSize { .. } => Ok(vec![Type::I32]),