Add support for roundtripping function names.

This commit is contained in:
Chris Fallin 2023-02-07 18:48:11 -08:00
parent 6f3d04d2b5
commit 8d0dc93930
4 changed files with 94 additions and 30 deletions

View file

@ -584,8 +584,8 @@ pub fn compile(module: &Module<'_>) -> anyhow::Result<Vec<u8>> {
let mut funcs = wasm_encoder::FunctionSection::new(); let mut funcs = wasm_encoder::FunctionSection::new();
for (func, func_decl) in module.funcs.entries().skip(num_func_imports) { for (func, func_decl) in module.funcs.entries().skip(num_func_imports) {
match func_decl { match func_decl {
FuncDecl::Import(_) => anyhow::bail!("Import comes after func with body: {}", func), FuncDecl::Import(_, _) => anyhow::bail!("Import comes after func with body: {}", func),
FuncDecl::Lazy(sig, _) | FuncDecl::Body(sig, _) => { FuncDecl::Lazy(sig, _, _) | FuncDecl::Body(sig, _, _) => {
funcs.function(sig.index() as u32); funcs.function(sig.index() as u32);
} }
} }
@ -702,17 +702,17 @@ pub fn compile(module: &Module<'_>) -> anyhow::Result<Vec<u8>> {
.par_iter() .par_iter()
.map(|(func, func_decl)| -> Result<FuncOrRawBytes> { .map(|(func, func_decl)| -> Result<FuncOrRawBytes> {
match func_decl { match func_decl {
FuncDecl::Lazy(_, reader) => { FuncDecl::Lazy(_, _name, reader) => {
let data = &module.orig_bytes[reader.range()]; let data = &module.orig_bytes[reader.range()];
Ok(FuncOrRawBytes::Raw(data)) Ok(FuncOrRawBytes::Raw(data))
} }
FuncDecl::Body(_, body) => { FuncDecl::Body(_, name, body) => {
log::debug!("Compiling {}", func); log::debug!("Compiling {} \"{}\"", func, name);
WasmFuncBackend::new(body)? WasmFuncBackend::new(body)?
.compile() .compile()
.map(|f| FuncOrRawBytes::Func(f)) .map(|f| FuncOrRawBytes::Func(f))
} }
FuncDecl::Import(_) => unreachable!("Should have skipped imports"), FuncDecl::Import(_, _) => unreachable!("Should have skipped imports"),
} }
}) })
.collect::<Result<Vec<FuncOrRawBytes<'_>>>>()?; .collect::<Result<Vec<FuncOrRawBytes<'_>>>>()?;
@ -741,6 +741,14 @@ pub fn compile(module: &Module<'_>) -> anyhow::Result<Vec<u8>> {
} }
into_mod.section(&data); 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()) Ok(into_mod.finish())
} }

View file

