1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147

use std::collections::HashSet;

use anyhow::{anyhow, Error};
use hotg_rune_runtime::Image;
use wasm3::{
    CallContext, Environment, Function, Module, ParsedModule, WasmArgs,
    WasmType, error::Trap,
};

const STACK_SIZE: u32 = 1024 * 16;

// FIXME: `wasm3`'s error type isn't Send or Sync since it contains a raw
// pointer, so it's not `anyhow`-compatible. We work around that by formatting
// it.
trait Wasm3ResultExt<T> {
    fn to_anyhow(self) -> Result<T, anyhow::Error>;
}

impl<T> Wasm3ResultExt<T> for Result<T, wasm3::error::Error> {
    fn to_anyhow(self) -> Result<T, anyhow::Error> {
        self.map_err(|e| anyhow!("{}", e))
    }
}

#[derive(Debug)]
pub struct Runtime {
    rt: wasm3::Runtime,
}

impl Runtime {
    pub fn load<I>(wasm: &[u8], image: I) -> Result<Self, Error>
    where
        I: for<'a> Image<Registrar<'a>>,
    {
        let env = Environment::new().to_anyhow()?;
        let module = ParsedModule::parse(&env, wasm).to_anyhow()?;

        let rt = env.create_runtime(STACK_SIZE).to_anyhow()?;

        log::debug!("Instantiating the WebAssembly module");
        let instance = rt.load_module(module).to_anyhow()?;

        log::debug!("Loading image");
        let mut registrar = Registrar::new(instance);
        image.initialize_imports(&mut registrar);

        // TODO: Rename the _manifest() method to _start() so it gets
        // automatically invoked while instantiating.
        let manifest: Function<(), i32> = registrar
            .into_module()
            .find_function("_manifest")
            .to_anyhow()?;
        manifest.call().to_anyhow()?;

        log::debug!("Loaded the Rune");

        Ok(Runtime { rt })
    }

    pub fn call(&mut self) -> Result<(), Error> {
        log::debug!("Running the rune");

        let call_func: Function<(i32, i32, i32), i32> =
            self.rt.find_function("_call").to_anyhow()?;

        // For some reason we pass in the RAND capability ID when it's meant
        // to be the Rune's responsibility to remember it. Similarly we are
        // passing in the sine model's output type as the "input_type" parameter
        // even though the model should know that.
        //
        // We should be able to change the _call function's signature once
        // hotg-ai/rune#28 lands.
        call_func.call(0, 0, 0).to_anyhow()?;

        Ok(())
    }
}

pub struct Registrar<'m>(RegistrarInner<'m>);

enum RegistrarInner<'m> {
    Module(Module<'m>),
    Tracing(HashSet<(String, String)>),
}

impl<'m> Registrar<'m> {
    pub fn new(module: Module<'m>) -> Self {
        Self(RegistrarInner::Module(module))
    }

    pub fn tracing() -> Self { Self(RegistrarInner::Tracing(HashSet::new())) }

    pub fn register_function<Args, Ret, F>(
        &mut self,
        namespace: &str,
        name: &str,
        f: F,
    ) -> &mut Self
    where
        Args: WasmArgs,
        Ret: WasmType,
        F: for<'cc> FnMut(CallContext<'cc>, Args) -> Result<Ret, Trap>
            + 'static,
    {
        match &mut self.0 {
            RegistrarInner::Module(module) => match module
                .link_closure(namespace, name, f)
            {
                Ok(()) => {},
                Err(wasm3::error::Error::FunctionNotFound) => {
                    // This error occurs when we try to link a function into the
                    // program that the program doesn't import. We
                    // just ignore that error here, since that is fine.
                },
                Err(e) => {
                    panic!(
                        "wasm3 register_function failed for `{}::{}`: {}",
                        namespace, name, e
                    );
                },
            },
            RegistrarInner::Tracing(trace) => {
                trace.insert((namespace.to_string(), name.to_string()));
            },
        }

        self
    }

    pub fn into_module(self) -> Module<'m> {
        match self.0 {
            RegistrarInner::Module(m) => m,
            RegistrarInner::Tracing(_) => {
                panic!("called `into_module` on tracing registrar")
            },
        }
    }

    pub fn into_trace(self) -> HashSet<(String, String)> {
        match self.0 {
            RegistrarInner::Module(_) => {
                panic!("called `into_trace` on non-tracing registrar")
            },
            RegistrarInner::Tracing(trace) => trace,
        }
    }
}