A Rust Calculator, Part 2: From Tokens to a Tree

Teaching the calculator order of operations

In part 1 I built the lexer, which takes raw text like this:

3 + 4 * (2 - 1) ** 2

and turns it into a flat list of tokens:

Number(3)
Plus
Number(4)
Star
Lpar
Number(2)
Minus
Number(1)
Rpar
Power
Number(2)

That felt like good progress, but a flat list is still not something we can calculate with. The next piece is the parser, and this is the part I was honestly the most unsure about before starting. It turns out to be less magical than I expected.

Why a flat list is not enough

Take a simpler expression:

3 + 4 * 2

If we just walked through the tokens from left to right and calculated as we went, we would do 3 + 4 first and get 7, then 7 * 2 and end up with 14.

The correct answer is, as those of you with big math skills know, 11, because multiplication binds tighter than addition. Somewhere, we just have to teach the calculator that as well.

So the job of the parser is to take the flat list and work out the structure that was implied all along.

The structure is a tree

I like trees, and it just so happens that the structure we want turns out to be a tree. For 3 + 4 * 2 it looks like this:

      +
     / \
    3   *
       / \
      4   2

In beautiful ascii drawing. Now I know this might not seem to make sense right now, so lets go through it. The + at the top means: my result is the thing on my left plus the thing on my right. The thing on the right happens to be another operation. To get an answer you have to work from the bottom up: 4 * 2 first, then 3 + 8.

The deeper something sits in the tree, the earlier it gets calculated. Precedence is not a rule we check while calculating, it is baked into the shape of the tree.

This kind of tree is usually called an AST, an abstract syntax tree.

Representing the tree in Rust

Just like the lexer needed a Token enum, the parser needs an Expr enum. My first attempt was this:

enum Expr {
    Number(f64),
    Add(Expr, Expr),
}

An expression is either a number, or an addition of two other expressions. That reads exactly like the definition of the tree above, so I was feeling pretty good about it. The compiler was not:

error[E0072]: recursive type `Expr` has infinite size
 --> main.rs:1:1
  |
1 | enum Expr {
  | ^^^^^^^^^
2 |     Number(f64),
3 |     Add(Expr, Expr),
  |         ---- recursive without indirection
  |
help: insert some indirection (e.g., a `Box`, `Rc`, or `&`) to break the cycle

The problem: Rust wants to know how big an Expr is, and an Expr that directly contains two more Exprs, which each contain two more, never stops growing. The type is infinitely large on paper.

The fix is the one the compiler itself suggests (shoutout to the rust compiler actually, I am impressed!). A Box<Expr> is a pointer to an Expr that lives on the heap, and a pointer always has the same known size. So the working version is:

#[derive(Debug)]
enum Expr {
    Number(f64),
    Add(Box<Expr>, Box<Expr>),
    Sub(Box<Expr>, Box<Expr>),
    Mul(Box<Expr>, Box<Expr>),
    Div(Box<Expr>, Box<Expr>),
    Pow(Box<Expr>, Box<Expr>),
}

I have to say, an error message that explains the problem and suggests you a fix is just nice!

Precedence as layers

Now for the actual parsing. The approach I am using is called recursive descent, and the trick behind it is to describe the expression in layers, one layer per precedence level:

expression = term, followed by any number of (+ or -) term
term       = power, followed by any number of (* or /) power
power      = primary, optionally followed by ** power
primary    = a number, or ( expression )

Reading that slowly: an expression is things added and subtracted together. But the things being added are not numbers, they are terms, and a term is things multiplied and divided together. And the things being multiplied are powers, and so on down.

The loosest operators live at the top layer and the tightest ones at the bottom, and that ordering alone is what produces the correct tree. When the expression layer asks for its left side, the term layer will greedily eat 4 * 2 as one unit before handing it back. That is the whole secret. Nobody ever compares precedences anywhere in the code.

Each layer becomes one function, and each function calls the layer below it.

Walking the tokens

The parser needs to walk through the token list the same way the lexer walked through characters, including the ability to look at the next token without consuming it. So Peekable returns from part 1:

use std::iter::Peekable;
use std::slice::Iter;

type TokenStream<'a> = Peekable<Iter<'a, Token>>;

The type line is just an alias so I do not have to write the full type in every function signature. It is the same peekable idea as before, except now it iterates over tokens instead of characters.

The top layer: + and -

The plan for the expression layer, in pseudocode:

Parse one term, call it left.

While the next token is + or -:
    Consume the operator.
    Parse another term, call it right.
    Combine: left becomes (left op right).

Return left.

In Rust:

fn parse_expression(stream: &mut TokenStream) -> Result<Expr, String> {
    let mut left = parse_term(stream)?;
    while let Some(token) = stream.peek() {
        match token {
            Token::Plus => {
                stream.next();
                let right = parse_term(stream)?;
                left = Expr::Add(Box::new(left), Box::new(right));
            }
            Token::Minus => {
                stream.next();
                let right = parse_term(stream)?;
                left = Expr::Sub(Box::new(left), Box::new(right));
            }
            _ => break,
        }
    }
    Ok(left)
}

Note the left = Expr::Add(...) line. Each time through the loop, everything built so far gets wrapped up as the left side of the new operation. For 1 + 2 + 3 that gives (1 + 2) + 3, so repeated operators group from the left, which is what you would expect from maths class.

If the next token is not + or -, the loop simply stops and returns what it has. It does not treat that as an error, because the token might belong to someone else, for example a ) that the parenthesis code below is waiting for.

