gobley_wasm_transformer/
lib.rs

1/*
2 * This Source Code Form is subject to the terms of the Mozilla Public
3 * License, v. 2.0. If a copy of the MPL was not distributed with this
4 * file, You can obtain one at http://mozilla.org/MPL/2.0/.
5 */
6
7pub mod import;
8pub mod stack;
9
10use std::collections::BTreeSet;
11
12use askama::Template;
13use base64::Engine;
14use walrus::{
15    ir::Value, ConstExpr, ElementItems, ElementKind, Export, ExportItem, Function, Global,
16    GlobalKind, Import, ImportKind, Module, ModuleGlobals, ValType,
17};
18
19use self::import::WasmFunctionImport;
20
21#[derive(Debug)]
22pub struct Transformer {
23    module: Module,
24    function_imports: Vec<WasmFunctionImport>,
25}
26
27#[derive(Template)]
28#[template(syntax = "kt", escape = "none", path = "js.kt")]
29pub struct KotlinJsRenderer<'a> {
30    package_name: Option<&'a str>,
31    base64: &'a str,
32    module: &'a Module,
33}
34
35impl<'a> KotlinJsRenderer<'a> {
36    fn import_modules(&self) -> Vec<String> {
37        let import_modules = self
38            .module
39            .imports
40            .iter()
41            .map(|i| &i.module)
42            .collect::<BTreeSet<_>>();
43
44        import_modules.iter().map(|i| i.to_string()).collect()
45    }
46
47    fn imports_from_module<'b>(
48        &'b self,
49        module: impl AsRef<str> + 'b,
50    ) -> impl Iterator<Item = &'a Import> + 'b {
51        self.module
52            .imports
53            .iter()
54            .filter(move |i| i.module == module.as_ref())
55    }
56
57    fn import_to_kt_signature(&self, import: &Import) -> String {
58        match import.kind {
59            ImportKind::Function(id) => self.function_to_kt_signature(self.module.funcs.get(id)),
60            ImportKind::Table(_) => "WebAssembly.Table".to_string(),
61            ImportKind::Memory(_) => "WebAssembly.Memory".to_string(),
62            ImportKind::Global(id) => Self::global_to_kt_signature(self.module.globals.get(id)),
63        }
64    }
65
66    fn import_to_function_table_entry_idx(&self, import: &Import) -> Option<usize> {
67        let ImportKind::Function(function_id) = import.kind else {
68            return None;
69        };
70
71        let Ok(Some(main_function_table)) = self.module.tables.main_function_table() else {
72            return None;
73        };
74
75        for element in self.module.elements.iter() {
76            let ElementItems::Functions(function_ids) = &element.items else {
77                continue;
78            };
79            let Some(offset) = function_ids.iter().position(|id| *id == function_id) else {
80                continue;
81            };
82            let ElementKind::Active {
83                table,
84                offset: element_offset,
85            } = &element.kind
86            else {
87                continue;
88            };
89            if main_function_table != *table {
90                continue;
91            }
92
93            fn get_usize_from_constexpr(
94                globals: &ModuleGlobals,
95                expr: &ConstExpr,
96            ) -> Option<usize> {
97                Some(match expr {
98                    ConstExpr::Value(value) => match value {
99                        Value::I32(i32) => *i32 as usize,
100                        Value::I64(i64) => *i64 as usize,
101                        Value::F32(f32) => *f32 as usize,
102                        Value::F64(f64) => *f64 as usize,
103                        Value::V128(v128) => *v128 as usize,
104                    },
105                    ConstExpr::Global(id) => {
106                        return match &globals.get(*id).kind {
107                            GlobalKind::Local(expr) => get_usize_from_constexpr(globals, expr),
108                            _ => None,
109                        }
110                    }
111                    _ => return None,
112                })
113            }
114
115            let Some(element_offset) =
116                get_usize_from_constexpr(&self.module.globals, element_offset)
117            else {
118                continue;
119            };
120
121            return Some(offset + element_offset);
122        }
123
124        None
125    }
126
127    fn exports(&self) -> impl Iterator<Item = &Export> {
128        self.module.exports.iter()
129    }
130
131    fn export_to_kt_signature(&self, export: &Export) -> String {
132        match export.item {
133            ExportItem::Function(id) => self.function_to_kt_signature(self.module.funcs.get(id)),
134            ExportItem::Table(_) => "WebAssembly.Table".to_string(),
135            ExportItem::Memory(_) => "WebAssembly.Memory".to_string(),
136            ExportItem::Global(id) => Self::global_to_kt_signature(self.module.globals.get(id)),
137        }
138    }
139
140    fn function_to_kt_signature(&self, function: &Function) -> String {
141        let ty = self.module.types.get(function.ty());
142        let mut output = String::new();
143        let mut first = true;
144        output.push('(');
145
146        for param_str in ty.params().iter().map(Self::map_val_type_to_kt) {
147            if !first {
148                output.push_str(", ");
149            }
150            first = false;
151            output.push_str(param_str);
152        }
153
154        output.push_str(") -> ");
155
156        if let Some(result) = ty.results().first() {
157            output.push_str(Self::map_val_type_to_kt(result));
158        } else {
159            output.push_str("Unit");
160        }
161
162        output
163    }
164
165    fn global_to_kt_signature(global: &Global) -> String {
166        let inner_ty = Self::map_val_type_to_kt(&global.ty);
167        format!("WebAssembly.Global<{inner_ty}>")
168    }
169
170    fn map_val_type_to_kt(ty: &ValType) -> &'static str {
171        match ty {
172            ValType::I32 => "Int",
173            ValType::F32 => "Float",
174            ValType::F64 => "Double",
175            _ => "Any",
176        }
177    }
178}
179
180impl Transformer {
181    pub fn new(input: &[u8], function_imports: Vec<WasmFunctionImport>) -> anyhow::Result<Self> {
182        Ok(Self {
183            module: Module::from_buffer(input)?,
184            function_imports,
185        })
186    }
187
188    fn transform(&mut self) -> anyhow::Result<()> {
189        self.inject_stack_pointer_shim()?;
190        self.inject_function_imports();
191        Ok(())
192    }
193
194    pub fn render_into_kt(mut self, package_name: Option<&str>) -> anyhow::Result<String> {
195        use base64::prelude::BASE64_STANDARD;
196
197        self.transform()?;
198
199        let wasm = self.module.emit_wasm();
200        let wasm_base64 = BASE64_STANDARD.encode(&wasm);
201        let module = Module::from_buffer(&wasm)?;
202        let renderer = KotlinJsRenderer {
203            package_name,
204            base64: &wasm_base64,
205            module: &module,
206        };
207        Ok(renderer.render()?)
208    }
209}
gobley_wasm_transformer/lib.rs

gobley_wasm_transformer/
lib.rs
