use std::convert::{TryFrom, TryInto}; use std::path::Path; use std::result; use inkwell::basic_block::BasicBlock; use inkwell::builder::Builder; pub use inkwell::context::Context; use inkwell::module::Module; use inkwell::support::LLVMString; use inkwell::types::FunctionType; use inkwell::values::{AnyValueEnum, BasicValueEnum, FunctionValue}; use inkwell::IntPredicate; use thiserror::Error; use crate::ast::{BinaryOperator, Binding, Decl, Expr, Fun, Ident, Literal, UnaryOperator}; use crate::common::env::Env; #[derive(Debug, PartialEq, Eq, Error)] pub enum Error { #[error("Undefined variable {0}")] UndefinedVariable(Ident<'static>), #[error("LLVM Error: {0}")] LLVMError(String), } impl From for Error { fn from(s: LLVMString) -> Self { Self::LLVMError(s.to_string()) } } pub type Result = result::Result; pub struct Codegen<'ctx, 'ast> { context: &'ctx Context, pub module: Module<'ctx>, builder: Builder<'ctx>, env: Env<'ast, AnyValueEnum<'ctx>>, function_stack: Vec>, identifier_counter: u32, } impl<'ctx, 'ast> Codegen<'ctx, 'ast> { pub fn new(context: &'ctx Context, module_name: &str) -> Self { let module = context.create_module(module_name); let builder = context.create_builder(); Self { context, module, builder, env: Default::default(), function_stack: Default::default(), identifier_counter: 0, } } pub fn new_function<'a>( &'a mut self, name: &str, ty: FunctionType<'ctx>, ) -> &'a FunctionValue<'ctx> { self.function_stack .push(self.module.add_function(name, ty, None)); let basic_block = self.append_basic_block("entry"); self.builder.position_at_end(basic_block); self.function_stack.last().unwrap() } pub fn finish_function(&mut self, res: &BasicValueEnum<'ctx>) -> FunctionValue<'ctx> { self.builder.build_return(Some(res)); self.function_stack.pop().unwrap() } pub fn append_basic_block(&self, name: &str) -> BasicBlock<'ctx> { self.context .append_basic_block(*self.function_stack.last().unwrap(), name) } pub fn codegen_expr(&mut self, expr: &'ast Expr<'ast>) -> Result> { match expr { Expr::Ident(id) => self .env .resolve(id) .cloned() .ok_or_else(|| Error::UndefinedVariable(id.to_owned())), Expr::Literal(Literal::Int(i)) => { let ty = self.context.i64_type(); Ok(AnyValueEnum::IntValue(ty.const_int(*i, false))) } Expr::UnaryOp { op, rhs } => { let rhs = self.codegen_expr(rhs)?; match op { UnaryOperator::Not => unimplemented!(), UnaryOperator::Neg => Ok(AnyValueEnum::IntValue( self.builder.build_int_neg(rhs.into_int_value(), "neg"), )), } } Expr::BinaryOp { lhs, op, rhs } => { let lhs = self.codegen_expr(lhs)?; let rhs = self.codegen_expr(rhs)?; match op { BinaryOperator::Add => Ok(AnyValueEnum::IntValue(self.builder.build_int_add( lhs.into_int_value(), rhs.into_int_value(), "add", ))), BinaryOperator::Sub => Ok(AnyValueEnum::IntValue(self.builder.build_int_sub( lhs.into_int_value(), rhs.into_int_value(), "add", ))), BinaryOperator::Mul => Ok(AnyValueEnum::IntValue(self.builder.build_int_sub( lhs.into_int_value(), rhs.into_int_value(), "add", ))), BinaryOperator::Div => { Ok(AnyValueEnum::IntValue(self.builder.build_int_signed_div( lhs.into_int_value(), rhs.into_int_value(), "add", ))) } BinaryOperator::Pow => unimplemented!(), BinaryOperator::Equ => { Ok(AnyValueEnum::IntValue(self.builder.build_int_compare( IntPredicate::EQ, lhs.into_int_value(), rhs.into_int_value(), "eq", ))) } BinaryOperator::Neq => todo!(), } } Expr::Let { bindings, body } => { self.env.push(); for Binding { ident, body, .. } in bindings { let val = self.codegen_expr(body)?; self.env.set(ident, val); } let res = self.codegen_expr(body); self.env.pop(); res } Expr::If { condition, then, else_, } => { let then_block = self.append_basic_block("then"); let else_block = self.append_basic_block("else"); let join_block = self.append_basic_block("join"); let condition = self.codegen_expr(condition)?; self.builder.build_conditional_branch( condition.into_int_value(), then_block, else_block, ); self.builder.position_at_end(then_block); let then_res = self.codegen_expr(then)?; self.builder.build_unconditional_branch(join_block); self.builder.position_at_end(else_block); let else_res = self.codegen_expr(else_)?; self.builder.build_unconditional_branch(join_block); self.builder.position_at_end(join_block); let phi = self.builder.build_phi(self.context.i64_type(), "join"); phi.add_incoming(&[ (&BasicValueEnum::try_from(then_res).unwrap(), then_block), (&BasicValueEnum::try_from(else_res).unwrap(), else_block), ]); Ok(phi.as_basic_value().into()) } Expr::Call { fun, args } => { if let Expr::Ident(id) = &**fun { let function = self .module .get_function(id.into()) .or_else(|| self.env.resolve(id)?.clone().try_into().ok()) .ok_or_else(|| Error::UndefinedVariable(id.to_owned()))?; let args = args .iter() .map(|arg| Ok(self.codegen_expr(arg)?.try_into().unwrap())) .collect::>>()?; Ok(self .builder .build_call(function, &args, "call") .try_as_basic_value() .left() .unwrap() .into()) } else { todo!() } } Expr::Fun(fun) => { let Fun { args, body } = &**fun; let fname = self.fresh_ident("f"); let cur_block = self.builder.get_insert_block().unwrap(); let env = self.env.save(); // TODO: closures let function = self.codegen_function(&fname, args, body)?; self.builder.position_at_end(cur_block); self.env.restore(env); Ok(function.into()) } Expr::Ascription { expr, .. } => self.codegen_expr(expr), } } pub fn codegen_function( &mut self, name: &str, args: &'ast [Ident<'ast>], body: &'ast Expr<'ast>, ) -> Result> { let i64_type = self.context.i64_type(); self.new_function( name, i64_type.fn_type( args.iter() .map(|_| i64_type.into()) .collect::>() .as_slice(), false, ), ); self.env.push(); for (i, arg) in args.iter().enumerate() { self.env.set( arg, self.cur_function().get_nth_param(i as u32).unwrap().into(), ); } let res = self.codegen_expr(body)?.try_into().unwrap(); self.env.pop(); Ok(self.finish_function(&res)) } pub fn codegen_decl(&mut self, decl: &'ast Decl<'ast>) -> Result<()> { match decl { Decl::Fun { name, body: Fun { args, body }, } => { self.codegen_function(name.into(), args, body)?; Ok(()) } } } pub fn codegen_main(&mut self, expr: &'ast Expr<'ast>) -> Result<()> { self.new_function("main", self.context.i64_type().fn_type(&[], false)); let res = self.codegen_expr(expr)?.try_into().unwrap(); self.finish_function(&res); Ok(()) } pub fn print_to_file

