This is an introduction to using LLVM and the inkwell crate to write a JIT compiled calculator in Rust.

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 the sin() function from libm)
• 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).

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.

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

At the moment, we've stubbed out the rust files with a bunch of extern crate and mod statements.

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".

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.

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);
}
#}

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.

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::*;
#}

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.

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,
        }
    }
}
#}

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 matching 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!()
    }
#}

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);
}
#}