@ -11,7 +11,9 @@ use anyhow::{bail, Result};
use fxhash::{FxHashMap, FxHashSet}; use fxhash::{FxHashMap, FxHashSet};
use log::trace; use log::trace;
use std::convert::TryFrom; 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<Module<'_>> { pub fn wasm_to_ir(bytes: &[u8]) -> Result<Module<'_>> {
let mut module = Module::with_orig_bytes(bytes); let mut module = Module::with_orig_bytes(bytes);
@ -74,7 +76,7 @@ fn handle_payload<'a>(
TypeRef::Func(sig_idx) => { TypeRef::Func(sig_idx) => {
let func = module let func = module
.funcs .funcs
.push(FuncDecl::Import(Signature::from(sig_idx))); .push(FuncDecl::Import(Signature::from(sig_idx), "".to_owned()));
*next_func += 1; *next_func += 1;
ImportKind::Func(func) ImportKind::Func(func)
} }
@ -136,17 +138,20 @@ fn handle_payload<'a>(
Payload::FunctionSection(reader) => { Payload::FunctionSection(reader) => {
for sig_idx in reader { for sig_idx in reader {
let sig_idx = Signature::from(sig_idx?); let sig_idx = Signature::from(sig_idx?);
module module.funcs.push(FuncDecl::Body(
.funcs sig_idx,
.push(FuncDecl::Body(sig_idx, FunctionBody::default())); "".to_owned(),
FunctionBody::default(),
));
} }
} }
Payload::CodeSectionEntry(body) => { Payload::CodeSectionEntry(body) => {
let func_idx = Func::new(*next_func); let func_idx = Func::new(*next_func);
*next_func += 1; *next_func += 1;
let my_sig = module.funcs[func_idx].sig(); let sig = module.funcs[func_idx].sig();
module.funcs[func_idx] = FuncDecl::Lazy(my_sig, body); let name = module.funcs[func_idx].name().to_owned();
module.funcs[func_idx] = FuncDecl::Lazy(sig, name, body);
} }
Payload::ExportSection(reader) => { Payload::ExportSection(reader) => {
for export in 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::CodeSectionStart { .. } => {}
Payload::Version { .. } => {} Payload::Version { .. } => {}
Payload::ElementSection(reader) => { Payload::ElementSection(reader) => {

View file

@ -209,16 +209,37 @@ impl<'a> Display for ModuleDisplay<'a> {
} }
for (func, func_decl) in self.0.funcs.entries() { for (func, func_decl) in self.0.funcs.entries() {
match func_decl { match func_decl {
FuncDecl::Body(sig, body) => { FuncDecl::Body(sig, name, body) => {
writeln!(f, " {}: {} = # {}", func, sig, sig_strs.get(&sig).unwrap())?; writeln!(
f,
" {} \"{}\": {} = # {}",
func,
name,
sig,
sig_strs.get(&sig).unwrap()
)?;
writeln!(f, "{}", body.display(" "))?; writeln!(f, "{}", body.display(" "))?;
} }
FuncDecl::Lazy(sig, reader) => { FuncDecl::Lazy(sig, name, reader) => {
writeln!(f, " {}: {} = # {}", func, sig, sig_strs.get(&sig).unwrap())?; writeln!(
f,
" {} \"{}\": {} = # {}",
func,
name,
sig,
sig_strs.get(&sig).unwrap()
)?;
writeln!(f, " # raw bytes (length {})", reader.range().len())?; writeln!(f, " # raw bytes (length {})", reader.range().len())?;
} }
FuncDecl::Import(sig) => { FuncDecl::Import(sig, name) => {
writeln!(f, " {}: {} # {}", func, sig, sig_strs.get(&sig).unwrap())?; writeln!(
f,
" {} \"{}\": {} # {}",
func,
name,
sig,
sig_strs.get(&sig).unwrap()
)?;
} }
} }
} }

View file

@ -6,15 +6,15 @@ use anyhow::Result;
#[derive(Clone, Debug)] #[derive(Clone, Debug)]
pub enum FuncDecl<'a> { pub enum FuncDecl<'a> {
Import(Signature), Import(Signature, String),
Lazy(Signature, wasmparser::FunctionBody<'a>), Lazy(Signature, String, wasmparser::FunctionBody<'a>),
Body(Signature, FunctionBody), Body(Signature, String, FunctionBody),
} }
impl<'a> FuncDecl<'a> { impl<'a> FuncDecl<'a> {
pub fn sig(&self) -> Signature { pub fn sig(&self) -> Signature {
match self { match self {
FuncDecl::Import(sig) => *sig, FuncDecl::Import(sig, ..) => *sig,
FuncDecl::Lazy(sig, ..) => *sig, FuncDecl::Lazy(sig, ..) => *sig,
FuncDecl::Body(sig, ..) => *sig, FuncDecl::Body(sig, ..) => *sig,
} }
@ -22,9 +22,9 @@ impl<'a> FuncDecl<'a> {
pub fn parse(&mut self, module: &Module) -> Result<()> { pub fn parse(&mut self, module: &Module) -> Result<()> {
match self { match self {
FuncDecl::Lazy(sig, body) => { FuncDecl::Lazy(sig, name, body) => {
let body = parse_body(module, *sig, body)?; let body = parse_body(module, *sig, body)?;
*self = FuncDecl::Body(*sig, body); *self = FuncDecl::Body(*sig, name.clone(), body);
Ok(()) Ok(())
} }
_ => Ok(()), _ => Ok(()),
@ -33,7 +33,7 @@ impl<'a> FuncDecl<'a> {
pub fn optimize(&mut self) { pub fn optimize(&mut self) {
match self { match self {
FuncDecl::Body(_, body) => { FuncDecl::Body(_, _, body) => {
body.optimize(); body.optimize();
} }
_ => {} _ => {}
@ -42,7 +42,7 @@ impl<'a> FuncDecl<'a> {
pub fn convert_to_max_ssa(&mut self) { pub fn convert_to_max_ssa(&mut self) {
match self { match self {
FuncDecl::Body(_, body) => { FuncDecl::Body(_, _, body) => {
body.convert_to_max_ssa(); body.convert_to_max_ssa();
} }
_ => {} _ => {}
@ -51,17 +51,33 @@ impl<'a> FuncDecl<'a> {
pub fn body(&self) -> Option<&FunctionBody> { pub fn body(&self) -> Option<&FunctionBody> {
match self { match self {
FuncDecl::Body(_, body) => Some(body), FuncDecl::Body(_, _, body) => Some(body),
_ => None, _ => None,
} }
} }
pub fn body_mut(&mut self) -> Option<&mut FunctionBody> { pub fn body_mut(&mut self) -> Option<&mut FunctionBody> {
match self { match self {
FuncDecl::Body(_, body) => Some(body), FuncDecl::Body(_, _, body) => Some(body),
_ => None, _ => 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)] #[derive(Clone, Debug, Default)]