(&self, path: P) -> Result<()> where P: AsRef, { Ok(self.module.print_to_file(path)?) } pub fn binary_to_file

(&self, path: P) -> Result<()> where P: AsRef, { if self.module.write_bitcode_to_path(path.as_ref()) { Ok(()) } else { Err(Error::LLVMError( "Error writing bitcode to output path".to_owned(), )) } } fn fresh_ident(&mut self, prefix: &str) -> String { self.identifier_counter += 1; format!("{}{}", prefix, self.identifier_counter) } fn cur_function(&self) -> &FunctionValue<'ctx> { self.function_stack.last().unwrap() } } #[cfg(test)] mod tests { use inkwell::execution_engine::JitFunction; use inkwell::OptimizationLevel; use super::*; fn jit_eval(expr: &str) -> anyhow::Result { let expr = crate::parser::expr(expr).unwrap().1; let context = Context::create(); let mut codegen = Codegen::new(&context, "test"); let execution_engine = codegen .module .create_jit_execution_engine(OptimizationLevel::None) .unwrap(); codegen.codegen_function("test", &[], &expr)?; unsafe { let fun: JitFunction T> = execution_engine.get_function("test")?; Ok(fun.call()) } } #[test] fn add_literals() { assert_eq!(jit_eval::("1 + 2").unwrap(), 3); } #[test] fn variable_shadowing() { assert_eq!( jit_eval::("let x = 1 in (let x = 2 in x) + x").unwrap(), 3 ); } #[test] fn eq() { assert_eq!( jit_eval::("let x = 1 in if x == 1 then 2 else 4").unwrap(), 2 ); } #[test] fn function_call() { let res = jit_eval::("let id = fn x = x in id 1").unwrap(); assert_eq!(res, 1); } }