From 8d0dc939300cc747f6d59ca336b80409f8fbb4f1 Mon Sep 17 00:00:00 2001 From: Chris Fallin Date: Tue, 7 Feb 2023 18:48:11 -0800 Subject: [PATCH] Add support for roundtripping function names. --- src/backend/mod.rs | 20 ++++++++++++++------ src/frontend.rs | 35 +++++++++++++++++++++++++++-------- src/ir/display.rs | 33 +++++++++++++++++++++++++++------ src/ir/func.rs | 36 ++++++++++++++++++++++++++---------- 4 files changed, 94 insertions(+), 30 deletions(-) diff --git a/src/backend/mod.rs b/src/backend/mod.rs index aac78f8..79db37b 100644 --- a/src/backend/mod.rs +++ b/src/backend/mod.rs @@ -584,8 +584,8 @@ pub fn compile(module: &Module<'_>) -> anyhow::Result> { let mut funcs = wasm_encoder::FunctionSection::new(); for (func, func_decl) in module.funcs.entries().skip(num_func_imports) { match func_decl { - FuncDecl::Import(_) => anyhow::bail!("Import comes after func with body: {}", func), - FuncDecl::Lazy(sig, _) | FuncDecl::Body(sig, _) => { + FuncDecl::Import(_, _) => anyhow::bail!("Import comes after func with body: {}", func), + FuncDecl::Lazy(sig, _, _) | FuncDecl::Body(sig, _, _) => { funcs.function(sig.index() as u32); } } @@ -702,17 +702,17 @@ pub fn compile(module: &Module<'_>) -> anyhow::Result> { .par_iter() .map(|(func, func_decl)| -> Result { match func_decl { - FuncDecl::Lazy(_, reader) => { + FuncDecl::Lazy(_, _name, reader) => { let data = &module.orig_bytes[reader.range()]; Ok(FuncOrRawBytes::Raw(data)) } - FuncDecl::Body(_, body) => { - log::debug!("Compiling {}", func); + FuncDecl::Body(_, name, body) => { + log::debug!("Compiling {} \"{}\"", func, name); WasmFuncBackend::new(body)? .compile() .map(|f| FuncOrRawBytes::Func(f)) } - FuncDecl::Import(_) => unreachable!("Should have skipped imports"), + FuncDecl::Import(_, _) => unreachable!("Should have skipped imports"), } }) .collect::>>>()?; @@ -741,6 +741,14 @@ pub fn compile(module: &Module<'_>) -> anyhow::Result> { } into_mod.section(&data); + let mut names = wasm_encoder::NameSection::new(); + let mut func_names = wasm_encoder::NameMap::new(); + for (func, decl) in module.funcs.entries() { + func_names.append(func.index() as u32, decl.name()); + } + names.functions(&func_names); + into_mod.section(&names); + Ok(into_mod.finish()) } diff --git a/src/frontend.rs b/src/frontend.rs index b6c9754..57b4dbb 100644 --- a/src/frontend.rs +++ b/src/frontend.rs @@ -11,7 +11,9 @@ use anyhow::{bail, Result}; use fxhash::{FxHashMap, FxHashSet}; use log::trace; use std::convert::TryFrom; -use wasmparser::{BlockType, DataKind, ExternalKind, Parser, Payload, TypeRef}; +use wasmparser::{ + BlockType, DataKind, ExternalKind, Name, NameSectionReader, Parser, Payload, TypeRef, +}; pub fn wasm_to_ir(bytes: &[u8]) -> Result> { let mut module = Module::with_orig_bytes(bytes); @@ -74,7 +76,7 @@ fn handle_payload<'a>( TypeRef::Func(sig_idx) => { let func = module .funcs - .push(FuncDecl::Import(Signature::from(sig_idx))); + .push(FuncDecl::Import(Signature::from(sig_idx), "".to_owned())); *next_func += 1; ImportKind::Func(func) } @@ -136,17 +138,20 @@ fn handle_payload<'a>( Payload::FunctionSection(reader) => { for sig_idx in reader { let sig_idx = Signature::from(sig_idx?); - module - .funcs - .push(FuncDecl::Body(sig_idx, FunctionBody::default())); + module.funcs.push(FuncDecl::Body( + sig_idx, + "".to_owned(), + FunctionBody::default(), + )); } } Payload::CodeSectionEntry(body) => { let func_idx = Func::new(*next_func); *next_func += 1; - let my_sig = module.funcs[func_idx].sig(); - module.funcs[func_idx] = FuncDecl::Lazy(my_sig, body); + let sig = module.funcs[func_idx].sig(); + let name = module.funcs[func_idx].name().to_owned(); + module.funcs[func_idx] = FuncDecl::Lazy(sig, name, body); } Payload::ExportSection(reader) => { for export in reader { @@ -193,7 +198,21 @@ fn handle_payload<'a>( } } } - Payload::CustomSection { .. } => {} + Payload::CustomSection(reader) if reader.name() == "name" => { + let name_reader = NameSectionReader::new(reader.data(), reader.data_offset())?; + for subsection in name_reader { + let subsection = subsection?; + match subsection { + Name::Function(names) => { + for name in names { + let name = name?; + module.funcs[Func::new(name.index as usize)].set_name(name.name); + } + } + _ => {} + } + } + } Payload::CodeSectionStart { .. } => {} Payload::Version { .. } => {} Payload::ElementSection(reader) => { diff --git a/src/ir/display.rs b/src/ir/display.rs index eb0284d..c717ba0 100644 --- a/src/ir/display.rs +++ b/src/ir/display.rs @@ -209,16 +209,37 @@ impl<'a> Display for ModuleDisplay<'a> { } for (func, func_decl) in self.0.funcs.entries() { match func_decl { - FuncDecl::Body(sig, body) => { - writeln!(f, " {}: {} = # {}", func, sig, sig_strs.get(&sig).unwrap())?; + FuncDecl::Body(sig, name, body) => { + writeln!( + f, + " {} \"{}\": {} = # {}", + func, + name, + sig, + sig_strs.get(&sig).unwrap() + )?; writeln!(f, "{}", body.display(" "))?; } - FuncDecl::Lazy(sig, reader) => { - writeln!(f, " {}: {} = # {}", func, sig, sig_strs.get(&sig).unwrap())?; + FuncDecl::Lazy(sig, name, reader) => { + writeln!( + f, + " {} \"{}\": {} = # {}", + func, + name, + sig, + sig_strs.get(&sig).unwrap() + )?; writeln!(f, " # raw bytes (length {})", reader.range().len())?; } - FuncDecl::Import(sig) => { - writeln!(f, " {}: {} # {}", func, sig, sig_strs.get(&sig).unwrap())?; + FuncDecl::Import(sig, name) => { + writeln!( + f, + " {} \"{}\": {} # {}", + func, + name, + sig, + sig_strs.get(&sig).unwrap() + )?; } } } diff --git a/src/ir/func.rs b/src/ir/func.rs index 7d6df34..cf76637 100644 --- a/src/ir/func.rs +++ b/src/ir/func.rs @@ -6,15 +6,15 @@ use anyhow::Result; #[derive(Clone, Debug)] pub enum FuncDecl<'a> { - Import(Signature), - Lazy(Signature, wasmparser::FunctionBody<'a>), - Body(Signature, FunctionBody), + Import(Signature, String), + Lazy(Signature, String, wasmparser::FunctionBody<'a>), + Body(Signature, String, FunctionBody), } impl<'a> FuncDecl<'a> { pub fn sig(&self) -> Signature { match self { - FuncDecl::Import(sig) => *sig, + FuncDecl::Import(sig, ..) => *sig, FuncDecl::Lazy(sig, ..) => *sig, FuncDecl::Body(sig, ..) => *sig, } @@ -22,9 +22,9 @@ impl<'a> FuncDecl<'a> { pub fn parse(&mut self, module: &Module) -> Result<()> { match self { - FuncDecl::Lazy(sig, body) => { + FuncDecl::Lazy(sig, name, body) => { let body = parse_body(module, *sig, body)?; - *self = FuncDecl::Body(*sig, body); + *self = FuncDecl::Body(*sig, name.clone(), body); Ok(()) } _ => Ok(()), @@ -33,7 +33,7 @@ impl<'a> FuncDecl<'a> { pub fn optimize(&mut self) { match self { - FuncDecl::Body(_, body) => { + FuncDecl::Body(_, _, body) => { body.optimize(); } _ => {} @@ -42,7 +42,7 @@ impl<'a> FuncDecl<'a> { pub fn convert_to_max_ssa(&mut self) { match self { - FuncDecl::Body(_, body) => { + FuncDecl::Body(_, _, body) => { body.convert_to_max_ssa(); } _ => {} @@ -51,17 +51,33 @@ impl<'a> FuncDecl<'a> { pub fn body(&self) -> Option<&FunctionBody> { match self { - FuncDecl::Body(_, body) => Some(body), + FuncDecl::Body(_, _, body) => Some(body), _ => None, } } pub fn body_mut(&mut self) -> Option<&mut FunctionBody> { match self { - FuncDecl::Body(_, body) => Some(body), + FuncDecl::Body(_, _, body) => Some(body), _ => None, } } + + pub fn name(&self) -> &str { + match self { + FuncDecl::Body(_, name, _) | FuncDecl::Lazy(_, name, _) | FuncDecl::Import(_, name) => { + &name[..] + } + } + } + + pub fn set_name(&mut self, new_name: &str) { + match self { + FuncDecl::Body(_, name, _) | FuncDecl::Lazy(_, name, _) | FuncDecl::Import(_, name) => { + *name = new_name.to_owned() + } + } + } } #[derive(Clone, Debug, Default)]