Skip to content

Commit

Permalink
Add --name-unnamed to wasm-tools print
Browse files Browse the repository at this point in the history
Sometimes I've wanted to twiddle with wasm files but I don't enjoy
having to update a bunch of indices. This option can be used to
automatically assign a name to everything in a file to more easily make
the file human editable.
  • Loading branch information
alexcrichton committed Feb 8, 2024
1 parent 97f4a1f commit 2e173cc
Show file tree
Hide file tree
Showing 3 changed files with 176 additions and 52 deletions.
195 changes: 151 additions & 44 deletions crates/wasmprinter/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@
use anyhow::{bail, Context, Result};
use std::collections::{HashMap, HashSet};
use std::fmt::{self, Write};
use std::marker;
use std::mem;
use std::path::Path;
use wasmparser::*;
Expand Down Expand Up @@ -54,6 +55,7 @@ pub struct Printer {
line: usize,
group_lines: Vec<usize>,
code_section_hints: Vec<(u32, Vec<(usize, BranchHint)>)>,
name_unnamed: bool,
}

#[derive(Default)]
Expand All @@ -66,18 +68,41 @@ struct CoreState {
tables: u32,
modules: u32,
instances: u32,
func_names: HashMap<u32, Naming>,
local_names: HashMap<(u32, u32), Naming>,
label_names: HashMap<(u32, u32), Naming>,
type_names: HashMap<u32, Naming>,
tag_names: HashMap<u32, Naming>,
table_names: HashMap<u32, Naming>,
memory_names: HashMap<u32, Naming>,
global_names: HashMap<u32, Naming>,
element_names: HashMap<u32, Naming>,
data_names: HashMap<u32, Naming>,
module_names: HashMap<u32, Naming>,
instance_names: HashMap<u32, Naming>,
func_names: NamingMap<u32, NameFunc>,
local_names: NamingMap<(u32, u32), NameLocal>,
label_names: NamingMap<(u32, u32), NameLabel>,
type_names: NamingMap<u32, NameType>,
tag_names: NamingMap<u32, NameTag>,
table_names: NamingMap<u32, NameTable>,
memory_names: NamingMap<u32, NameMemory>,
global_names: NamingMap<u32, NameGlobal>,
element_names: NamingMap<u32, NameElem>,
data_names: NamingMap<u32, NameData>,
module_names: NamingMap<u32, NameModule>,
instance_names: NamingMap<u32, NameInstance>,
}

/// A map of index-to-name for tracking what are the contents of the name
/// section.
///
/// The type parameter `T` is either `u32` for most index-based maps or a `(u32,
/// u32)` for label/local maps where there are two levels of indices.
///
/// The type parameter `K` is a static description/namespace for what kind of
/// item is contained within this map. That's used by some helper methods to
/// synthesize reasonable names automatically.
struct NamingMap<T, K> {
index_to_name: HashMap<T, Naming>,
_marker: marker::PhantomData<K>,
}

impl<T, K> Default for NamingMap<T, K> {
fn default() -> NamingMap<T, K> {
NamingMap {
index_to_name: HashMap::new(),
_marker: marker::PhantomData,
}
}
}

#[derive(Default)]
Expand All @@ -87,11 +112,11 @@ struct ComponentState {
instances: u32,
components: u32,
values: u32,
type_names: HashMap<u32, Naming>,
func_names: HashMap<u32, Naming>,
component_names: HashMap<u32, Naming>,
instance_names: HashMap<u32, Naming>,
value_names: HashMap<u32, Naming>,
type_names: NamingMap<u32, NameType>,
func_names: NamingMap<u32, NameFunc>,
component_names: NamingMap<u32, NameComponent>,
instance_names: NamingMap<u32, NameInstance>,
value_names: NamingMap<u32, NameValue>,
}

