componentize_mbt/
lib.rs

1use std::collections::HashMap;
2use std::path::PathBuf;
3use std::{fs, mem};
4
5use wast::core::{
6    Func, FuncKind, Import, InlineExport, Instruction, Memory, Module, ModuleField, ModuleKind,
7};
8use wast::parser::ParseBuffer;
9use wast::token::Index;
10use wast::Wat;
11use wit_bindgen_core::wit_parser::{Resolve, WorldId};
12use wit_bindgen_core::{Files, WorldGenerator};
13use wit_component::{embed_component_metadata, ComponentEncoder, StringEncoding};
14
15#[cfg_attr(feature = "clap", derive(clap::Args))]
16pub struct Opts {
17    #[cfg_attr(feature = "clap", arg(long))]
18    wat: PathBuf,
19}
20
21impl Opts {
22    pub fn run(
23        &self,
24        resolve: Resolve,
25        world: WorldId,
26        out_dir: Option<PathBuf>,
27    ) -> anyhow::Result<()> {
28        let wat = fs::read_to_string(&self.wat)?;
29        let wasm = componentize(&wat, resolve, world)?;
30        let target = self.wat.with_extension("wasm");
31        let target = match out_dir {
32            Some(out_dir) => out_dir.join(target.file_name().unwrap()),
33            None => target,
34        };
35        fs::write(&target, wasm)?;
36        println!("Write to {target:?}");
37        Ok(())
38    }
39}
40
41pub fn componentize(wat: &str, resolve: Resolve, world: WorldId) -> anyhow::Result<Vec<u8>> {
42    // 再跑一遍 bindgen 取回 exported_symbols,等到 MoonBit 支持自定义 FFI export 名称后,删除此 HACK
43    let mut gen = wit_bindgen_mbt::MoonBit::default();
44    gen.generate(&resolve, world, &mut Files::default())?;
45    let exported_symbols = gen.exported_symbols;
46
47    // 一些用于 lift/lower 的 ABI 函数,等到 MoonBit 能直接支持了,就可以删除了
48    let mut impls = HashMap::from([
49        (
50            "rael.memory_copy",
51            ParseBuffer::new(
52                "func $rael.memory_copy \
53                        (param $dest i32) (param $src i32) (param $len i32) \
54                        (memory.copy \
55                            (local.get $dest) (local.get $src) (local.get $len)\
56                        )",
57            )?,
58        ),
59        (
60            "rael.load_i32",
61            ParseBuffer::new(
62                "func $rael.load_i32 \
63                        (param $ptr i32) (result i32) \
64                        (i32.load (local.get $ptr))",
65            )?,
66        ),
67        (
68            "rael.load_i64",
69            ParseBuffer::new(
70                "func $rael.load_i64 \
71                        (param $ptr i32) (result i64) \
72                        (i64.load (local.get $ptr))",
73            )?,
74        ),
75        (
76            "rael.bytes_data",
77            ParseBuffer::new(
78                "func $rael.bytes_data \
79                        (param $str i32) (result i32) \
80                        (i32.add (local.get $str) (i32.const 4))",
81            )?,
82        ),
83        (
84            "moonbit.string_data",
85            ParseBuffer::new(
86                "func $moonbit.string_data \
87                        (param $str i32) (result i32) \
88                        (i32.add (local.get $str) (i32.const 4))",
89            )?,
90        ),
91        ("printc", ParseBuffer::new("func $printc (param $ptr i32)")?),
92    ]);
93    let mut builtins = HashMap::new();
94    let mut realloc = Some(ParseBuffer::new(
95        "func (export \"cabi_realloc\") \
96            (param i32) (param i32) (param i32) (param $len i32) (result i32) \
97            (call $rael.malloc (local.get $len))",
98    )?);
99
100    let buf = ParseBuffer::new(wat)?;
101    let mut ast = wast::parser::parse(&buf)?;
102    let mut start = None;
103    match &mut ast {
104        Wat::Module(Module {
105            kind: ModuleKind::Text(ref mut fields),
106            ..
107        }) => {
108            fields.retain_mut(|field| match field {
109                ModuleField::Import(Import {
110                    module: "spectest", ..
111                }) => false,
112                ModuleField::Memory(Memory {
113                    exports: InlineExport { names },
114                    ..
115                }) if names == &vec!["moonbit.memory"] => {
116                    names[0] = "memory";
117                    true
118                }
119                ModuleField::Func(Func {
120                    kind: FuncKind::Inline { expression, .. },
121                    exports,
122                    ty,
123                    ..
124                }) => {
125                    if exports.names.len() == 1 {
126                        if let Some((_, key)) = exports.names[0].split_once("::") {
127                            if let Some((symbol, has_rv)) = exported_symbols.get(key) {
128                                exports.names[0] = symbol;
129
130                                // 似乎是 MoonBit 的一个 bug —— pub fn 没有返回值,但 WASM func 却返回 i32
131                                if !has_rv {
132                                    if let Some(ty) = &mut ty.inline {
133                                        if ty.results.len() == 1 {
134                                            ty.results = Box::new([]);
135                                            let mut instrs = Default::default();
136                                            mem::swap(&mut expression.instrs, &mut instrs);
137                                            let mut instrs = instrs.into_vec();
138                                            instrs.push(Instruction::Drop);
139                                            expression.instrs = instrs.into();
140                                        }
141                                    }
142                                }
143                            }
144                        }
145                    }
146                    for instr in expression.instrs.iter() {
147                        match instr {
148                            Instruction::Call(Index::Id(id)) => {
149                                let name = id.name();
150                                if !builtins.contains_key(name) {
151                                    if let Some((name, imp)) = impls.remove_entry(name) {
152                                        builtins.insert(name, imp);
153                                    }
154                                }
155                                if name == "rael.malloc" {
156                                    if let Some(realloc) = realloc.take() {
157                                        builtins.insert(name, realloc);
158                                    }
159                                }
160                            }
161                            _ => {}
162                        }
163                    }
164                    true
165                }
166                ModuleField::Export(e) if e.name == "_start" => {
167                    start = Some(e.item);
168                    false
169                }
170                _ => true,
171            });
172
173            for buf in builtins.values() {
174                let field: ModuleField = wast::parser::parse(buf)?;
175                fields.push(field);
176            }
177            fields.push(ModuleField::Start(start.expect("_start")));
178        }
179        _ => {}
180    }
181
182    let mut buf = ast.encode()?;
183    embed_component_metadata(&mut buf, &resolve, world, StringEncoding::UTF8)?;
184    Ok(ComponentEncoder::default().module(&buf)?.encode()?)
185}