jonaskruckenberg.de
Next Generation WASM Microkernel Operating System
1// Copyright 2025 Jonas Kruckenberg
2//
3// Licensed under the Apache License, Version 2.0, <LICENSE-APACHE or
4// http://apache.org/licenses/LICENSE-2.0> or the MIT license <LICENSE-MIT or
5// http://opensource.org/licenses/MIT>, at your option. This file may not be
6// copied, modified, or distributed except according to those terms.
7
8mod compile_key;
9mod compiled_function;
10
11use crate::wasm::Engine;
12use crate::wasm::builtins::BuiltinFunctionIndex;
13use crate::wasm::compile::compiled_function::{RelocationTarget, TrapInfo};
14use crate::wasm::indices::{DefinedFuncIndex, ModuleInternedTypeIndex};
15use crate::wasm::translate::{
16 FunctionBodyData, ModuleTranslation, ModuleTypes, TranslatedModule, WasmFuncType,
17};
18use crate::wasm::trap::TrapKind;
19use crate::wasm::vm::{CodeObject, MmapVec};
20use alloc::boxed::Box;
21use alloc::collections::BTreeMap;
22use alloc::format;
23use alloc::string::String;
24use alloc::vec::Vec;
25use compile_key::CompileKey;
26pub use compiled_function::CompiledFunction;
27use cranelift_codegen::control::ControlPlane;
28use cranelift_entity::{EntitySet, PrimaryMap};
29use hashbrown::HashSet;
30
31/// Namespace corresponding to wasm functions, the index is the index of the
32/// defined function that's being referenced.
33pub const NS_WASM_FUNC: u32 = 0;
34
35/// Namespace for builtin function trampolines. The index is the index of the
36/// builtin that's being referenced.
37pub const NS_BUILTIN: u32 = 1;
38
39pub trait Compiler {
40 /// Returns the target triple this compiler is configured for
41 fn triple(&self) -> &target_lexicon::Triple;
42
43 fn text_section_builder(
44 &self,
45 capacity: usize,
46 ) -> Box<dyn cranelift_codegen::TextSectionBuilder>;
47
48 /// Compile the translated WASM function `index` within `translation`.
49 fn compile_function(
50 &self,
51 translation: &ModuleTranslation<'_>,
52 index: DefinedFuncIndex,
53 data: FunctionBodyData<'_>,
54 types: &ModuleTypes,
55 ) -> crate::Result<CompiledFunction>;
56
57 /// Compile a trampoline for calling the `index` WASM function through the
58 /// array-calling convention used by host code to call into WASM.
59 fn compile_array_to_wasm_trampoline(
60 &self,
61 translation: &ModuleTranslation<'_>,
62 types: &ModuleTypes,
63 index: DefinedFuncIndex,
64 ) -> crate::Result<CompiledFunction>;
65
66 // Compile a trampoline for calling the a(host-defined) function through the array
67 /// calling convention used to by WASM to call into host code.
68 fn compile_wasm_to_array_trampoline(
69 &self,
70 wasm_func_ty: &WasmFuncType,
71 ) -> crate::Result<CompiledFunction>;
72
73 /// Compile a trampoline for calling the `index` builtin function from WASM.
74 fn compile_wasm_to_builtin(
75 &self,
76 index: BuiltinFunctionIndex,
77 ) -> crate::Result<CompiledFunction>;
78}
79
80/// A position within an original source file,
81#[derive(Clone, Copy, Debug, PartialEq, Eq)]
82pub struct FilePos(u32);
83
84impl Default for FilePos {
85 fn default() -> Self {
86 Self(u32::MAX)
87 }
88}
89
90impl FilePos {
91 pub(crate) fn new(pos: u32) -> Self {
92 Self(pos)
93 }
94
95 pub fn file_offset(self) -> Option<u32> {
96 if self.0 == u32::MAX {
97 None
98 } else {
99 Some(self.0)
100 }
101 }
102}
103
104#[derive(Debug)]
105pub struct CompiledFunctionInfo {
106 /// The [`FunctionLoc`] indicating the location of this function in the text
107 /// section of the compilation artifact.
108 pub wasm_func_loc: FunctionLoc,
109 /// A trampoline for host callers (e.g. `Func::wrap`) calling into this function (if needed).
110 pub array_to_wasm_trampoline: Option<FunctionLoc>,
111 pub start_srcloc: FilePos,
112}
113
114/// Description of where a function is located in the text section of a
115/// compiled image.
116#[derive(Debug, Copy, Clone)]
117pub struct FunctionLoc {
118 /// The byte offset from the start of the text section where this
119 /// function starts.
120 pub start: u32,
121 /// The byte length of this function's function body.
122 pub length: u32,
123}
124
125pub type CompileInput<'a> =
126 Box<dyn FnOnce(&dyn Compiler) -> crate::Result<CompileOutput> + Send + 'a>;
127
128pub struct CompileInputs<'a>(Vec<CompileInput<'a>>);
129
130impl<'a> CompileInputs<'a> {
131 pub fn from_module(
132 translation: &'a ModuleTranslation,
133 types: &'a ModuleTypes,
134 function_body_data: PrimaryMap<DefinedFuncIndex, FunctionBodyData<'a>>,
135 ) -> Self {
136 let mut inputs: Vec<CompileInput> = Vec::new();
137
138 for (def_func_index, function_body_data) in function_body_data {
139 // push the "main" function compilation job
140 inputs.push(Box::new(move |compiler| {
141 let symbol = format!("wasm[0]::function[{}]", def_func_index.as_u32());
142 tracing::debug!("compiling {symbol}...");
143
144 let function = compiler.compile_function(
145 translation,
146 def_func_index,
147 function_body_data,
148 types,
149 )?;
150
151 Ok(CompileOutput {
152 key: CompileKey::wasm_function(def_func_index),
153 function,
154 symbol,
155 })
156 }));
157
158 // Compile a host->wasm trampoline for every function that are flags as "escaping"
159 // and could therefore theoretically be called by native code.
160 let func_index = translation.module.func_index(def_func_index);
161 if translation.module.functions[func_index].is_escaping() {
162 inputs.push(Box::new(move |compiler| {
163 let symbol =
164 format!("wasm[0]::array_to_wasm_trampoline[{}]", func_index.as_u32());
165 tracing::debug!("compiling {symbol}...");
166
167 let function = compiler.compile_array_to_wasm_trampoline(
168 translation,
169 types,
170 def_func_index,
171 )?;
172
173 Ok(CompileOutput {
174 key: CompileKey::array_to_wasm_trampoline(def_func_index),
175 function,
176 symbol,
177 })
178 }));
179 }
180 }
181
182 let mut trampoline_types_seen = HashSet::new();
183 for (_func_type_index, trampoline_type_index) in types.trampoline_types() {
184 let is_new = trampoline_types_seen.insert(trampoline_type_index);
185 if !is_new {
186 continue;
187 }
188
189 let trampoline_func_ty = types
190 .get_wasm_type(trampoline_type_index)
191 .unwrap()
192 .unwrap_func();
193 inputs.push(Box::new(move |compiler| {
194 let symbol = format!(
195 "signatures[{}]::wasm_to_array_trampoline",
196 trampoline_type_index.as_u32()
197 );
198 tracing::debug!("compiling {symbol}...");
199
200 let function = compiler.compile_wasm_to_array_trampoline(trampoline_func_ty)?;
201
202 Ok(CompileOutput {
203 key: CompileKey::wasm_to_array_trampoline(trampoline_type_index),
204 function,
205 symbol,
206 })
207 }));
208 }
209
210 Self(inputs)
211 }
212
213 pub fn compile(self, compiler: &dyn Compiler) -> crate::Result<UnlinkedCompileOutputs> {
214 let mut outputs = self
215 .0
216 .into_iter()
217 .map(|f| f(compiler))
218 .collect::<Result<Vec<_>, _>>()?;
219
220 compile_required_builtin_trampolines(compiler, &mut outputs)?;
221
222 let mut indices: BTreeMap<u32, BTreeMap<CompileKey, usize>> = BTreeMap::new();
223 for (index, output) in outputs.iter().enumerate() {
224 indices
225 .entry(output.key.kind())
226 .or_default()
227 .insert(output.key, index);
228 }
229
230 Ok(UnlinkedCompileOutputs { indices, outputs })
231 }
232}
233
234fn compile_required_builtin_trampolines(
235 compiler: &dyn Compiler,
236 outputs: &mut Vec<CompileOutput>,
237) -> crate::Result<()> {
238 let mut builtins = EntitySet::new();
239 let mut new_jobs: Vec<CompileInput<'_>> = Vec::new();
240
241 let builtin_indices = outputs
242 .iter()
243 .flat_map(|output| output.function.relocations())
244 .filter_map(|reloc| match reloc.target {
245 RelocationTarget::Wasm(_) => None,
246 RelocationTarget::Builtin(index) => Some(index),
247 });
248
249 let compile_builtin = |builtin: BuiltinFunctionIndex| -> CompileInput {
250 Box::new(move |compiler: &dyn Compiler| {
251 let symbol = format!("wasm_builtin_{}", builtin.name());
252 tracing::debug!("compiling {symbol}...");
253 Ok(CompileOutput {
254 key: CompileKey::wasm_to_builtin_trampoline(builtin),
255 symbol,
256 function: compiler.compile_wasm_to_builtin(builtin)?,
257 })
258 })
259 };
260
261 for index in builtin_indices {
262 if builtins.insert(index) {
263 new_jobs.push(compile_builtin(index));
264 }
265 }
266
267 outputs.extend(
268 new_jobs
269 .into_iter()
270 .map(|f| f(compiler))
271 .collect::<Result<Vec<_>, _>>()?,
272 );
273
274 Ok(())
275}
276
277#[derive(Debug)]
278pub struct CompileOutput {
279 pub key: CompileKey,
280 pub function: CompiledFunction,
281 pub symbol: String,
282}
283
284#[derive(Debug)]
285pub struct UnlinkedCompileOutputs {
286 indices: BTreeMap<u32, BTreeMap<CompileKey, usize>>,
287 outputs: Vec<CompileOutput>,
288}
289
290impl UnlinkedCompileOutputs {
291 pub fn link_and_finish(
292 mut self,
293 engine: &Engine,
294 module: &TranslatedModule,
295 do_mmap: impl FnOnce(Vec<u8>) -> crate::Result<MmapVec<u8>>,
296 ) -> crate::Result<CodeObject> {
297 let mut text_builder = engine.compiler().text_section_builder(self.outputs.len());
298 let mut ctrl_plane = ControlPlane::default();
299 let mut locs = Vec::new(); // TODO get a capacity value for this
300 let mut traps = TrapsBuilder::default();
301
302 for output in &self.outputs {
303 let body = output.function.buffer();
304 let alignment = output.function.alignment();
305 let body_len = body.len() as u64;
306 let off = text_builder.append(true, body, alignment, &mut ctrl_plane);
307
308 tracing::debug!(
309 "Function {}: {off:#x}..{:#x} align: {alignment}",
310 output.symbol,
311 off + body_len
312 );
313
314 for r in output.function.relocations() {
315 let target = match r.target {
316 RelocationTarget::Wasm(callee_index) => {
317 let def_func_index = module.defined_func_index(callee_index).unwrap();
318
319 self.indices[&CompileKey::WASM_FUNCTION_KIND]
320 [&CompileKey::wasm_function(def_func_index)]
321 }
322 RelocationTarget::Builtin(index) => {
323 self.indices[&CompileKey::WASM_TO_BUILTIN_TRAMPOLINE_KIND]
324 [&CompileKey::wasm_to_builtin_trampoline(index)]
325 }
326 };
327
328 // Ensure that we actually resolved the relocation
329 let resolved =
330 text_builder.resolve_reloc(off + u64::from(r.offset), r.kind, r.addend, target);
331 debug_assert!(resolved);
332 }
333
334 let loc = FunctionLoc {
335 start: u32::try_from(off).unwrap(),
336 length: u32::try_from(body_len).unwrap(),
337 };
338
339 traps.push_traps(loc, output.function.traps());
340 locs.push(loc);
341 }
342
343 let wasm_functions = self
344 .indices
345 .remove(&CompileKey::WASM_FUNCTION_KIND)
346 .unwrap_or_default()
347 .into_iter();
348
349 let mut host_to_wasm_trampolines = self
350 .indices
351 .remove(&CompileKey::ARRAY_TO_WASM_TRAMPOLINE_KIND)
352 .unwrap_or_default();
353
354 let funcs = wasm_functions
355 .map(|(key, index)| {
356 let host_to_wasm_trampoline_key =
357 CompileKey::array_to_wasm_trampoline(DefinedFuncIndex::from_u32(key.index));
358 let host_to_wasm_trampoline = host_to_wasm_trampolines
359 .remove(&host_to_wasm_trampoline_key)
360 .map(|index| locs[index]);
361
362 CompiledFunctionInfo {
363 start_srcloc: self.outputs[index].function.metadata().start_srcloc,
364 wasm_func_loc: locs[index],
365 array_to_wasm_trampoline: host_to_wasm_trampoline,
366 }
367 })
368 .collect();
369
370 // collect up all the WASM->host trampolines we compiled
371 let wasm_to_host_trampolines: Vec<_> = self
372 .indices
373 .remove(&CompileKey::WASM_TO_ARRAY_TRAMPOLINE_KIND)
374 .unwrap_or_default()
375 .into_iter()
376 .map(|(key, index)| (ModuleInternedTypeIndex::from_u32(key.index), locs[index]))
377 .collect();
378
379 let (trap_offsets, traps) = traps.finish();
380 let mmap_vec = do_mmap(text_builder.finish(&mut ctrl_plane))?;
381
382 Ok(CodeObject::new(
383 mmap_vec,
384 trap_offsets,
385 traps,
386 wasm_to_host_trampolines,
387 funcs,
388 ))
389 }
390}
391
392#[derive(Default)]
393struct TrapsBuilder {
394 offsets: Vec<u32>,
395 traps: Vec<TrapKind>,
396 last_offset: u32,
397}
398
399impl TrapsBuilder {
400 pub fn push_traps(
401 &mut self,
402 func: FunctionLoc,
403 traps: impl ExactSizeIterator<Item = TrapInfo>,
404 ) {
405 self.offsets.reserve_exact(traps.len());
406 self.traps.reserve_exact(traps.len());
407
408 for trap in traps {
409 let pos = func.start + trap.offset;
410 debug_assert!(pos >= self.last_offset);
411 // sanity check to make sure everything is sorted.
412 // otherwise we won't be able to use lookup later.
413 self.offsets.push(pos);
414 self.traps.push(trap.trap);
415 self.last_offset = pos;
416 }
417
418 self.last_offset = func.start + func.length;
419 }
420
421 pub fn finish(self) -> (Vec<u32>, Vec<TrapKind>) {
422 (self.offsets, self.traps)
423 }
424}