struct State {
Expand Down Expand Up @@ -136,6 +161,20 @@ impl Printer {
self.print_skeleton = print;
}

/// Assign names to all unnamed items.
///
/// If enabled then any previously unnamed item will have a name synthesized
/// that looks like `$#func10` for example. The leading `#` indicates that
/// it's `wasmprinter`-generated. The `func` is the namespace of the name
/// and provides extra context about the item when referenced. The 10 is the
/// local index of the item.
///
/// Note that if the resulting text output is converted back to binary the
/// resulting `name` custom section will not be the same as before.
pub fn name_unnamed(&mut self, enable: bool) {
self.name_unnamed = enable;
}

/// Registers a custom `printer` function to get invoked whenever a custom
/// section of name `section` is seen.
///
Expand Down Expand Up @@ -552,8 +591,8 @@ impl Printer {
}

fn register_names(&mut self, state: &mut State, names: NameSectionReader<'_>) -> Result<()> {
fn indirect_name_map(
into: &mut HashMap<(u32, u32), Naming>,
fn indirect_name_map<K>(
into: &mut NamingMap<(u32, u32), K>,
names: IndirectNameMap<'_>,
name: &str,
) -> Result<()> {
Expand All @@ -567,7 +606,7 @@ impl Printer {
};
for naming in indirect.names {
let naming = naming?;
into.insert(
into.index_to_name.insert(
(indirect.index, naming.index),
Naming::new(naming.name, naming.index, name, used.as_mut()),
);
Expand Down Expand Up @@ -811,8 +850,7 @@ impl Printer {
// we need to be careful to terminate previous param blocks and open
// a new one if that's the case with a named parameter.
for (i, param) in ty.params().iter().enumerate() {
let name = names_for.and_then(|n| state.core.local_names.get(&(n, i as u32)));
params.start_local(name, &mut self.result);
params.start_local(names_for.unwrap_or(u32::MAX), i as u32, self, state);
self.print_valtype(*param)?;
params.end_local(&mut self.result);
}
Expand Down Expand Up @@ -1171,8 +1209,7 @@ impl Printer {
self.newline(offset);
first = false;
}
let name = state.core.local_names.get(&(func_idx, params + local_idx));
locals.start_local(name, &mut self.result);
locals.start_local(func_idx, params + local_idx, self, state);
self.print_valtype(ty)?;
locals.end_local(&mut self.result);
local_idx += 1;
Expand Down Expand Up @@ -1345,28 +1382,55 @@ impl Printer {
Ok(())
}

fn print_idx(&mut self, names: &HashMap<u32, Naming>, idx: u32) -> Result<()> {
fn print_idx<K>(&mut self, names: &NamingMap<u32, K>, idx: u32) -> Result<()>
where
K: NamingNamespace,
{
self._print_idx(&names.index_to_name, idx, K::desc())
}

fn _print_idx(&mut self, names: &HashMap<u32, Naming>, idx: u32, desc: &str) -> Result<()> {
match names.get(&idx) {
Some(name) => write!(self.result, "${}", name.identifier())?,
None => write!(self.result, "{}", idx)?,
None if self.name_unnamed => write!(self.result, "$#{desc}{idx}")?,
None => write!(self.result, "{idx}")?,
}
Ok(())
}

fn print_local_idx(&mut self, state: &State, func: u32, idx: u32) -> Result<()> {
match state.core.local_names.get(&(func, idx)) {
match state.core.local_names.index_to_name.get(&(func, idx)) {
Some(name) => write!(self.result, "${}", name.identifier())?,
None if self.name_unnamed => write!(self.result, "$#local{idx}")?,
None => write!(self.result, "{}", idx)?,
}
Ok(())
}

fn print_name(&mut self, names: &HashMap<u32, Naming>, cur_idx: u32) -> Result<()> {
if let Some(name) = names.get(&cur_idx) {
name.write(&mut self.result);
self.result.push(' ');
fn print_name<K>(&mut self, names: &NamingMap<u32, K>, cur_idx: u32) -> Result<()>
where
K: NamingNamespace,
{
self._print_name(&names.index_to_name, cur_idx, K::desc())
}

fn _print_name(
&mut self,
names: &HashMap<u32, Naming>,
cur_idx: u32,
desc: &str,
) -> Result<()> {
match names.get(&cur_idx) {
Some(name) => {
name.write(&mut self.result);
self.result.push(' ');
}
None if self.name_unnamed => {
write!(self.result, "$#{desc}{cur_idx} ")?;
}
None => {}
}
write!(self.result, "(;{};)", cur_idx)?;
write!(self.result, "(;{cur_idx};)")?;
Ok(())
}

Expand Down Expand Up @@ -2764,35 +2828,46 @@ impl NamedLocalPrinter {
}
}

fn start_local(&mut self, name: Option<&Naming>, dst: &mut String) {
fn start_local(&mut self, func: u32, local: u32, dst: &mut Printer, state: &State) {
let name = state.core.local_names.index_to_name.get(&(func, local));

// Named locals must be in their own group, so if we have a name we need
// to terminate the previous group.
if name.is_some() && self.in_group {
dst.push(')');
dst.result.push(')');
self.in_group = false;
}

if self.first {
self.first = false;
} else {
dst.push(' ');
dst.result.push(' ');
}

// Next we either need a separator if we're already in a group or we
// need to open a group for our new local.
if !self.in_group {
dst.push('(');
dst.push_str(self.group_name);
dst.push(' ');
dst.result.push('(');
dst.result.push_str(self.group_name);
dst.result.push(' ');
self.in_group = true;
}

// Print the optional name if given...
if let Some(name) = name {
name.write(dst);
dst.push(' ');
match name {
Some(name) => {
name.write(&mut dst.result);
dst.result.push(' ');
self.end_group_after_local = true;
}
None if dst.name_unnamed => {
dst.result.push_str(&format!("$#local{local} "));
self.end_group_after_local = true;
}
None => {
self.end_group_after_local = false;
}
}
self.end_group_after_local = name.is_some();
}

fn end_local(&mut self, dst: &mut String) {
Expand Down Expand Up @@ -3019,11 +3094,43 @@ impl Naming {
}
}

fn name_map(into: &mut HashMap<u32, Naming>, names: NameMap<'_>, name: &str) -> Result<()> {
/// Helper trait for the `NamingMap` type's `K` type parameter.
trait NamingNamespace {
fn desc() -> &'static str;
}

macro_rules! naming_namespaces {
($(struct $name:ident => $desc:tt)*) => ($(
struct $name;

impl NamingNamespace for $name {
fn desc() -> &'static str { $desc }
}
)*)
}

naming_namespaces! {
struct NameFunc => "func"
struct NameModule => "module"
struct NameInstance => "instance"
struct NameGlobal => "global"
struct NameMemory => "memory"
struct NameLocal => "local"
struct NameLabel => "label"
struct NameTable => "table"
struct NameValue => "value"
struct NameType => "type"
struct NameData => "data"
struct NameElem => "elem"
struct NameComponent => "component"
struct NameTag => "tag"
}

fn name_map<K>(into: &mut NamingMap<u32, K>, names: NameMap<'_>, name: &str) -> Result<()> {
let mut used = HashSet::new();
for naming in names {
let naming = naming?;
into.insert(
into.index_to_name.insert(
naming.index,
Naming::new(naming.name, naming.index, name, Some(&mut used)),
);
Expand Down
24 changes: 16 additions & 8 deletions crates/wasmprinter/src/operator.rs
Original file line number Diff line number Diff line change
Expand Up @@ -68,12 +68,18 @@ impl<'a, 'b> PrintOperator<'a, 'b> {
}

let key = (self.state.core.funcs, self.label);
let has_name = if let Some(name) = self.state.core.label_names.get(&key) {
self.printer.result.push_str(" ");
name.write(&mut self.printer.result);
true
} else {
false
let has_name = match self.state.core.label_names.index_to_name.get(&key) {
Some(name) => {
self.printer.result.push_str(" ");
name.write(&mut self.printer.result);
true
}
None if self.printer.name_unnamed => {
let depth = self.cur_depth();
write!(self.result(), " $#label{depth}")?;
true
}
None => false,
};
match ty {
BlockType::Empty => {}
Expand Down Expand Up @@ -127,7 +133,7 @@ impl<'a, 'b> PrintOperator<'a, 'b> {
.and_then(|idx| self.label_indices.get(idx as usize).copied())
.and_then(|label_idx| {
let key = (self.state.core.funcs, label_idx);
self.state.core.label_names.get(&key)
self.state.core.label_names.index_to_name.get(&key)
});

// This is a bit tricky, but if there's a shallower label than
Expand All @@ -139,7 +145,7 @@ impl<'a, 'b> PrintOperator<'a, 'b> {
let name = name.and_then(|name| {
for other_label in self.label_indices[i as usize..].iter() {
let key = (self.state.core.funcs, *other_label);
if let Some(other) = self.state.core.label_names.get(&key) {
if let Some(other) = self.state.core.label_names.index_to_name.get(&key) {
if name.name == other.name {
return None;
}
Expand All @@ -151,6 +157,8 @@ impl<'a, 'b> PrintOperator<'a, 'b> {
match name {
Some(name) => name.write(&mut self.printer.result),

None if self.printer.name_unnamed => write!(self.result(), "$#label{i}")?,

// If this label has some name also print its pseudo-name as
// `@N` to help match things up in the text format.
None => write!(self.result(), "{depth} (;@{i};)")?,
Expand Down
9 changes: 9 additions & 0 deletions src/bin/wasm-tools/print.rs
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,14 @@ pub struct Opts {
/// replaced with "..." instead of printing their actual contents.
#[clap(long)]
skeleton: bool,

/// Ensure all wasm items have `$`-based names, even if they don't have an
/// entry in the `name` section.
///
/// This option, when enabled, will synthesize names for any item which
/// doesn't previously have a name.
#[clap(long)]
name_unnamed: bool,
}

impl Opts {
Expand All @@ -29,6 +37,7 @@ impl Opts {
let mut printer = wasmprinter::Printer::new();
printer.print_offsets(self.print_offsets);
printer.print_skeleton(self.skeleton);
printer.name_unnamed(self.name_unnamed);
let wat = printer.print(&wasm)?;
self.io.output(wasm_tools::Output::Wat(&wat))?;
Ok(())
Expand Down

0 comments on commit 2e173cc

Please sign in to comment.