The term layer for * and / is the exact same shape, one level down: it parses powers instead of terms, and matches Token::Star and Token::Slash instead. I will not repeat it here, it is in the full listing at the end.

Power is different

The ** operator gets its own layer, and it has one quirk: it groups from the right. In maths notation,

2 ** 3 ** 2

means 2 ** (3 ** 2) which is 2 ** 9 = 512, not (2 ** 3) ** 2 which would be 64.

A while loop like the one above would give the wrong grouping. Instead, the recipe uses itself for the right-hand side:

Parse one primary, call it base.

If the next token is **:
    Consume it.
    Parse a whole power (this same recipe, from the top), call it exponent.
    Combine: base ** exponent.

Return base.

That "parse a whole power" step in the middle is the function calling itself. In Rust:

fn parse_power(stream: &mut TokenStream) -> Result<Expr, String> {
    let base = parse_primary(stream)?;
    if let Some(Token::Power) = stream.peek() {
        stream.next();
        let exponent = parse_power(stream)?;
        return Ok(Expr::Pow(Box::new(base), Box::new(exponent)));
    }
    Ok(base)
}

When the exponent is itself a power expression, the recursive call swallows all of it before the outer Pow is built. So the rightmost ** ends up deepest in the tree, and gets calculated first. Checking with the finished parser:

2 ** 3 ** 2   becomes   Pow(Number(2.0), Pow(Number(3.0), Number(2.0)))

Which is exactly the 2 ** (3 ** 2) shape we wanted. I had to draw this on paper before I believed it.

The bottom layer: numbers and parentheses

At the bottom of the stack sits parse_primary, which handles the two kinds of atoms: a plain number, or a parenthesised expression.

Look at the next token.

If it is a number:
    That number is the expression.

If it is (:
    Parse a whole expression.
    Expect a ) right after it.

Anything else:
    Return an error.

In Rust:

fn parse_primary(stream: &mut TokenStream) -> Result<Expr, String> {
    match stream.next() {
        Some(Token::Number(value)) => Ok(Expr::Number(*value)),
        Some(Token::Lpar) => {
            let inner = parse_expression(stream)?;
            match stream.next() {
                Some(Token::Rpar) => Ok(inner),
                Some(other) => Err(format!("Expected a closing parenthesis, found {:?}", other)),
                None => Err("Expected a closing parenthesis, found the end of the input".to_string()),
            }
        }
        Some(other) => Err(format!("Expected a number or a parenthesis, found {:?}", other)),
        None => Err("Expected a number or a parenthesis, found the end of the input".to_string()),
    }
}

