From f8cdb67a9c4d933702b8e6b5bb296e2671cfcbfb Mon Sep 17 00:00:00 2001 From: Alex Crichton Date: Tue, 19 Nov 2019 11:50:07 -0800 Subject: [PATCH] Allow custom sections to specify gc roots This commit adds APIs necessary to allow custom sections to reference core wasm functions/memories/etc and provide gc roots into them. This is intended to be used, for example, with wasm interface types. --- src/module/custom.rs | 11 +++ src/passes/gc.rs | 4 +- src/passes/mod.rs | 2 +- src/passes/used.rs | 156 +++++++++++++++++++++++-------------------- 4 files changed, 98 insertions(+), 75 deletions(-) diff --git a/src/module/custom.rs b/src/module/custom.rs index 13cc9875..98de4037 100644 --- a/src/module/custom.rs +++ b/src/module/custom.rs @@ -1,5 +1,6 @@ //! Working with custom sections. +use crate::passes::Roots; use crate::tombstone_arena::{Id, Tombstone, TombstoneArena}; use crate::CodeTransform; use crate::IdsToIndices; @@ -29,6 +30,16 @@ pub trait CustomSection: WalrusAny + Debug + Send + Sync { /// payload. `walrus` will handle these for you. fn data(&self, ids_to_indices: &IdsToIndices) -> Cow<[u8]>; + /// Add any core wasm roots to the provided `roots` argument. + /// + /// This function will add any referenced core wasm items into the `Roots` + /// array provided. + /// + /// The default provided method does nothing. + fn gc_add_roots(&self, roots: &mut Roots) { + drop(roots); + } + /// Apply the given code transformations to this custom section. /// /// If the module was not configured with `preserve_code_transform = true`, diff --git a/src/passes/gc.rs b/src/passes/gc.rs index 87a050b9..f4831810 100644 --- a/src/passes/gc.rs +++ b/src/passes/gc.rs @@ -4,13 +4,13 @@ //! internally and can be safely removed. use crate::map::IdHashSet; -use crate::passes::Used; +use crate::passes::used::Used; use crate::{ImportKind, Module}; use id_arena::Id; /// Run GC passes over the module specified. pub fn run(m: &mut Module) { - let used = Used::new(m, m.exports.iter().map(|e| e.id())); + let used = Used::new(m); let mut unused_imports = Vec::new(); for import in m.imports.iter() { diff --git a/src/passes/mod.rs b/src/passes/mod.rs index 7664a6de..43f5c1aa 100644 --- a/src/passes/mod.rs +++ b/src/passes/mod.rs @@ -3,4 +3,4 @@ pub mod gc; mod used; pub mod validate; -pub use self::used::Used; +pub use self::used::Roots; diff --git a/src/passes/used.rs b/src/passes/used.rs index d65e66db..5f681f3b 100644 --- a/src/passes/used.rs +++ b/src/passes/used.rs @@ -1,12 +1,72 @@ use crate::ir::*; use crate::map::IdHashSet; -use crate::{ - ActiveDataLocation, Data, DataId, DataKind, Element, ExportId, ExportItem, Function, InitExpr, -}; +use crate::{ActiveDataLocation, Data, DataId, DataKind, Element, ExportItem, Function, InitExpr}; use crate::{FunctionId, FunctionKind, Global, GlobalId}; use crate::{GlobalKind, ImportKind, Memory, MemoryId, Table, TableId}; use crate::{Module, TableKind, Type, TypeId}; +/// Set of all root used items in a wasm module. +#[derive(Debug, Default)] +pub struct Roots { + tables: Vec, + funcs: Vec, + globals: Vec, + memories: Vec, + datas: Vec, + used: Used, +} + +impl Roots { + /// Creates a new set of empty roots. + pub fn new() -> Roots { + Roots::default() + } + + /// Adds a new function to the set of roots + pub fn push_func(&mut self, func: FunctionId) -> &mut Roots { + if self.used.funcs.insert(func) { + log::trace!("function is used: {:?}", func); + self.funcs.push(func); + } + self + } + + /// Adds a new table to the set of roots + pub fn push_table(&mut self, table: TableId) -> &mut Roots { + if self.used.tables.insert(table) { + log::trace!("table is used: {:?}", table); + self.tables.push(table); + } + self + } + + /// Adds a new memory to the set of roots + pub fn push_memory(&mut self, memory: MemoryId) -> &mut Roots { + if self.used.memories.insert(memory) { + log::trace!("memory is used: {:?}", memory); + self.memories.push(memory); + } + self + } + + /// Adds a new global to the set of roots + pub fn push_global(&mut self, global: GlobalId) -> &mut Roots { + if self.used.globals.insert(global) { + log::trace!("global is used: {:?}", global); + self.globals.push(global); + } + self + } + + fn push_data(&mut self, data: DataId) -> &mut Roots { + if self.used.data.insert(data) { + log::trace!("data is used: {:?}", data); + self.datas.push(data); + } + self + } +} + /// Finds the things within a module that are used. /// /// This is useful for implementing something like a linker's `--gc-sections` so @@ -32,29 +92,21 @@ pub struct Used { impl Used { /// Construct a new `Used` set for the given module. - pub fn new(module: &Module, roots: R) -> Used - where - R: IntoIterator, - { + pub fn new(module: &Module) -> Used { log::debug!("starting to calculate used set"); - let mut used = Used::default(); - let mut stack = UsedStack { - used: &mut used, - functions: Vec::new(), - tables: Vec::new(), - globals: Vec::new(), - memories: Vec::new(), - datas: Vec::new(), - }; - - for r in roots { - match module.exports.get(r).item { + let mut stack = Roots::default(); + + // All exports are roots + for export in module.exports.iter() { + match export.item { ExportItem::Function(f) => stack.push_func(f), ExportItem::Table(t) => stack.push_table(t), ExportItem::Memory(m) => stack.push_memory(m), ExportItem::Global(g) => stack.push_global(g), - } + }; } + + // The start function is an implicit root as well if let Some(f) = module.start { stack.push_func(f); } @@ -85,13 +137,19 @@ impl Used { } } + // And finally ask custom sections for their roots + for (_id, section) in module.customs.iter() { + section.gc_add_roots(&mut stack); + } + // Iteratively visit all items until our stack is empty - while stack.functions.len() > 0 + while stack.funcs.len() > 0 || stack.tables.len() > 0 || stack.memories.len() > 0 || stack.globals.len() > 0 + || stack.datas.len() > 0 { - while let Some(f) = stack.functions.pop() { + while let Some(f) = stack.funcs.pop() { let func = module.funcs.get(f); stack.used.types.insert(func.ty()); @@ -151,61 +209,15 @@ impl Used { } } - used - } -} - -struct UsedStack<'a> { - used: &'a mut Used, - functions: Vec, - tables: Vec, - memories: Vec, - globals: Vec, - datas: Vec, -} - -impl UsedStack<'_> { - fn push_func(&mut self, f: FunctionId) { - log::trace!("function is used: {:?}", f); - if self.used.funcs.insert(f) { - self.functions.push(f); - } - } - - fn push_table(&mut self, f: TableId) { - log::trace!("table is used: {:?}", f); - if self.used.tables.insert(f) { - self.tables.push(f); - } - } - - fn push_global(&mut self, f: GlobalId) { - log::trace!("global is used: {:?}", f); - if self.used.globals.insert(f) { - self.globals.push(f); - } - } - - fn push_memory(&mut self, f: MemoryId) { - log::trace!("memory is used: {:?}", f); - if self.used.memories.insert(f) { - self.memories.push(f); - } - } - - fn push_data(&mut self, d: DataId) { - log::trace!("data is used: {:?}", d); - if self.used.data.insert(d) { - self.datas.push(d); - } + stack.used } } -struct UsedVisitor<'a, 'b> { - stack: &'a mut UsedStack<'b>, +struct UsedVisitor<'a> { + stack: &'a mut Roots, } -impl<'expr> Visitor<'expr> for UsedVisitor<'_, '_> { +impl<'expr> Visitor<'expr> for UsedVisitor<'_> { fn visit_function_id(&mut self, &func: &FunctionId) { self.stack.push_func(func); }