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
148
149
150
151
152
153
154
155
156

//! `llama-wasm` allows for LLVm modules to be executed in a WASM runtime using Lightbeam

use std::collections::BTreeMap;

#[derive(Debug, thiserror::Error)]
pub enum Error {
    #[error("Target is not WASM")]
    InvalidTarget,

    #[error("Func not found: {0}")]
    FuncNotFound(String),

    #[error("Llama: {0}")]
    Llama(#[from] llama::Error),

    #[error("Lightbeam")]
    Lightbeam(String),
}

/// Exectuable module type
pub use lightbeam::ExecutableModule as Exec;

/// WASM context
pub struct Wasm<'a> {
    exec: Exec,
    func_map: BTreeMap<String, usize>,
    _module: &'a llama::Module<'a>,
    _codegen: llama::Codegen,
}

impl<'a> Wasm<'a> {
    /// Create a new WASM context with the given LLVM module and export names
    pub fn new<'b>(
        module: &'a llama::Module,
        exports: impl AsRef<[&'b str]>,
    ) -> Result<Wasm<'a>, Error> {
        if !module.target()?.to_ascii_lowercase().starts_with("wasm") {
            return Err(Error::InvalidTarget);
        }

        let mut func_map = BTreeMap::new();

        let exports = exports.as_ref();

        let codegen = llama::Codegen::new(&module, exports, true)?;

        let mut index = 0;
        let symbols = codegen.symbols();
        for sym in symbols {
            if exports.contains(&sym.as_str()) {
                func_map.insert(sym.clone(), index);
            }
            index += 1;
        }

        let exec = lightbeam::translate(codegen.as_ref())
            .map_err(|x| Error::Lightbeam(format!("{:?}", x)))?;
        Ok(Wasm {
            _module: module,
            _codegen: codegen,
            exec,
            func_map,
        })
    }

    /// Get index of named export
    pub fn index(&self, name: impl AsRef<str>) -> Option<u32> {
        self.func_map.get(name.as_ref()).map(|x| *x as u32)
    }

    /// Get executable module
    pub fn exec(&self) -> &Exec {
        &self.exec
    }
}

/// Call a function within a WASM context:
/// call!(module_name.function_name(i32, i64))
#[macro_export]
macro_rules! call {
    ($wasm:ident.$name:ident($($arg:expr),*$(,)?)) => {
        $wasm.exec().execute_func($wasm.index(stringify!($name)).expect("Invalid function"), ($($arg,)*)).map_err(|x| $crate::Error::Lightbeam(format!("{:?}", x)))
    };
}

#[cfg(test)]
mod tests {
    use crate::*;

    #[test]
    fn codegen() -> Result<(), Error> {
        let context = llama::Context::new()?;
        let mut module = llama::Module::new(&context, "test")?;
        module.set_wasm32();

        let builder = llama::Builder::new(&context)?;

        let i32 = llama::Type::int(&context, 32)?;

        let ft = llama::FuncType::new(i32, &[i32, i32])?;
        module.declare_function(&builder, "testing_sub", ft, |f| {
            let params = f.params();
            let a = builder.sub(&params[0], &params[1], "a")?;
            builder.ret(&a)
        })?;

        let ft = llama::FuncType::new(i32, &[i32, i32])?;
        module.declare_function(&builder, "testing", ft, |f| {
            let params = f.params();
            let a = builder.add(&params[0], &params[1], "a")?;
            builder.ret(&a)
        })?;

        println!("{}", module);

        let wasm = Wasm::new(&module, &["testing"])?;
        println!("{:?}", wasm.func_map);

        let x: i32 = call!(wasm.testing(1i32, 2i32))?;
        assert_eq!(x, 3);
        Ok(())
    }

    #[test]
    fn test_for_loop() -> Result<(), Error> {
        use llama::*;

        let ctx = Context::new()?;
        let mut module = Module::new(&ctx, "test_for_loop")?;
        module.set_wasm32();
        let builder = Builder::new(&ctx)?;

        let i64 = Type::int(&ctx, 64)?;
        let ft = FuncType::new(i64, &[i64])?;
        module.declare_function(&builder, "testing", ft, |f| {
            let params = f.params();
            let one = Const::int_sext(i64, 1)?;
            let f = builder.for_loop(
                Const::int_sext(i64, 0)?,
                |x| builder.icmp(Icmp::LLVMIntSLE, x, &params[0], "cond"),
                |x| builder.add(x, one, "add"),
                |x| Ok(*x),
            )?;
            builder.ret(f)
        })?;

        println!("{}", module);

        let wasm = Wasm::new(&module, &["testing"])?;
        println!("{:?}", wasm.func_map);

        let x: i64 = call!(wasm.testing(9i64))?;
        assert_eq!(x, 9);
        Ok(())
    }
}