use eyre::{bail, Result};
use std::collections::HashMap;
use wasmparser::{
BinaryReader, Imports, NameSectionReader, Payload, Subsection, Subsections, TypeRef,
};
pub use wasmparser::{
Data, Element, Export, FuncType, FunctionBody, Global, Import, MemoryType, Table, TagType,
};
use crate::{
dep_graph::DepNode,
magic_constants::{
FeatureTag as WsFeature, SubsectionTag as WsSubsectionTag, Version as WsVersion,
},
reloc::{RelocInfo, RelocInfoParser},
};
pub struct CustomSection<'a> {
pub name: &'a str,
pub data_offset: usize,
pub data: &'a [u8],
}
pub type FuncTypeId = usize;
pub type InputFuncId = usize;
pub type TableId = usize;
pub type ImportId = usize;
pub type ExportId = usize;
pub type MemoryId = usize;
pub type GlobalId = usize;
pub type ElementId = usize;
pub type DataSegmentId = usize;
pub type TagId = usize;
#[derive(Debug)]
pub struct ImportedFunc {
pub type_id: FuncTypeId,
pub import_id: ImportId,
}
#[derive(Debug)]
pub struct DefinedFunc<'a> {
pub type_id: FuncTypeId,
pub body: FunctionBody<'a>,
}
#[derive(Default, Clone)]
pub struct Names<'a> {
pub module: Option<&'a str>,
pub functions: HashMap<InputFuncId, &'a str>,
pub locals: HashMap<InputFuncId, wasmparser::NameMap<'a>>,
pub labels: HashMap<InputFuncId, wasmparser::NameMap<'a>>,
pub types: HashMap<FuncTypeId, &'a str>,
pub tables: HashMap<TableId, &'a str>,
pub memories: HashMap<MemoryId, &'a str>,
pub globals: HashMap<GlobalId, &'a str>,
pub elements: HashMap<ElementId, &'a str>,
pub data_segments: HashMap<DataSegmentId, &'a str>,
pub tags: HashMap<TagId, &'a str>,
}
fn convert_name_map<'a>(name_map: wasmparser::NameMap<'a>) -> Result<HashMap<usize, &'a str>> {
name_map
.into_iter()
.map(|r| r.map(|naming| (naming.index as usize, naming.name)))
.collect::<Result<HashMap<usize, &'a str>, _>>()
.map_err(|e| e.into())
}
fn convert_indirect_name_map<'a>(
indirect_name_map: wasmparser::IndirectNameMap<'a>,
) -> Result<HashMap<usize, wasmparser::NameMap<'a>>> {
indirect_name_map
.into_iter()
.map(|r| -> Result<(usize, wasmparser::NameMap<'a>)> {
let indirect_naming = r?;
Ok((indirect_naming.index as usize, indirect_naming.names))
})
.collect::<Result<HashMap<_, _>, _>>()
}
impl<'a> Names<'a> {
fn new(data: &'a [u8], original_offset: usize) -> Result<Self> {
let mut names: Self = Default::default();
for part in NameSectionReader::new(BinaryReader::new(data, original_offset)) {
use wasmparser::Name;
match part? {
Name::Module { name, .. } => {
names.module = Some(name);
}
Name::Function(name_map) => {
names.functions = convert_name_map(name_map)?;
}
Name::Local(indirect_name_map) => {
names.locals = convert_indirect_name_map(indirect_name_map)?;
}
Name::Label(indirect_name_map) => {
names.labels = convert_indirect_name_map(indirect_name_map)?;
}
Name::Type(name_map) => {
names.types = convert_name_map(name_map)?;
}
Name::Table(name_map) => {
names.tables = convert_name_map(name_map)?;
}
Name::Memory(name_map) => {
names.memories = convert_name_map(name_map)?;
}
Name::Global(name_map) => {
names.globals = convert_name_map(name_map)?;
}
Name::Data(name_map) => {
names.data_segments = convert_name_map(name_map)?;
}
Name::Element(name_map) => {
names.elements = convert_name_map(name_map)?;
}
Name::Tag(name_map) => {
names.tags = convert_name_map(name_map)?;
}
Name::Field(_name_map) => {
bail!("Field names not supported");
}
Name::Unknown { ty, .. } => {
bail!("Unknown name subsection: {:?}", ty);
}
}
}
Ok(names)
}
}
pub type InputOffset = usize;
pub use crate::reloc::SymbolIndex;
#[derive(Default)]
pub struct InputModule<'a> {
pub raw: &'a [u8],
pub types: Vec<FuncType>,
pub imports: Vec<Import<'a>>,
pub imported_tags_num: usize,
pub tags: Vec<TagType>,
pub imported_globals_num: usize,
pub globals: Vec<Global<'a>>,
pub imported_memories_num: usize,
pub memories: Vec<MemoryType>,
pub imported_tables_num: usize,
pub tables: Vec<Table<'a>>,
pub start: Option<InputFuncId>,
pub elements: Vec<Element<'a>>,
pub data_segments: Vec<Data<'a>>,
pub imported_funcs: Vec<ImportedFunc>,
pub defined_funcs: Vec<DefinedFunc<'a>>,
pub exports: Vec<Export<'a>>,
pub custom_sections: Vec<CustomSection<'a>>,
pub names: Names<'a>,
pub imported_func_map: HashMap<ImportId, InputFuncId>,
pub reloc_info: RelocInfo<'a>,
pub options: Options,
}
pub enum Strictness {
IntegrationTesting,
Lenient,
}
impl<'a> InputModule<'a> {
pub fn parse(wasm: &'a [u8], strict: Strictness) -> Result<Self> {
let mut module = Self {
raw: wasm,
..Default::default()
};
let mut reloc_info = RelocInfoParser::default();
let mut function_types: Vec<FuncTypeId> = Vec::new();
let parser = wasmparser::Parser::new(0);
for payload in parser.parse_all(wasm) {
let payload = payload?;
reloc_info.visit_payload(&payload)?;
match payload {
Payload::Version { .. } => {}
Payload::TypeSection(reader) => {
module.types = reader
.into_iter_err_on_gc_types()
.collect::<Result<Vec<_>, _>>()?;
}
Payload::ImportSection(reader) => {
let gathered = &mut module.imports;
for imports in reader.into_iter() {
match imports? {
Imports::Single(_, import) => gathered.push(import),
Imports::Compact1 { module, items } => {
for item in items.into_iter() {
let wasmparser::ImportItemCompact { name, ty } = item?;
gathered.push(Import { module, name, ty })
}
}
Imports::Compact2 { module, ty, names } => {
for name in names.into_iter() {
let name = name?;
gathered.push(Import { module, name, ty })
}
}
}
}
}
Payload::FunctionSection(reader) => {
function_types = reader
.into_iter()
.map(|t| t.map(|id| id as FuncTypeId))
.collect::<Result<Vec<_>, _>>()?;
}
Payload::TableSection(reader) => {
module.tables = reader.into_iter().collect::<Result<Vec<_>, _>>()?;
}
Payload::MemorySection(reader) => {
module.memories = reader.into_iter().collect::<Result<Vec<_>, _>>()?;
}
Payload::TagSection(reader) => {
module.tags = reader.into_iter().collect::<Result<Vec<_>, _>>()?;
}
Payload::GlobalSection(reader) => {
module.globals = reader.into_iter().collect::<Result<Vec<_>, _>>()?;
}
Payload::ExportSection(reader) => {
module.exports = reader.into_iter().collect::<Result<Vec<_>, _>>()?;
}
Payload::StartSection { func, .. } => {
module.start = Some(func as usize);
}
Payload::ElementSection(reader) => {
module.elements = reader.into_iter().collect::<Result<Vec<_>, _>>()?;
}
Payload::DataCountSection { .. } => {}
Payload::DataSection(reader) => {
module.data_segments = reader.into_iter().collect::<Result<Vec<_>, _>>()?;
}
Payload::CodeSectionStart { .. } => {}
Payload::CodeSectionEntry(body) => {
let index = module.defined_funcs.len();
module.defined_funcs.push(DefinedFunc {
type_id: function_types[index],
body,
});
}
Payload::CustomSection(reader) => {
module.custom_sections.push(CustomSection {
name: reader.name(),
data: reader.data(),
data_offset: reader.data_offset(),
});
}
Payload::End(_) => {}
section => {
bail!("Unknown section: {:?}", section);
}
}
}
let mut num_split_sections_found = 0;
for section in module.custom_sections.iter() {
if section.name == "name" {
module.names = Names::new(section.data, section.data_offset)?;
}
if section.name == crate::magic_constants::LINK_SECTION {
let rdr = BinaryReader::new(section.data, section.data_offset);
let () = read_wasm_split_section(rdr, &mut module.options)?;
num_split_sections_found += 1;
}
}
if matches!(strict, Strictness::IntegrationTesting) && num_split_sections_found != 1 {
bail!(
"Wrong number of custom sections for wasm-split found, expected 1 got {}",
num_split_sections_found
);
}
module.reloc_info = reloc_info.finish(&module)?;
for (import_id, import) in module.imports.iter().enumerate() {
let import_id = import_id as ImportId;
match import.ty {
TypeRef::Func(func_type_id) | TypeRef::FuncExact(func_type_id) => {
let func_id = module.imported_funcs.len();
module.imported_funcs.push(ImportedFunc {
import_id,
type_id: func_type_id as usize,
});
module.imported_func_map.insert(import_id, func_id);
}
TypeRef::Table(_table_type) => module.imported_tables_num += 1,
TypeRef::Memory(_memory_type) => module.imported_memories_num += 1,
TypeRef::Global(_global_type) => module.imported_globals_num += 1,
TypeRef::Tag(_tag_type) => module.imported_tags_num += 1,
}
}
Ok(module)
}
pub fn func_type_id(&self, func_id: InputFuncId) -> FuncTypeId {
if func_id < self.imported_funcs.len() {
self.imported_funcs[func_id].type_id
} else {
self.defined_funcs[func_id - self.imported_funcs.len()].type_id
}
}
pub fn main_memory(&self) -> DepNode {
DepNode::Memory(0)
}
pub fn indirect_function_table(&self) -> DepNode {
DepNode::Table(self.reloc_info.indirect_table)
}
}
enum WsReadVersion {
Known(WsVersion),
Unknown,
}
enum WsReadFeature {
Known(WsFeature),
Unknown,
}
enum WasmSplitSubsection<'a> {
Version(WsReadVersion),
Feature(WsReadFeature),
Unknown(BinaryReader<'a>),
}
fn read_version(reader: &mut BinaryReader<'_>) -> wasmparser::Result<WsReadVersion> {
let version = reader.read_u8()?;
Ok(match version {
vers if vers == (WsVersion::Version1 as u8) => WsReadVersion::Known(WsVersion::Version1),
_ => WsReadVersion::Unknown,
})
}
fn read_feature(reader: &mut BinaryReader<'_>) -> wasmparser::Result<WsReadFeature> {
let feature = reader.read_u8()?;
Ok(match feature {
feat if feat == WsFeature::CfgDebugAssertions as u8 => {
WsReadFeature::Known(WsFeature::CfgDebugAssertions)
}
_ => WsReadFeature::Unknown,
})
}
impl<'a> Subsection<'a> for WasmSplitSubsection<'a> {
fn from_reader(id: u8, mut reader: BinaryReader<'a>) -> wasmparser::Result<Self> {
match id {
id if id == WsSubsectionTag::Version as u8 => {
let version = read_version(&mut reader)?;
Ok(Self::Version(version))
}
id if id == WsSubsectionTag::Feature as u8 => {
let feature = read_feature(&mut reader)?;
Ok(Self::Feature(feature))
}
_ => Ok(Self::Unknown(reader)),
}
}
}
#[derive(Default)]
pub struct Options {
pub debug_assertions: bool,
}
fn read_wasm_split_section(rdr: BinaryReader<'_>, options: &mut Options) -> Result<()> {
let subsections = Subsections::<WasmSplitSubsection>::new(rdr);
for subsection in subsections {
match subsection? {
WasmSplitSubsection::Version(WsReadVersion::Unknown) => bail!(
r#"Input file was linked with a `wasm_split_helpers` version that is not recognized by this tool. \
Please upgrade your `wasm_split_cli` version and ensure you have the latest version of the relevant \
build tools installed."#
),
WasmSplitSubsection::Version(WsReadVersion::Known(WsVersion::Version1)) => {
}
WasmSplitSubsection::Feature(WsReadFeature::Unknown) => {}
WasmSplitSubsection::Feature(WsReadFeature::Known(WsFeature::CfgDebugAssertions)) => {
options.debug_assertions = true;
}
WasmSplitSubsection::Unknown(reader) => {
let _ = reader;
}
}
}
Ok(())
}