The parenthesis case is my favourite part of the whole parser. When it meets a (, it calls parse_expression, the function at the very top of the stack. The entire machinery runs again on whatever is inside the brackets, and the result is handed back as a single finished unit.

That one call is the whole implementation of "parentheses override precedence". Nested parentheses work too, to any depth, without a single extra line of code. The recursion just handles it.

Catching leftovers

One thing that I checked:

3 4

parsed without complaint. parse_expression read the 3, saw that the next token was not an operator, and politely returned Number(3.0). The 4 was just left hanging.

So the public entry point checks that the parser actually consumed everything:

fn parse(tokens: &[Token]) -> Result<Expr, String> {
    let mut stream = tokens.iter().peekable();
    let expr = parse_expression(&mut stream)?;
    if let Some(leftover) = stream.next() {
        return Err(format!("Unexpected token after the expression: {:?}", leftover));
    }
    Ok(expr)
}

With that in place, the error messages come out reasonably helpful:

(3 + 4     ->  Expected a closing parenthesis, found the end of the input
3 + * 4    ->  Expected a number or a parenthesis, found Star
3 4        ->  Unexpected token after the expression: Number(4.0)

Not exactly compiler-quality diagnostics, but a lot better than silently computing the wrong thing, which can be really helpful in a calculator.

Putting it all together

Here is the parser in full. Together with the Token enum and the lex function from part 1, it makes a complete program you can run:

//The tree that the parser builds. Every operator owns the two things it works on.
#[derive(Debug)]
enum Expr {
    Number(f64),
    Add(Box<Expr>, Box<Expr>),
    Sub(Box<Expr>, Box<Expr>),
    Mul(Box<Expr>, Box<Expr>),
    Div(Box<Expr>, Box<Expr>),
    Pow(Box<Expr>, Box<Expr>),
}

//The parser peeks at tokens the same way the lexer peeked at characters
type TokenStream<'a> = Peekable<Iter<'a, Token>>;

// The Parser, takes the token vector and turns it into a tree
fn parse(tokens: &[Token]) -> Result<Expr, String> {
    let mut stream = tokens.iter().peekable();
    let expr = parse_expression(&mut stream)?;
    //If there are tokens left over, the input was not one single expression
    if let Some(leftover) = stream.next() {
        return Err(format!("Unexpected token after the expression: {:?}", leftover));
    }
    Ok(expr)
}

//Lowest precedence: + and -
fn parse_expression(stream: &mut TokenStream) -> Result<Expr, String> {
    let mut left = parse_term(stream)?;
    while let Some(token) = stream.peek() {
        match token {
            Token::Plus => {
                stream.next();
                let right = parse_term(stream)?;
                left = Expr::Add(Box::new(left), Box::new(right));
            }
            Token::Minus => {
                stream.next();
                let right = parse_term(stream)?;
                left = Expr::Sub(Box::new(left), Box::new(right));
            }
            _ => break,
        }
    }
    Ok(left)
}

//Middle precedence: * and /
fn parse_term(stream: &mut TokenStream) -> Result<Expr, String> {
    let mut left = parse_power(stream)?;
    while let Some(token) = stream.peek() {
        match token {
            Token::Star => {
                stream.next();
                let right = parse_power(stream)?;
                left = Expr::Mul(Box::new(left), Box::new(right));
            }
            Token::Slash => {
                stream.next();
                let right = parse_power(stream)?;
                left = Expr::Div(Box::new(left), Box::new(right));
            }
            _ => break,
        }
    }
    Ok(left)
}

//Highest precedence operator: **, and it groups from the right
fn parse_power(stream: &mut TokenStream) -> Result<Expr, String> {
    let base = parse_primary(stream)?;
    if let Some(Token::Power) = stream.peek() {
        stream.next();
        let exponent = parse_power(stream)?; //The recursion here is what makes ** right-associative
        return Ok(Expr::Pow(Box::new(base), Box::new(exponent)));
    }
    Ok(base)
}

//The bottom: a number, or a whole new expression in parentheses
fn parse_primary(stream: &mut TokenStream) -> Result<Expr, String> {
    match stream.next() {
        Some(Token::Number(value)) => Ok(Expr::Number(*value)),
        Some(Token::Lpar) => {
            let inner = parse_expression(stream)?;
            match stream.next() {
                Some(Token::Rpar) => Ok(inner),
                Some(other) => Err(format!("Expected a closing parenthesis, found {:?}", other)),
                None => Err("Expected a closing parenthesis, found the end of the input".to_string()),
            }
        }
        Some(other) => Err(format!("Expected a number or a parenthesis, found {:?}", other)),
        None => Err("Expected a number or a parenthesis, found the end of the input".to_string()),
    }
}

fn main() {
    match lex("3 + 4 * (2 - 1) ** 2") {
        Ok(tokens) => println!("{:#?}", parse(&tokens)),
        Err(e) => println!("Lexer error: {e}"),
    }
}

You will also need these two lines at the top of the file for the TokenStream alias:

use std::iter::Peekable;
use std::slice::Iter;

Running it on our expression from the start prints the tree, and thanks to {:#?} it is even indented so you can see the shape:

Ok(
    Add(
        Number(
            3.0,
        ),
        Mul(
            Number(
                4.0,
            ),
            Pow(
                Sub(
                    Number(
                        2.0,
                    ),
                    Number(
                        1.0,
                    ),
                ),
                Number(
                    2.0,
                ),
            ),
        ),
    ),
)

Reading it from the inside out: 2 - 1 sits deepest, then it gets raised to the power of 2, then multiplied by 4, and finally added to 3.

The next part is the evaluator, which walks this tree and finally produces an actual number. Compared to what the parser had to figure out, I suspect it will feel almost easy.