hotg_rune_wasm3_runtime/
lib.rs

1use std::collections::HashSet;
2
3use anyhow::{anyhow, Error};
4use hotg_rune_runtime::Image;
5use wasm3::{
6    CallContext, Environment, Function, Module, ParsedModule, WasmArgs,
7    WasmType, error::Trap,
8};
9
10const STACK_SIZE: u32 = 1024 * 16;
11
12// FIXME: `wasm3`'s error type isn't Send or Sync since it contains a raw
13// pointer, so it's not `anyhow`-compatible. We work around that by formatting
14// it.
15trait Wasm3ResultExt<T> {
16    fn to_anyhow(self) -> Result<T, anyhow::Error>;
17}
18
19impl<T> Wasm3ResultExt<T> for Result<T, wasm3::error::Error> {
20    fn to_anyhow(self) -> Result<T, anyhow::Error> {
21        self.map_err(|e| anyhow!("{}", e))
22    }
23}
24
25#[derive(Debug)]
26pub struct Runtime {
27    rt: wasm3::Runtime,
28}
29
30impl Runtime {
31    pub fn load<I>(wasm: &[u8], image: I) -> Result<Self, Error>
32    where
33        I: for<'a> Image<Registrar<'a>>,
34    {
35        let env = Environment::new().to_anyhow()?;
36        let module = ParsedModule::parse(&env, wasm).to_anyhow()?;
37
38        let rt = env.create_runtime(STACK_SIZE).to_anyhow()?;
39
40        log::debug!("Instantiating the WebAssembly module");
41        let instance = rt.load_module(module).to_anyhow()?;
42
43        log::debug!("Loading image");
44        let mut registrar = Registrar::new(instance);
45        image.initialize_imports(&mut registrar);
46
47        // TODO: Rename the _manifest() method to _start() so it gets
48        // automatically invoked while instantiating.
49        let manifest: Function<(), i32> = registrar
50            .into_module()
51            .find_function("_manifest")
52            .to_anyhow()?;
53        manifest.call().to_anyhow()?;
54
55        log::debug!("Loaded the Rune");
56
57        Ok(Runtime { rt })
58    }
59
60    pub fn call(&mut self) -> Result<(), Error> {
61        log::debug!("Running the rune");
62
63        let call_func: Function<(i32, i32, i32), i32> =
64            self.rt.find_function("_call").to_anyhow()?;
65
66        // For some reason we pass in the RAND capability ID when it's meant
67        // to be the Rune's responsibility to remember it. Similarly we are
68        // passing in the sine model's output type as the "input_type" parameter
69        // even though the model should know that.
70        //
71        // We should be able to change the _call function's signature once
72        // hotg-ai/rune#28 lands.
73        call_func.call(0, 0, 0).to_anyhow()?;
74
75        Ok(())
76    }
77}
78
79pub struct Registrar<'m>(RegistrarInner<'m>);
80
81enum RegistrarInner<'m> {
82    Module(Module<'m>),
83    Tracing(HashSet<(String, String)>),
84}
85
86impl<'m> Registrar<'m> {
87    pub fn new(module: Module<'m>) -> Self {
88        Self(RegistrarInner::Module(module))
89    }
90
91    pub fn tracing() -> Self { Self(RegistrarInner::Tracing(HashSet::new())) }
92
93    pub fn register_function<Args, Ret, F>(
94        &mut self,
95        namespace: &str,
96        name: &str,
97        f: F,
98    ) -> &mut Self
99    where
100        Args: WasmArgs,
101        Ret: WasmType,
102        F: for<'cc> FnMut(CallContext<'cc>, Args) -> Result<Ret, Trap>
103            + 'static,
104    {
105        match &mut self.0 {
106            RegistrarInner::Module(module) => match module
107                .link_closure(namespace, name, f)
108            {
109                Ok(()) => {},
110                Err(wasm3::error::Error::FunctionNotFound) => {
111                    // This error occurs when we try to link a function into the
112                    // program that the program doesn't import. We
113                    // just ignore that error here, since that is fine.
114                },
115                Err(e) => {
116                    panic!(
117                        "wasm3 register_function failed for `{}::{}`: {}",
118                        namespace, name, e
119                    );
120                },
121            },
122            RegistrarInner::Tracing(trace) => {
123                trace.insert((namespace.to_string(), name.to_string()));
124            },
125        }
126
127        self
128    }
129
130    pub fn into_module(self) -> Module<'m> {
131        match self.0 {
132            RegistrarInner::Module(m) => m,
133            RegistrarInner::Tracing(_) => {
134                panic!("called `into_module` on tracing registrar")
135            },
136        }
137    }
138
139    pub fn into_trace(self) -> HashSet<(String, String)> {
140        match self.0 {
141            RegistrarInner::Module(_) => {
142                panic!("called `into_trace` on non-tracing registrar")
143            },
144            RegistrarInner::Tracing(trace) => trace,
145        }
146    }
147}