llama_wasm/
lib.rs

1//! `llama-wasm` allows for LLVm modules to be executed in a WASM runtime using Lightbeam
2
3use std::collections::BTreeMap;
4
5#[derive(Debug, thiserror::Error)]
6pub enum Error {
7    #[error("Target is not WASM")]
8    InvalidTarget,
9
10    #[error("Func not found: {0}")]
11    FuncNotFound(String),
12
13    #[error("Llama: {0}")]
14    Llama(#[from] llama::Error),
15
16    #[error("Lightbeam")]
17    Lightbeam(String),
18}
19
20/// Exectuable module type
21pub use lightbeam::ExecutableModule as Exec;
22
23/// WASM context
24pub struct Wasm<'a> {
25    exec: Exec,
26    func_map: BTreeMap<String, usize>,
27    _module: &'a llama::Module<'a>,
28    _codegen: llama::Codegen,
29}
30
31impl<'a> Wasm<'a> {
32    /// Create a new WASM context with the given LLVM module and export names
33    pub fn new<'b>(
34        module: &'a llama::Module,
35        exports: impl AsRef<[&'b str]>,
36    ) -> Result<Wasm<'a>, Error> {
37        if !module.target()?.to_ascii_lowercase().starts_with("wasm") {
38            return Err(Error::InvalidTarget);
39        }
40
41        let mut func_map = BTreeMap::new();
42
43        let exports = exports.as_ref();
44
45        let codegen = llama::Codegen::new(&module, exports, true)?;
46
47        let mut index = 0;
48        let symbols = codegen.symbols();
49        for sym in symbols {
50            if exports.contains(&sym.as_str()) {
51                func_map.insert(sym.clone(), index);
52            }
53            index += 1;
54        }
55
56        let exec = lightbeam::translate(codegen.as_ref())
57            .map_err(|x| Error::Lightbeam(format!("{:?}", x)))?;
58        Ok(Wasm {
59            _module: module,
60            _codegen: codegen,
61            exec,
62            func_map,
63        })
64    }
65
66    /// Get index of named export
67    pub fn index(&self, name: impl AsRef<str>) -> Option<u32> {
68        self.func_map.get(name.as_ref()).map(|x| *x as u32)
69    }
70
71    /// Get executable module
72    pub fn exec(&self) -> &Exec {
73        &self.exec
74    }
75}
76
77/// Call a function within a WASM context:
78/// call!(module_name.function_name(i32, i64))
79#[macro_export]
80macro_rules! call {
81    ($wasm:ident.$name:ident($($arg:expr),*$(,)?)) => {
82        $wasm.exec().execute_func($wasm.index(stringify!($name)).expect("Invalid function"), ($($arg,)*)).map_err(|x| $crate::Error::Lightbeam(format!("{:?}", x)))
83    };
84}
85
86#[cfg(test)]
87mod tests {
88    use crate::*;
89
90    #[test]
91    fn codegen() -> Result<(), Error> {
92        let context = llama::Context::new()?;
93        let mut module = llama::Module::new(&context, "test")?;
94        module.set_wasm32();
95
96        let builder = llama::Builder::new(&context)?;
97
98        let i32 = llama::Type::int(&context, 32)?;
99
100        let ft = llama::FuncType::new(i32, &[i32, i32])?;
101        module.declare_function(&builder, "testing_sub", ft, |f| {
102            let params = f.params();
103            let a = builder.sub(&params[0], &params[1], "a")?;
104            builder.ret(&a)
105        })?;
106
107        let ft = llama::FuncType::new(i32, &[i32, i32])?;
108        module.declare_function(&builder, "testing", ft, |f| {
109            let params = f.params();
110            let a = builder.add(&params[0], &params[1], "a")?;
111            builder.ret(&a)
112        })?;
113
114        println!("{}", module);
115
116        let wasm = Wasm::new(&module, &["testing"])?;
117        println!("{:?}", wasm.func_map);
118
119        let x: i32 = call!(wasm.testing(1i32, 2i32))?;
120        assert_eq!(x, 3);
121        Ok(())
122    }
123
124    #[test]
125    fn test_for_loop() -> Result<(), Error> {
126        use llama::*;
127
128        let ctx = Context::new()?;
129        let mut module = Module::new(&ctx, "test_for_loop")?;
130        module.set_wasm32();
131        let builder = Builder::new(&ctx)?;
132
133        let i64 = Type::int(&ctx, 64)?;
134        let ft = FuncType::new(i64, &[i64])?;
135        module.declare_function(&builder, "testing", ft, |f| {
136            let params = f.params();
137            let one = Const::int_sext(i64, 1)?;
138            let f = builder.for_loop(
139                Const::int_sext(i64, 0)?,
140                |x| builder.icmp(Icmp::LLVMIntSLE, x, &params[0], "cond"),
141                |x| builder.add(x, one, "add"),
142                |x| Ok(*x),
143            )?;
144            builder.ret(f)
145        })?;
146
147        println!("{}", module);
148
149        let wasm = Wasm::new(&module, &["testing"])?;
150        println!("{:?}", wasm.func_map);
151
152        let x: i64 = call!(wasm.testing(9i64))?;
153        assert_eq!(x, 9);
154        Ok(())
155    }
156}
llama_wasm/lib.rs

llama_wasm/
lib.rs
