1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148

use anyhow::{anyhow, Context, Result};
use std::borrow::Cow;
use std::collections::HashMap;
use walrus::{passes::Roots, CustomSection, IdsToIndices, IndicesToIds, Module};
use wit_schema_version::SECTION_NAME;

#[derive(Debug, Default)]
pub struct WasmInterfaceTypes {
    pub types: Types,
    pub imports: Imports,
    pub implements: Implements,
    pub exports: Exports,
    pub funcs: Funcs,
}

mod exports;
mod funcs;
mod implements;
mod imports;
mod types;
pub use self::exports::*;
pub use self::funcs::*;
pub use self::implements::*;
pub use self::imports::*;
pub use self::types::*;

impl CustomSection for WasmInterfaceTypes {
    fn name(&self) -> &str {
        SECTION_NAME
    }

    fn data(&self, indices: &IdsToIndices) -> Cow<'_, [u8]> {
        let mut writer = wit_writer::Writer::new();
        let mut wids = WitIdsToIndices::default();
        self.encode_types(&mut writer, &mut wids);
        self.encode_imports(&mut writer, &mut wids);
        self.encode_funcs(&mut writer, &mut wids, indices);
        self.encode_exports(&mut writer, &wids);
        self.encode_implements(&mut writer, &wids, indices);
        writer.into_payload().into()
    }

    fn add_gc_roots(&self, roots: &mut Roots) {
        for i in self.implements.iter() {
            roots.push_func(i.core_func);
        }
        for f in self.funcs.iter() {
            let instrs = match &f.kind {
                FuncKind::Local(instrs) => instrs,
                _ => continue,
            };
            for instr in instrs {
                match instr {
                    Instruction::CallCore(f) | Instruction::DeferCallCore(f) => {
                        roots.push_func(*f);
                    }
                    Instruction::MemoryToString(mem) => {
                        roots.push_memory(*mem);
                    }
                    Instruction::StringToMemory { mem, malloc } => {
                        roots.push_memory(*mem).push_func(*malloc);
                    }
                    _ => {}
                }
            }
        }
    }
}

#[derive(Default)]
struct WitIndicesToIds {
    types: Vec<TypeId>,
    funcs: Vec<FuncId>,
}

impl WitIndicesToIds {
    fn ty(&self, ty: u32) -> Result<TypeId> {
        self.types
            .get(ty as usize)
            .cloned()
            .ok_or_else(|| anyhow!("adapter type index out of bounds: {}", ty))
    }

    fn func(&self, ty: u32) -> Result<FuncId> {
        self.funcs
            .get(ty as usize)
            .cloned()
            .ok_or_else(|| anyhow!("adapter func index out of bounds: {}", ty))
    }
}

#[derive(Default)]
struct WitIdsToIndices {
    types: HashMap<TypeId, u32>,
    funcs: HashMap<FuncId, u32>,
}

impl WitIdsToIndices {
    fn push_ty(&mut self, ty: TypeId) {
        self.types.insert(ty, self.types.len() as u32);
    }

    fn ty(&self, ty: TypeId) -> u32 {
        self.types
            .get(&ty)
            .cloned()
            .unwrap_or_else(|| panic!("reference to dead type found {:?}", ty))
    }

    fn push_func(&mut self, func: FuncId) {
        self.funcs.insert(func, self.funcs.len() as u32);
    }

    fn func(&self, f: FuncId) -> u32 {
        self.funcs
            .get(&f)
            .cloned()
            .unwrap_or_else(|| panic!("reference to dead function found {:?}", f))
    }
}

/// Callback for the `ModuleConfig::on_parse` function in `walrus` to act as a
/// convenience to parse the wasm interface types custom section, if present.
pub fn on_parse(module: &mut Module, ids: &IndicesToIds) -> Result<()> {
    let section = match module.customs.remove_raw(SECTION_NAME) {
        Some(s) => s,
        None => return Ok(()),
    };
    let mut parser = wit_parser::Parser::new(0, &section.data)
        .context("failed parsing wasm interface types header")?;
    let mut section = WasmInterfaceTypes::default();
    let mut wids = WitIndicesToIds::default();
    while !parser.is_empty() {
        let s = parser
            .section()
            .context("failed parsing wasm interface types section header")?;
        match s {
            wit_parser::Section::Type(t) => section.parse_types(t, &mut wids)?,
            wit_parser::Section::Import(t) => section.parse_imports(t, &mut wids)?,
            wit_parser::Section::Func(t) => section.parse_funcs(t, ids, &mut wids)?,
            wit_parser::Section::Implement(t) => section.parse_implements(t, ids, &mut wids)?,
            wit_parser::Section::Export(t) => section.parse_exports(t, &mut wids)?,
        }
    }

    module.customs.add(section);
    Ok(())
}