Add initial JIT compilation

Cargo.toml

@@ -6,6 +6,9 @@ authors = ["Marceline Cramer <mars@tebibyte.media>"]
 license = "GPL-3.0-or-later"
 
 [dependencies]
+cranelift = "0.78.0"
+cranelift-module = "0.78.0"
+cranelift-jit = "0.78.0"
 logos = "0.12.0"
 pest = "2"
 pest_derive = "2"

src/jit/engine.rs

@@ -0,0 +1,95 @@
+// Copyright (c) 2022 Marceline Cramer
+// SPDX-License-Identifier: GPL-3.0-or-later
+
+use crate::parse::ast;
+use cranelift::codegen::{self, Context};
+use cranelift::frontend::FunctionBuilderContext;
+use cranelift::prelude::*;
+use cranelift_jit::{JITBuilder, JITModule};
+use cranelift_module::{DataContext, Linkage, Module};
+
+pub struct Engine {
+    builder_context: FunctionBuilderContext,
+    ctx: Context,
+    data_ctx: DataContext,
+    module: JITModule,
+}
+
+impl Engine {
+    pub fn new() -> Self {
+        let builder = JITBuilder::new(cranelift_module::default_libcall_names());
+        let module = JITModule::new(builder);
+        Self {
+            builder_context: FunctionBuilderContext::new(),
+            ctx: module.make_context(),
+            data_ctx: DataContext::new(),
+            module,
+        }
+    }
+
+    pub fn compile(&mut self, expr: &ast::Expr) -> Result<*const u8, String> {
+        self.translate(expr)?;
+
+        let name = "dummy_function";
+        let id = self
+            .module
+            .declare_function(name, Linkage::Export, &self.ctx.func.signature)
+            .map_err(|e| e.to_string())?;
+
+        let mut trap_sink = codegen::binemit::NullTrapSink {};
+        let mut stack_map_sink = codegen::binemit::NullStackMapSink {};
+        self.module
+            .define_function(id, &mut self.ctx, &mut trap_sink, &mut stack_map_sink)
+            .map_err(|e| e.to_string())?;
+
+        self.module.clear_context(&mut self.ctx);
+        self.module.finalize_definitions();
+
+        let code = self.module.get_finalized_function(id);
+        Ok(code)
+    }
+
+    pub fn translate(&mut self, expr: &ast::Expr) -> Result<(), String> {
+        let int = self.module.target_config().pointer_type();
+
+        self.ctx.func.signature.returns.push(AbiParam::new(int));
+
+        let mut builder = FunctionBuilder::new(&mut self.ctx.func, &mut self.builder_context);
+
+        let entry_block = builder.create_block();
+        builder.append_block_params_for_function_params(entry_block);
+        builder.switch_to_block(entry_block);
+        builder.seal_block(entry_block);
+
+        let mut trans = super::translate::FunctionTranslator { int, builder };
+
+        let return_value = trans.translate_expr(expr);
+        trans.builder.ins().return_(&[return_value]);
+
+        println!("{}", trans.builder.func.display());
+        trans.builder.finalize();
+
+        Ok(())
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn simple_math() {
+        use crate::parse::{
+            lexer::{Lexer, Token},
+            rd::RecursiveDescent,
+        };
+
+        let source = "(1 + 2) * (3 / (4 + 5));";
+        let lexer = Lexer::new(source);
+        let mut rd = RecursiveDescent::new(lexer);
+        let (expr, tail) = rd.build_expr();
+        let expr = expr.expect("Failed to parse expression");
+        assert_eq!(tail, Token::Semicolon);
+        Engine::new().compile(&expr).unwrap();
+    }
+}

src/jit/mod.rs

@@ -0,0 +1,5 @@
+// Copyright (c) 2022 Marceline Cramer
+// SPDX-License-Identifier: GPL-3.0-or-later
+
+pub mod engine;
+pub mod translate;

src/jit/translate.rs

@@ -0,0 +1,43 @@
+// Copyright (c) 2022 Marceline Cramer
+// SPDX-License-Identifier: GPL-3.0-or-later
+
+use cranelift::prelude::*;
+use crate::parse::ast;
+
+pub struct FunctionTranslator<'a> {
+    pub int: types::Type,
+    pub builder: FunctionBuilder<'a>,
+}
+
+impl<'a> FunctionTranslator<'a> {
+    pub fn translate_expr(&mut self, expr: &ast::Expr) -> Value {
+        use ast::Expr::*;
+        match expr {
+            Literal(ast::Literal::DecimalInteger(literal)) => {
+                // TODO parse integers while building ast so that codegen doesn't have to
+                let val: i64 = literal.parse().unwrap();
+                self.builder.ins().iconst(self.int, val)
+            }
+            BinaryOp(op, terms) => {
+                let lhs = self.translate_expr(&terms.0);
+                let rhs = self.translate_expr(&terms.1);
+                self.translate_binary_op(op, lhs, rhs)
+            }
+            // TODO the AST doesn't need this either
+            Group(expr) => self.translate_expr(expr),
+            _ => unimplemented!("Expression: {:#?}", expr),
+        }
+    }
+
+    pub fn translate_binary_op(&mut self, op: &ast::BinaryOp, lhs: Value, rhs: Value) -> Value {
+        use ast::BinaryOp::*;
+        let ins = self.builder.ins();
+        match op {
+            Add => ins.iadd(lhs, rhs),
+            Sub => ins.isub(lhs, rhs),
+            Mul => ins.imul(lhs, rhs),
+            Div => ins.udiv(lhs, rhs),
+            _ => unimplemented!("Binary operation: {:#?}", op),
+        }
+    }
+}

src/lib.rs

@@ -4,6 +4,7 @@
 #[macro_use]
 extern crate pest_derive;
 
+pub mod jit;
 pub mod namevec;
 pub mod parse;
 pub mod symbols;