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
8use alloc::vec;
9use alloc::vec::Vec;
10use core::ops::Range;
11use core::ptr::NonNull;
12use core::slice;
13
14use anyhow::Context;
15use cranelift_entity::PrimaryMap;
16use kmem::VirtualAddress;
17
18use crate::mem::{AddressSpace, Mmap};
19use crate::wasm::TrapKind;
20use crate::wasm::compile::{CompiledFunctionInfo, FunctionLoc};
21use crate::wasm::indices::{DefinedFuncIndex, ModuleInternedTypeIndex};
22use crate::wasm::vm::{MmapVec, VMWasmCallFunction};
23
24#[derive(Debug)]
25pub struct CodeObject {
26 mmap: Mmap,
27 len: usize,
28 published: bool,
29
30 trap_offsets: Vec<u32>,
31 traps: Vec<TrapKind>,
32 wasm_to_host_trampolines: Vec<(ModuleInternedTypeIndex, FunctionLoc)>,
33 function_info: PrimaryMap<DefinedFuncIndex, CompiledFunctionInfo>,
34}
35
36impl CodeObject {
37 pub fn empty() -> Self {
38 Self {
39 mmap: Mmap::new_empty(),
40 len: 0,
41 published: false,
42 trap_offsets: vec![],
43 traps: vec![],
44 wasm_to_host_trampolines: vec![],
45 function_info: PrimaryMap::new(),
46 }
47 }
48
49 pub fn new(
50 mmap_vec: MmapVec<u8>,
51 trap_offsets: Vec<u32>,
52 traps: Vec<TrapKind>,
53 wasm_to_host_trampolines: Vec<(ModuleInternedTypeIndex, FunctionLoc)>,
54 function_info: PrimaryMap<DefinedFuncIndex, CompiledFunctionInfo>,
55 ) -> Self {
56 let (mmap, size) = mmap_vec.into_parts();
57 Self {
58 mmap,
59 len: size,
60 published: false,
61 trap_offsets,
62 traps,
63 wasm_to_host_trampolines,
64 function_info,
65 }
66 }
67
68 pub fn publish(&mut self, aspace: &mut AddressSpace) -> crate::Result<()> {
69 debug_assert!(!self.published);
70 self.published = true;
71
72 if self.mmap.is_empty() {
73 tracing::warn!("Compiled module has no code to publish");
74 return Ok(());
75 }
76
77 // Switch the executable portion from readonly to read/execute.
78 self.mmap
79 .make_executable(aspace, true)
80 .context("Failed to mark mmap'ed region as executable")?;
81
82 Ok(())
83 }
84
85 pub fn text(&self) -> &[u8] {
86 let base = self.mmap.as_ptr();
87 if base.is_null() {
88 &[]
89 } else {
90 // Safety: we have checked the slice is valid (both above and through construction)
91 unsafe { slice::from_raw_parts(self.mmap.as_ptr(), self.len) }
92 }
93 }
94
95 #[inline]
96 pub fn text_range(&self) -> Range<VirtualAddress> {
97 let start = self.mmap.range().start;
98
99 start..start.add(self.len)
100 }
101
102 pub fn resolve_function_loc(&self, func_loc: FunctionLoc) -> usize {
103 let text_range = self.text_range();
104 let addr = text_range.start.get() + func_loc.start as usize;
105
106 tracing::trace!(
107 "resolve_function_loc {func_loc:?}, text {:?} => {:?}",
108 self.mmap.as_ptr(),
109 addr,
110 );
111
112 // Assert the function location actually lies in our text section
113 debug_assert!(
114 text_range.start.get() <= addr
115 && text_range.end.get()
116 >= addr.saturating_add(usize::try_from(func_loc.length).unwrap())
117 );
118
119 addr
120 }
121
122 pub fn lookup_trap_code(&self, text_offset: usize) -> Option<TrapKind> {
123 let text_offset = u32::try_from(text_offset).unwrap();
124
125 let index = self
126 .trap_offsets
127 .binary_search_by_key(&text_offset, |val| *val)
128 .ok()?;
129
130 Some(self.traps[index])
131 }
132
133 pub(crate) fn function_info(&self) -> &PrimaryMap<DefinedFuncIndex, CompiledFunctionInfo> {
134 &self.function_info
135 }
136
137 pub fn wasm_to_host_trampoline(
138 &self,
139 sig: ModuleInternedTypeIndex,
140 ) -> NonNull<VMWasmCallFunction> {
141 let Ok(idx) = self
142 .wasm_to_host_trampolines
143 .binary_search_by_key(&sig, |entry| entry.0)
144 else {
145 panic!("missing trampoline for {sig:?}")
146 };
147
148 let (_, loc) = self.wasm_to_host_trampolines[idx];
149
150 NonNull::new(self.resolve_function_loc(loc) as *mut VMWasmCallFunction).unwrap()
151 }
152}