Introduction
This is an introduction to using LLVM and the inkwell crate to write a JIT compiled calculator in Rust.
Roadmap
By the end of this endeavour we want to have a command-line calculator which can
- Do all the basic arithmetic operations (
5 * (7+8)
) - Have access to a bunch of pre-defined constants (
2 * PI / 3
) - Call mathematical functions from the C math library (
sin(2*PI/3)
calls thesin()
function fromlibm
) - Create our own variables (
angle = 3 * PI / 4
)
If there is time we might even try to define our own functions. It'd also be
pretty cool to compile the code as a shared library (*.so
or DLL) so it can be
linked into other programs.
To do this, our calculator will need to run several phases
- Parse the input into its AST (Abstract Syntax Tree) representation
- Use inkwell to turn this AST into a LLVM Module (a single unit of
compilation in LLVM) and define a top level
calc_main()
function - JIT compile this
Module
- Call the
calc_main
(possibly passing in arguments) and print out the result
For simplicity of implementation, the only data type our language will know
about is the double
(a 64-bit floating point number).
Parsing
The first step in creating our calculator is turning a stream of text provided by the user into something more computer-friendly. This structure is usually referred to as an Abstract Syntax Tree and is essentially just a tree where each leaf node is an "atom" (the smallest possible construct in a language, usually constants or identifiers). All non-leaf nodes then correspond to the compound constructs such as binary operators or function calls.
To make things easier we'll be using lalrpop to generate our parsing code and
construct the AST. If you've never heard of lalrpop
I highly recommend you
check out their guide.
Setting Up Lalrpop
To use lalrpop
we'll need to add it to our dependencies and set up the build
script. While we're at it, lets also make sure we've added inkwell
and
failure
as dependencies (for LLVM bindings and error handling respectively).
First lets create a new cargo project. We'll structure it as a main calc
crate
with a small binary that just parses the command line arguments and sets
everything up before calling into the central crate to run everything.
$ cargo new calc
Then update Cargo.toml
:
# Cargo.toml
[package]
name = "calc"
version = "0.1.0"
authors = ["Michael Bryan <michaelfbryan@gmail.com>"]
build = "build.rs"
[dependencies]
inkwell = { git = "https://github.com/TheDan64/inkwell", features = ["llvm3-7"] }
failure = "0.1.1"
lalrpop-util = "0.14.0"
regex = "0.2.7"
[build-dependencies]
lalrpop = "0.14.0"
And the build script:
// build.rs extern crate lalrpop; fn main() { lalrpop::process_root().unwrap(); }
With the lalrpop build system set up we can lay out the crate's skeleton. It's usually a good idea to break each phase of a compiler (because that's what we're effectively making) out into their own modules, so here's the tentative directory structure:
- /
- bin/
- yalc.rs
- src/
- lib.rs
- syntax/
- mod.rs
- ast.rs
- grammar.lalrpop
- bin/
At the moment, we've stubbed out the rust files with a bunch of extern crate
and mod
statements.
The Language Grammar
Now we've got a lot of the boilerplate set up, we can start trying to figure out what our language's grammar should look like.
The easiest way to do this is by writing out a bunch of example use cases.
# This is a comment
5 * (3+4) # You can do the usual arithmetic stuff
x = 3*PI/4 # and read/write variables
y = sin(x)^2 # plus call functions
While this language won't be turing complete (we don't have conditionals or loops), it should be a fairly decent calculator.
Once you have several examples the next step is to formalize the language grammar to make it easier to parse. This is usually done by writing a bunch of "rules" in [Backus-Naur Form][bnf].
expr := <term>
| "(" <expr> ")"
| <function-call>
term := <factor>
| <term> "+" <term>
| <term> "-" <term>
factor := NUMBER
| IDENTIFIER
| <factor> "*" <factor>
| <factor> "/" <factor>
function-call := IDENTIFIER "(" <arg-list> ")"
arg-list := EPSILON
| <expr> ("," <expr>)*
To put it in human terms, we would read the first rule as saying "an expr is either a term, an expr surrounded by parentheses, or a function call".
The AST
Now we have an idea of the language's syntax and the various elements in it, we can define an Abstract Syntax Tree for it.
At the very bottom of the tree is the Atom
. This is either a number literal or
an identifier.
# #![allow(unused_variables)] #fn main() { #[derive(Debug, Clone, PartialEq)] pub enum Atom { Number(f64), Ident(String), } #}
To make constructing an Atom
easier, you probably want to implement From<T>
for f64
, String
, and &'a str
.
Next up is the BinaryOp
. This is just a container which holds its left and
right arguments, plus the operation that was used.
# #![allow(unused_variables)] #fn main() { #[derive(Debug, Clone, PartialEq)] pub struct BinaryOp { pub op: Op, pub left: Expr, pub right: Expr, } #[derive(Debug, Clone, Copy, PartialEq)] pub enum Op { Add, Divide, Multiply, Subtract, } #}
If you were paying attention, you will have seen that the type of a BinaryOp
's
left
operand is Expr
. This will be our language's top-level construct and is
implemented as a simple enum.
# #![allow(unused_variables)] #fn main() { #[derive(Debug, Clone, PartialEq)] pub enum Expr { FunctionCall(FunctionCall), Atom(Atom), BinaryOp(Box<BinaryOp>), } #}
The last thing we need to define is a FunctionCall
. This is just a thing that
has a name and a bunch of arguments.
# #![allow(unused_variables)] #fn main() { #[derive(Debug, Clone, PartialEq)] pub struct FunctionCall { pub name: String, pub arguments: Vec<Expr>, } #}
It is recommended to sprinkle the ast
module with implementations of From
or
similar constructors/helper functions to make working with an AST and creation
easier.
Note: Assignment nodes have been left as an exercise for the reader. They're not overly difficult to add to the language, in fact, there's a way to add them without needing to define any new types.
Writing grammar.lalrpop
Now we've got some types to work with we can write a grammar which lalrpop
will use when generating the parser.
The top of the grammar.lalrpop
will be inserted into the generated file as-is,
making it the perfect place to insert the import statements we'll need.
At the moment we only need to put a single line here.
# #![allow(unused_variables)] #fn main() { use syntax::ast::{Expr, Atom, BinaryOp, FunctionCall}; #}
Next we tell lalrpop
that the grammar section has started
# #![allow(unused_variables)] #fn main() { grammar; #}
Our grammar is composed of expressions which are built up from a bunch of factors and terms. This lets us quite naturally break the grammar up into three rules.
# #![allow(unused_variables)] #fn main() { pub Expr: Expr = { <l:Expr> "+" <r:Factor> => BinaryOp::add(l, r).into(), <l:Expr> "-" <r:Factor> => BinaryOp::sub(l, r).into(), Factor, }; Factor: Expr = { <l:Factor> "*" <r:Term> => BinaryOp::mult(l, r).into(), <l:Factor> "/" <r:Term> => BinaryOp::div(l, r).into(), Term, }; Term: Expr = { "(" <e:Expr> ")" => e, Atom => Expr::Atom(<>), <f:FunctionCall> => f.into(), }; #}
There's also the rule for function calls:
# #![allow(unused_variables)] #fn main() { pub FunctionCall: FunctionCall = { <i:ident> "(" <a:CommaSeparated<Expr>> ")" => FunctionCall::new(i, a), }; CommaSeparated<T>: Vec<T> = { <v:(<T> ",")*> <e:T?> => match e { None => v, Some(e) => { let mut v = v; v.push(e); v } } }; #}
And finally, we define the rules for parsing an Atom
.
# #![allow(unused_variables)] #fn main() { pub Atom: Atom = { num => Atom::Number(<>), ident => Atom::Ident(<>), }; num: f64 = { <s:r"[0-9]+(\.[0-9]+)?"> => s.parse().unwrap(), }; ident: String = { <i:r"[a-zA-Z][a-zA-Z0-9_-]*"> => i.to_string(), }; #}
As a sanity check, we should add some tests to make sure the language's grammar parses correctly.
First up, we'll test for parsing atoms.
# #![allow(unused_variables)] #fn main() { #[test] fn parse_a_number_atom() { let src = "3.14"; let should_be = Atom::Number(3.14); let got = grammar::parse_Atom(src).unwrap(); assert_eq!(got, should_be); } #[test] fn parse_an_identifier() { let src = "x"; let should_be = Atom::Ident(String::from(src)); let got = grammar::parse_Atom(src).unwrap(); assert_eq!(got, should_be); } #}
Then a binary op,
# #![allow(unused_variables)] #fn main() { #[test] fn parse_a_multiply() { let src = "a * 5"; let should_be = BinaryOp::mult(Atom::Ident(String::from("a")).into(), Atom::Number(5.0).into()); let should_be = Expr::from(should_be); let got = grammar::parse_Expr(src).unwrap(); assert_eq!(got, should_be); } #}
And we should also add a test for function calls.
# #![allow(unused_variables)] #fn main() { #[test] fn parse_a_function_call() { let src = "sin(90.0)"; let should_be = FunctionCall::new("sin", vec![Expr::Atom(Atom::Number(90.0))]); let got = grammar::parse_FunctionCall(src).unwrap(); assert_eq!(got, should_be); } #}
So far our tests have checked individual grammar rules in isolation. To ensure operator precedence is encoded correctly in the language's grammar we'll need to create a non-trivial example and make sure it gives us exactly the parse tree we'd expect.
# #![allow(unused_variables)] #fn main() { #[test] fn complex_parse_tree() { let src = "5 + (3-2) * x - sin(90.0)"; let should_be = BinaryOp::sub( BinaryOp::add(Atom::from(5).into(), BinaryOp::mult( BinaryOp::sub(Atom::from(3).into(), Atom::from(2).into()).into(), Atom::from("x").into(), ).into()).into(), FunctionCall::new("sin", vec![Atom::Number(90.0).into()]).into() ); let should_be = Expr::from(should_be); let got = grammar::parse_Expr(src).unwrap(); assert_eq!(got, should_be); } #}
Creating an AST Visitor
Now that we can parse source code into an AST we need a mechanism for traversing
the tree. The way this is commonly done in Rust is by defining a Visitor
trait
which, by default, will recursively walk the tree.
A good example of this is syn::visit::Visit
.
The Visitor Trait
Our AST is nowhere near as complex as the Rust AST, so we shouldn't require
as many methods as syn
's Visit
trait.
# #![allow(unused_variables)] #fn main() { pub trait Visitor { fn visit_expr(&mut self, e: &Expr) { walk_expr(self, e); } fn visit_binary_op(&mut self, b: &BinaryOp) { walk_binary_op(self, b); } fn visit_function_call(&mut self, f: &FunctionCall) { walk_function_call(self, f); } fn visit_atom(&mut self, atom: &Atom) {} } #}
We also define several walk_*()
functions that will allow the Visitor
to
recursively visit each node in the tree using the default traversal order. By
making them pub
we allow users to use them to continue walking the tree when
they've done something at a particular node.
# #![allow(unused_variables)] #fn main() { pub fn walk_expr<V: Visitor + ?Sized>(visitor: &mut V, e: &Expr) { match *e { Expr::Atom(ref a) => visitor.visit_atom(a), _ => unimplemented!() } } pub fn walk_binary_op<V: Visitor + ?Sized>(visitor: &mut V, b: &BinaryOp) { visitor.visit_expr(&b.left); visitor.visit_expr(&b.right); } pub fn walk_function_call<V: Visitor + ?Sized>(visitor: &mut V, f: &FunctionCall) { for arg in &f.arguments { visitor.visit_expr(arg); } } #}
Don't forget to update mod.rs
so this visit
module is included in the crate.
# #![allow(unused_variables)] #fn main() { // src/syntax/mod.rs mod ast; mod grammar; pub mod visit; pub use self::ast::*; #}
Converting to LLVM IR
Now we've got a more computer-friendly representation of our program we need to convert it to LLVM's Intermediate Representation. This IR can come in several forms, a human-readable "assembly", a compiled bitcode, or an in-memory tree of objects (similar to our current AST).
LLVM uses the Module
as its base compilation unit (think of it as a single
*.c
file), with a Module
containing several functions, datatypes, constants,
or global variables.
Because our simple calculator doesn't allow you to declare functions, we're
going to throw everything into one big main function with the signature
fn calc_main() -> f64
. This way when we JIT compile the program we can call
into the calc_main()
function to execute everything. It also means it's quite
trivial to compile the program into a shared library (*.so
or DLL) so other
programs can call it.
This conversion process is done by recursively walking the parsed AST, turning each node into the corresponding LLVM instruction(s).
Later on, the Visitor
trait will be used to do some minor type-checking by
looking for variables uses/assignments and function calls.
The Compiler Struct
LLVM uses a Context
for the various internal states and variables involved
during the compilation process, which we'll encapsulate in a Compiler
struct.
The Compiler
will hold an IR Builder
and a cached LLVM FloatType
representing a double
.
# #![allow(unused_variables)] #fn main() { pub struct Compiler<'ctx> { ctx: &'ctx Context, logger: Logger, builder: Builder, double: FloatType, } #}
The constructor for Compiler
isn't overly exciting:
# #![allow(unused_variables)] #fn main() { impl<'ctx> Compiler<'ctx> { pub fn new(ctx: &'ctx Context) -> Compiler<'ctx> { Compiler::new_with_logger(ctx, &Logger::root(Discard, o!())) } pub fn new_with_logger(ctx: &'ctx Context, logger: &Logger) -> Compiler<'ctx> { let logger = logger.new(o!("phase" => "trans")); let double = ctx.f64_type(); let builder = ctx.create_builder(); Compiler { ctx, builder, logger, double, } } } #}
Super Basic Compilation
The compilation process is actually quite simple. We take in an AST and
recursively visit each node, generating the corresponding LLVM IR. To begin
with, we'll hard-code the module to be "calc"
and compile our one and only
function.
For this first pass we're going to take several short-cuts (noticeable by the
use of unimplemented!()
) so we can get the initial compiler working.
# #![allow(unused_variables)] #fn main() { /// Compile an AST tree to a LLVM `Module`. pub fn compile(&self, ast: &Expr) -> Module { const CALC_ENTRYPOINT: &'static str = "calc_main"; let mut module = self.ctx.create_module("calc"); self.compile_function(&mut module, CALC_ENTRYPOINT, ast); module } #}
In LLVM, you create a function by first declaring its signature, then add one or
more basic blocks (contiguous set of instructions without any branching or
jumps). The entry point of every function is typically named "entry"
.
# #![allow(unused_variables)] #fn main() { fn compile_function(&self, module: &mut Module, name: &str, body: &Expr) -> FunctionValue { // hard-code all functions to be `fn() -> f64` let sig = self.double.fn_type(&[], false); let func = module.add_function(name, &sig, None); let entry = func.append_basic_block("entry"); self.builder.position_at_end(&entry); let ret = self.compile_expr(body); self.builder.build_return(Some(&ret)); func } #}
Compiling an Expr
is just a case of match
ing on the type of expression and
calling the corresponding compile_*()
method.
# #![allow(unused_variables)] #fn main() { fn compile_expr(&self, expr: &Expr) -> FloatValue { match *expr { Expr::Atom(ref atom) => self.compile_atom(atom), Expr::BinaryOp(ref op) => self.compile_binary_op(op), Expr::FunctionCall(ref call) => self.compile_function_call(call), } } #}
We're not going to worry about variables just yet, so compiling atoms is just a case of emitting a constant.
# #![allow(unused_variables)] #fn main() { fn compile_atom(&self, atom: &Atom) -> FloatValue { match *atom { Atom::Number(n) => self.double.const_float(n), _ => unimplemented!(), } } #}
Compiling a binary op is also fairly straightforward, we need to match
on the
type of operation and then use the Builder
's build_float_*()
methods to
emit the corresponding LLVM IR.
There's a little twist to this step though. In order to compile a binary
operation we need to give LLVM its two operands. This means we'll need to
recursively call compile_expr()
on op.left
and op.right
before the match
bit.
# #![allow(unused_variables)] #fn main() { fn compile_binary_op(&self, op: &BinaryOp) -> FloatValue { let left = self.compile_expr(&op.left); let right = self.compile_expr(&op.right); match op.op { Op::Add => self.builder.build_float_add(&left, &right, "add"), Op::Subtract => self.builder.build_float_sub(&left, &right, "sub"), Op::Multiply => self.builder.build_float_mul(&left, &right, "mul"), Op::Divide => self.builder.build_float_div(&left, &right, "div"), } } #}
Compiling function calls requires us to do a type-checking pass beforehand, if you skip type-checking there's a good chance someone will use the wrong number of parameters and leave the world in an inconsistent state (usually resulting in a segfault).
Type-checking and symbol table generation will be done in a later chapter, so we
can leave it unimplemented!()
for now.
# #![allow(unused_variables)] #fn main() { fn compile_function_call(&self, call: &FunctionCall) -> FloatValue { unimplemented!() } #}
Testing The Code Generation
So far we can support binary operations and double
constants. It's not overly
much, but our calc
tool can already do everything a normal desk calculator
can. We just need to ask LLVM to JIT-compile and execute our code.
Once we've parsed the source text into an AST, the JIT-compilation process consists of:
- Initialize a
Target
- Create an LLVM
Context
- Compile the AST into a
Module
- Create a JIT execution engine based on the module
- Get a pointer to the JIT-compiled
calc_main
function - Run it
While this may sound long and compilicated, it's maybe a dozen lines and one
unsafe
block at most.
# #![allow(unused_variables)] #fn main() { pub type CalcMain = unsafe extern "C" fn() -> f64; fn execute(src: &str) -> Result<f64, Error> { Target::initialize_native(&InitializationConfig::default())?; let ast = ::syntax::parse(src)?; let ctx = Context::create(); let module = Compiler::new(&ctx).compile(&ast); let ee = module .create_jit_execution_engine(OptimizationLevel::None)?; unsafe { let calc_main = ee.get_function::<CalcMain>("calc_main")?; Ok(calc_main()) } } #}
While it's not 100% production-ready yet, we can use the above execute()
function to start testing some basic inputs.
# #![allow(unused_variables)] #fn main() { #[test] fn execute_some_binary_ops() { let inputs = vec![ ("1+1", 2.0), ("1-1", 0.0), ("2*4.5", 9.0), ("100.0/3", 100.0 / 3.0), ]; for (src, should_be) in inputs { let got = execute(src).unwrap(); assert_eq!(got, should_be); } } #[test] fn execute_a_more_complex_statement() { let src = "5 * (100 + 3) / 9 - 2.5"; let should_be = 5.0 * (100.0 + 3.0) / 9.0 - 2.5; let got = execute(src).unwrap(); assert_eq!(got, should_be); } #}