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 crate::wasm::vm::CodeObject;
9use alloc::collections::BTreeMap;
10use alloc::sync::Arc;
11use spin::{OnceLock, RwLock};
12
13fn global_code() -> &'static RwLock<GlobalRegistry> {
14 static GLOBAL_CODE: OnceLock<RwLock<GlobalRegistry>> = OnceLock::new();
15 GLOBAL_CODE.get_or_init(Default::default)
16}
17
18type GlobalRegistry = BTreeMap<usize, (usize, Arc<CodeObject>)>;
19
20/// Find which registered region of code contains the given program counter, and
21/// what offset that PC is within that module's code.
22pub fn lookup_code(pc: usize) -> Option<(Arc<CodeObject>, usize)> {
23 let all_modules = global_code().read();
24 let (_end, (start, module)) = all_modules.range(pc..).next()?;
25 let text_offset = pc.checked_sub(*start)?;
26 Some((module.clone(), text_offset))
27}
28
29/// Registers a new region of code.
30///
31/// Must not have been previously registered and must be `unregister`'d to
32/// prevent leaking memory.
33///
34/// This is required to enable traps to work correctly since the signal handler
35/// will lookup in the `GLOBAL_CODE` list to determine which a particular pc
36/// is a trap or not.
37pub fn register_code(code: &Arc<CodeObject>) {
38 let text = code.text();
39 if text.is_empty() {
40 return;
41 }
42 let start = text.as_ptr() as usize;
43 let end = start + text.len() - 1;
44 let prev = global_code().write().insert(end, (start, code.clone()));
45 assert!(prev.is_none());
46}
47
48/// Unregisters a code mmap from the global map.
49///
50/// Must have been previously registered with `register`.
51pub fn unregister_code(code: &Arc<CodeObject>) {
52 let text = code.text();
53 if text.is_empty() {
54 return;
55 }
56 let end = (text.as_ptr() as usize) + text.len() - 1;
57 let code = global_code().write().remove(&end);
58 assert!(code.is_some());
59}