Bratish Goswami

Personal website and blog

Trying a Few Small Optimizations

The lowered IR still asks the program to add two constants at runtime. We already know both operands, so we can do that work during compilation. That's a useful place to start optimizing.

We'll fold constant expressions, propagate known values, and remove unused instructions where that preserves behavior. These few transformations give us enough to study before considering loop vectorization or machine-specific tuning.

Why this phase matters

Reducing the instruction count and the number of live values makes the backend's job easier, as well as potentially speeding up the program. Optimization can make debugging harder, though. Production compilers need to track source mappings carefully as instructions are transformed or removed.

Production optimizers also consider loops, aliasing, inlining, vector instructions, caches, branch prediction, and undefined behavior. Devirtualization, escape analysis, and target-specific cost models add further work. rustc optimizes MIR and delegates much of its backend optimization to LLVM. We'll keep to the few passes we can follow here.

What we are building here

The optimizer belongs in the ir module we added to main.rs for lowering. It takes that lowered IrProgram and changes its instructions in place. The driver can then pass the result straight to code generation, or skip optimization and pass the original IR.

We are adding an optimization pipeline over the IR from lowering. A pass is a transformation or analysis that runs over a program representation. Some passes change the IR. Some only collect facts. For explanation, we can separate the steps like this; the companion combines folding and propagation in one forward scan:

lowered IR
  -> constant folding
  -> simple constant propagation
  -> dead code elimination
  -> optimized IR

Constant folding evaluates operations whose inputs are known constants. If the IR says %2 = add %0, %1, and %0 is 5 while %1 is 3, the compiler can replace the add with %2 = const 8.

Here is a tiny source snippet:

fn main() {
    let x: i32 = 5 + 3;
    println!("{}", x);
}

Lowering gives us:

bb0:
  %0 = const 5
  %1 = const 3
  %2 = add %0, %1
  store x, %2
  %3 = load x
  println %3
  return

After constant folding:

bb0:
  %0 = const 5
  %1 = const 3
  %2 = const 8
  store x, %2
  %3 = load x
  println %3
  return

The addition has gone, leaving %0 and %1 unused. We'll remove them with dead code elimination after folding.

The constant folder scans instructions with a map from value IDs to known constants. When both operands are present, it uses the same checked arithmetic as the IR runner:

impl BinOp {
    pub fn checked(self, a: i32, b: i32) -> Option<i32> {
        match self {
            Self::Add => a.checked_add(b),
            Self::Sub => a.checked_sub(b),
            Self::Mul => a.checked_mul(b),
            Self::Div => a.checked_div(b),
        }
    }
}

For 1 / 0, there is no integer result to substitute. MicroRust checks arithmetic in every build mode: checked returns None for overflow, zero division, or signed division overflow. The folder leaves those instructions in place for the runner or native backend to report at runtime, rather than evaluating a failing operation inside the compiler.

Folding depends on the language's arithmetic rules. C's undefined behavior permits some changes that Java's defined arithmetic rules out. JavaScript has coercions and NaN; Rust's ordinary overflow behavior differs between debug and release modes. Rust may also reject constant failures through lints that MicroRust doesn't implement.

Simple constant propagation extends the idea. If a local is stored once with a constant and then loaded, the load can become a constant:

  %2 = const 8
  store x, %2
  %3 = load x

can become:

  %2 = const 8
  store x, %2
  %3 = const 8

MicroRust locals are immutable after let, so propagation needs no assignment analysis. Mutation, aliases, references, and closures would complicate it. A C pointer may modify memory indirectly; Python and Ruby operations can call user code. JavaScript engines may speculate and deoptimize when assumptions fail. Rust's ownership rules clarify some aliasing, though the optimizer still needs to establish that a change is valid.

Dead code elimination removes unused results only when the instruction has no side effects: changes beyond its returned value, such as printing or writing memory. An unknown call may have effects too.

const and a valid local load can be removed when unused. An arithmetic operation that has not been folded can still fail under our checked policy, so we conservatively keep it. println and every call are also retained.

This distinction is recorded on Instr. Here, "side effects" includes possible runtime failure:

pub fn has_side_effects(&self) -> bool {
    // Checked arithmetic can trap even if its result is never used.
    !matches!(self, Self::Const { .. } | Self::LoadLocal { .. })
}

Stores need one more check. Assigning to a local that nobody reads has no observable effect in MicroRust: these locals hold integers, and other functions cannot access them. Removing such an assignment is called dead-store elimination. We can include it in the same backward walk by tracking which locals still need a value.

The pass starts with the returned value. A retained load marks its local as needed. When we reach the store that supplies that local, we keep it and remove the local from the needed set. A store with no later retained read can go. Other instructions follow the rule above: keep a needed result or an operation with possible effects.

Here is the walk with both sets of uses:

pub fn eliminate_dead_code(block: &mut BasicBlock) {
    let mut used = HashSet::new();
    let mut needed_locals = HashSet::new();
    if let Terminator::Return(Some(value)) = block.terminator {
        used.insert(value);
    }
    let mut kept = Vec::new();
    for item in block.instrs.drain(..).rev() {
        let live = item
            .instr
            .destination()
            .is_some_and(|dst| used.contains(&dst));
        let keep = match &item.instr {
            Instr::StoreLocal { local, .. } => needed_locals.remove(local),
            _ => live || item.instr.has_side_effects(),
        };
        if keep {
            if let Instr::LoadLocal { local, .. } = &item.instr {
                needed_locals.insert(*local);
            }
            used.extend(item.instr.inputs());
            kept.push(item);
        }
    }
    kept.reverse();
    block.instrs = kept;
}

Only loads we keep mark a local as needed. That lets the walk remove chains of unused bindings too. After folding and propagation, the load of x has become a constant, so its store is no longer needed. Cleanup leaves:

bb0:
  %3 = const 8
  println %3
  return

Removing a store does not necessarily remove its initializer. If an unused binding calls a function or performs arithmetic that might fail, the walk retains those instructions. It also uses local IDs, so two bindings named x remain distinct. This analysis relies on our straight-line functions and lack of references; branches or aliases would need more work.

The canonical program benefits in the first line:

  %0 = const 5
  %1 = const 3
  %2 = add %0, %1
  store x, %2

Folding replaces the addition with a constant. Showing just that result and its store:

  %2 = const 8
  store x, %2

Propagation replaces the later reads of x and y with constants, allowing cleanup to remove their stores as well. The store of result remains because its value comes from a call and is still read. Inlining would expose the body of add at the call site, allowing more constant propagation when the arguments are known. That's another pass we could add later.

An unknown divisor leaves us with a possible runtime failure, so we retain the division unless the language requires a compile-time error. We also retain calls whose unused results may hide printing, mutation, panics, allocation, or observations of time.

An optimization must preserve observable behavior under the language's rules. For MicroRust, that includes output, return values, evaluation order, and arithmetic failure. Other languages have more to account for: C's undefined behavior affects which transformations are permitted, a JavaScript property read might invoke a getter, and Python addition might call user-defined methods. Rust has destructor timing, panic unwinding, and interactions with unsafe code. Those details determine whether deleting an apparently unused operation is legal.

Pass order affects the result. Dead code elimination can't remove operands made unused by folding until folding has run. Production compilers repeat passes or use worklists to pick up opportunities created by earlier transformations. LLVM also has to manage the analyses those passes depend on, including when a change invalidates their results, as well as target knowledge and heuristics.

Checking that the program still prints 24 won't tell us whether folding happened. We can also test the before and after IR:

before:
  %0 = const 5
  %1 = const 3
  %2 = add %0, %1
  return %2

after:
  %2 = const 8
  return %2

This checks folding and the removal of the unused constants. Separate tests check unused bindings, shadowed names, and initializers whose calls or arithmetic failures must survive. Comparing IR also lets us inspect the stores and calls that remain. Renumbering instructions later may require a less literal comparison.

Optimization levels trade compilation time for runtime speed. rustc -C opt-level=0 skips many expensive passes and generally makes instructions easier to relate to the source. Higher levels ask LLVM to do more work. MicroRust has one optimization mode, with a switch to disable it when we need to inspect unoptimized IR.

When we add the native build command in the linking chapter, it will use the same switch. That gives the driver a simple shape:

./microrustc build example.mrs           # lower, optimize, generate
./microrustc build --no-opt example.mrs  # lower, generate

The unoptimized program should behave the same way. Keeping that path available helps us catch backend errors that constant folding might otherwise hide.

We test both paths. Compiler teams also compare optimization levels and backends, using assertions and IR verifiers to catch mistakes. Our verifier checks that values are defined before use and blocks have terminators. Running it before and after optimization catches malformed IR before code generation.

Run this version

This version adds folding, propagation, and dead-code cleanup to the ir module in our existing main.rs. The driver verifies the IR before and after these passes. --no-opt gives us a way to compare them with the lowering result from the last chapter.

Keep the same source files:

microrust/
|-- main.rs
`-- example.mrs

Replace main.rs with the complete listing. It already contains the earlier phases. Keep any golden.mrs file beside it; there is no new chapter directory to create.

main.rs: compiler through optimization

// Chapter 07: add constant folding and cleanup.
use std::fmt;

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct Span {
    pub start: usize,
    pub end: usize,
}

#[derive(Debug, Clone)]
pub struct Error {
    pub phase: &'static str,
    pub span: Span,
    pub message: String,
}

impl Error {
    pub fn new(phase: &'static str, span: Span, message: impl Into<String>) -> Self {
        Self {
            phase,
            span,
            message: message.into(),
        }
    }

    pub fn render(&self, path: &str, source: &str) -> String {
        let start = self.span.start.min(source.len());
        let prefix = &source[..start];
        let line = prefix.bytes().filter(|b| *b == b'\n').count() + 1;
        let column = prefix.rsplit('\n').next().unwrap_or("").chars().count() + 1;
        let text = source.lines().nth(line - 1).unwrap_or("");
        format!(
            "{}\n --> {path}:{line}:{column}\n{line} | {text}\n  | {}^",
            self,
            " ".repeat(column - 1)
        )
    }
}

impl fmt::Display for Error {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "{} error: {}", self.phase, self.message)
    }
}
impl std::error::Error for Error {}

// Later phases will use the source spans retained in the AST.
#[allow(dead_code)]
mod syntax {
    use crate::{Error, Span};

    #[derive(Debug, Clone, Copy, PartialEq, Eq)]
    pub enum TokenKind {
        Fn,
        Let,
        Return,
        Ident,
        Int,
        String,
        I32,
        Plus,
        Minus,
        Star,
        Slash,
        Colon,
        Semicolon,
        Comma,
        LParen,
        RParen,
        LBrace,
        RBrace,
        Arrow,
        Eq,
        Bang,
        Eof,
    }

    #[derive(Debug, Clone)]
    pub struct Token {
        pub kind: TokenKind,
        pub lexeme: String,
        pub span: Span,
    }

    pub fn ident_or_keyword(text: &str) -> TokenKind {
        match text {
            "fn" => TokenKind::Fn,
            "let" => TokenKind::Let,
            "return" => TokenKind::Return,
            "i32" => TokenKind::I32,
            _ => TokenKind::Ident,
        }
    }

    pub fn lex(source: &str) -> Result<Vec<Token>, Error> {
        let bytes = source.as_bytes();
        let mut pos = 0;
        let mut tokens = Vec::new();
        while pos < bytes.len() {
            if bytes[pos].is_ascii_whitespace() {
                pos += 1;
                continue;
            }
            if source[pos..].starts_with("//!")
                || (source[pos..].starts_with("///") && !source[pos..].starts_with("////"))
            {
                return Err(Error::new(
                    "lex",
                    Span {
                        start: pos,
                        end: pos + 3,
                    },
                    "doc comments are outside MicroRust",
                ));
            }
            if source[pos..].starts_with("//") {
                while pos < bytes.len() && bytes[pos] != b'\n' {
                    pos += 1;
                }
                continue;
            }
            let start = pos;
            let c = bytes[pos];
            pos += 1;
            let kind = if c.is_ascii_alphabetic() || c == b'_' {
                while pos < bytes.len()
                    && (bytes[pos].is_ascii_alphanumeric() || bytes[pos] == b'_')
                {
                    pos += 1;
                }
                ident_or_keyword(&source[start..pos])
            } else if c.is_ascii_digit() {
                while pos < bytes.len() && bytes[pos].is_ascii_digit() {
                    pos += 1;
                }
                TokenKind::Int
            } else {
                match c {
                    b'+' => TokenKind::Plus,
                    b'-' if bytes.get(pos) == Some(&b'>') => {
                        pos += 1;
                        TokenKind::Arrow
                    }
                    b'-' => TokenKind::Minus,
                    b'*' => TokenKind::Star,
                    b'/' if bytes.get(pos) == Some(&b'*') => {
                        return Err(Error::new(
                            "lex",
                            Span {
                                start,
                                end: pos + 1,
                            },
                            "block comments are outside MicroRust; use //",
                        ))
                    }
                    b'/' => TokenKind::Slash,
                    b':' => TokenKind::Colon,
                    b';' => TokenKind::Semicolon,
                    b',' => TokenKind::Comma,
                    b'(' => TokenKind::LParen,
                    b')' => TokenKind::RParen,
                    b'{' => TokenKind::LBrace,
                    b'}' => TokenKind::RBrace,
                    b'=' => TokenKind::Eq,
                    b'!' => TokenKind::Bang,
                    b'"' => {
                        while pos < bytes.len() && bytes[pos] != b'"' {
                            if !bytes[pos].is_ascii() || matches!(bytes[pos], b'\\' | b'\n' | b'\r')
                            {
                                return Err(Error::new(
                                    "lex",
                                    Span { start, end: pos },
                                    "only the unescaped format string \"{}\" is supported",
                                ));
                            }
                            pos += 1;
                        }
                        if pos == bytes.len() {
                            return Err(Error::new(
                                "lex",
                                Span { start, end: pos },
                                "unterminated string",
                            ));
                        }
                        pos += 1;
                        TokenKind::String
                    }
                    _ => {
                        let ch = source[start..].chars().next().unwrap();
                        return Err(Error::new(
                            "lex",
                            Span {
                                start,
                                end: start + ch.len_utf8(),
                            },
                            format!("unexpected character '{ch}'"),
                        ));
                    }
                }
            };
            tokens.push(Token {
                kind,
                lexeme: source[start..pos].into(),
                span: Span { start, end: pos },
            });
        }
        tokens.push(Token {
            kind: TokenKind::Eof,
            lexeme: String::new(),
            span: Span {
                start: pos,
                end: pos,
            },
        });
        Ok(tokens)
    }

    #[derive(Debug, Clone, Copy, PartialEq, Eq)]
    pub enum Type {
        I32,
        Unit,
        Never,
    }

    #[derive(Debug, Clone, Copy, PartialEq, Eq)]
    pub enum BinOp {
        Add,
        Sub,
        Mul,
        Div,
    }

    #[derive(Debug, Clone)]
    pub struct Program {
        pub functions: Vec<Function>,
    }

    #[derive(Debug, Clone)]
    pub struct Function {
        pub name: String,
        pub span: Span,
        pub params: Vec<Param>,
        pub return_type: Type,
        pub body: Block,
    }

    #[derive(Debug, Clone)]
    pub struct Param {
        pub name: String,
        pub span: Span,
    }

    #[derive(Debug, Clone)]
    pub struct Block {
        pub statements: Vec<Stmt>,
        pub tail: Option<Box<Expr>>,
        pub span: Span,
    }

    #[derive(Debug, Clone)]
    pub enum Stmt {
        Let {
            name: String,
            span: Span,
            value: Expr,
        },
        Return {
            value: Option<Expr>,
            span: Span,
        },
        Expr(Expr),
    }

    #[derive(Debug, Clone)]
    pub struct Expr {
        pub kind: ExprKind,
        pub span: Span,
    }

    #[derive(Debug, Clone)]
    pub enum ExprKind {
        Int(i32),
        Name(String),
        Binary {
            op: BinOp,
            left: Box<Expr>,
            right: Box<Expr>,
        },
        Call {
            callee: String,
            args: Vec<Expr>,
        },
        Println {
            value: Box<Expr>,
        },
        Block(Block),
    }

    impl BinOp {
        pub fn checked(self, a: i32, b: i32) -> Option<i32> {
            match self {
                Self::Add => a.checked_add(b),
                Self::Sub => a.checked_sub(b),
                Self::Mul => a.checked_mul(b),
                Self::Div => a.checked_div(b),
            }
        }
    }

    pub struct Parser {
        tokens: Vec<Token>,
        pos: usize,
    }

    impl Parser {
        pub fn new(tokens: Vec<Token>) -> Self {
            Self { tokens, pos: 0 }
        }
        fn peek(&self) -> &Token {
            &self.tokens[self.pos]
        }
        fn advance(&mut self) -> Token {
            let t = self.peek().clone();
            if t.kind != TokenKind::Eof {
                self.pos += 1;
            }
            t
        }
        fn error(&self, message: &str) -> Error {
            Error::new("parse", self.peek().span, message)
        }
        fn eat(&mut self, kind: TokenKind) -> bool {
            if self.peek().kind == kind {
                self.advance();
                true
            } else {
                false
            }
        }
        fn expect(&mut self, kind: TokenKind) -> Result<Token, Error> {
            if self.peek().kind == kind {
                Ok(self.advance())
            } else {
                Err(self.error(&format!("expected {kind:?}, found {:?}", self.peek().kind)))
            }
        }
        fn ident(&mut self) -> Result<Token, Error> {
            let token = self.expect(TokenKind::Ident)?;
            // A restricted scanner must not turn Rust keywords into valid user names.
            const RESERVED: &[&str] = &[
                "_", "as", "async", "await", "break", "const", "continue", "crate", "dyn", "else",
                "enum", "extern", "false", "for", "if", "impl", "in", "loop", "match", "mod",
                "move", "mut", "pub", "ref", "self", "Self", "static", "struct", "super", "trait",
                "true", "type", "unsafe", "use", "where", "while", "abstract", "become", "box",
                "do", "final", "macro", "override", "priv", "typeof", "unsized", "virtual",
                "yield", "try", "gen",
            ];
            if RESERVED.contains(&token.lexeme.as_str()) {
                return Err(Error::new(
                    "parse",
                    token.span,
                    "reserved Rust word is not an identifier",
                ));
            }
            Ok(token)
        }
        pub fn parse(mut self) -> Result<Program, Error> {
            let mut functions = Vec::new();
            while self.peek().kind != TokenKind::Eof {
                functions.push(self.parse_function()?);
            }
            Ok(Program { functions })
        }
        fn parse_function(&mut self) -> Result<Function, Error> {
            self.expect(TokenKind::Fn)?;
            let name = self.ident()?;
            self.expect(TokenKind::LParen)?;
            let mut params = Vec::new();
            while self.peek().kind != TokenKind::RParen {
                let param = self.ident()?;
                self.expect(TokenKind::Colon)?;
                self.expect(TokenKind::I32)?;
                params.push(Param {
                    name: param.lexeme,
                    span: param.span,
                });
                if !self.eat(TokenKind::Comma) {
                    break;
                }
            }
            self.expect(TokenKind::RParen)?;
            let return_type = if self.eat(TokenKind::Arrow) {
                self.expect(TokenKind::I32)?;
                Type::I32
            } else {
                Type::Unit
            };
            let body = self.parse_block()?;
            Ok(Function {
                name: name.lexeme,
                span: name.span,
                params,
                return_type,
                body,
            })
        }
        fn parse_block(&mut self) -> Result<Block, Error> {
            let start = self.expect(TokenKind::LBrace)?.span.start;
            let mut statements = Vec::new();
            let mut tail = None;
            while self.peek().kind != TokenKind::RBrace {
                if self.peek().kind == TokenKind::Eof {
                    return Err(self.error("unclosed block"));
                }
                if self.eat(TokenKind::Let) {
                    let name = self.ident()?;
                    self.expect(TokenKind::Colon)?;
                    self.expect(TokenKind::I32)?;
                    self.expect(TokenKind::Eq)?;
                    let value = self.parse_expr(0)?;
                    self.expect(TokenKind::Semicolon)?;
                    statements.push(Stmt::Let {
                        name: name.lexeme,
                        span: name.span,
                        value,
                    });
                } else if self.peek().kind == TokenKind::Return {
                    let span = self.advance().span;
                    let value = if self.peek().kind == TokenKind::Semicolon {
                        None
                    } else {
                        Some(self.parse_expr(0)?)
                    };
                    self.expect(TokenKind::Semicolon)?;
                    statements.push(Stmt::Return { value, span });
                } else {
                    let value = self.parse_expr(0)?;
                    if self.eat(TokenKind::Semicolon) {
                        statements.push(Stmt::Expr(value));
                    } else {
                        tail = Some(Box::new(value));
                        break;
                    }
                }
            }
            let end = self.expect(TokenKind::RBrace)?.span.end;
            Ok(Block {
                statements,
                tail,
                span: Span { start, end },
            })
        }
        fn current_infix(&self) -> Option<(BinOp, u8, u8)> {
            match self.peek().kind {
                TokenKind::Plus => Some((BinOp::Add, 20, 21)),
                TokenKind::Minus => Some((BinOp::Sub, 20, 21)),
                TokenKind::Star => Some((BinOp::Mul, 30, 31)),
                TokenKind::Slash => Some((BinOp::Div, 30, 31)),
                _ => None,
            }
        }
        fn parse_expr(&mut self, min_bp: u8) -> Result<Expr, Error> {
            let mut left = self.parse_prefix()?;
            while let Some((op, left_bp, right_bp)) = self.current_infix() {
                if left_bp < min_bp {
                    break;
                }
                self.advance();
                let right = self.parse_expr(right_bp)?;
                let span = Span {
                    start: left.span.start,
                    end: right.span.end,
                };
                left = Expr {
                    kind: ExprKind::Binary {
                        op,
                        left: Box::new(left),
                        right: Box::new(right),
                    },
                    span,
                };
            }
            Ok(left)
        }
        fn parse_prefix(&mut self) -> Result<Expr, Error> {
            match self.peek().kind {
                TokenKind::Int => {
                    let t = self.advance();
                    let n = t.lexeme.parse::<i32>().map_err(|_| {
                        Error::new("parse", t.span, "integer literal is outside 0..=2147483647")
                    })?;
                    Ok(Expr {
                        kind: ExprKind::Int(n),
                        span: t.span,
                    })
                }
                TokenKind::Ident => {
                    let name = self.ident()?;
                    let kind = if self.eat(TokenKind::Bang) {
                        if name.lexeme != "println" {
                            return Err(Error::new(
                                "parse",
                                name.span,
                                "only println! is supported",
                            ));
                        }
                        self.expect(TokenKind::LParen)?;
                        let format = self.expect(TokenKind::String)?;
                        if format.lexeme != "\"{}\"" {
                            return Err(Error::new(
                                "parse",
                                format.span,
                                "expected format string \"{}\"",
                            ));
                        }
                        self.expect(TokenKind::Comma)?;
                        let value = self.parse_expr(0)?;
                        self.expect(TokenKind::RParen)?;
                        ExprKind::Println {
                            value: Box::new(value),
                        }
                    } else if self.eat(TokenKind::LParen) {
                        let mut args = Vec::new();
                        while self.peek().kind != TokenKind::RParen {
                            args.push(self.parse_expr(0)?);
                            if !self.eat(TokenKind::Comma) {
                                break;
                            }
                        }
                        self.expect(TokenKind::RParen)?;
                        ExprKind::Call {
                            callee: name.lexeme,
                            args,
                        }
                    } else {
                        ExprKind::Name(name.lexeme)
                    };
                    let end = self.tokens[self.pos - 1].span.end;
                    Ok(Expr {
                        kind,
                        span: Span {
                            start: name.span.start,
                            end,
                        },
                    })
                }
                TokenKind::LParen => {
                    self.advance();
                    let expr = self.parse_expr(0)?;
                    self.expect(TokenKind::RParen)?;
                    Ok(expr)
                }
                TokenKind::LBrace => {
                    let block = self.parse_block()?;
                    let span = block.span;
                    Ok(Expr {
                        kind: ExprKind::Block(block),
                        span,
                    })
                }
                _ => Err(self.error("expected expression")),
            }
        }
    }
}

#[allow(dead_code)]
mod semantics {
    use crate::syntax::{Block, Expr, ExprKind, Program, Stmt, Type};
    use crate::{Error, Span};
    use std::collections::{HashMap, HashSet};

    #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
    pub struct SymbolId(pub usize);

    #[derive(Debug, Clone)]
    pub struct Symbol {
        pub name: String,
        pub kind: SymbolKind,
        pub span: Span,
    }

    #[derive(Debug, Clone, Copy, PartialEq, Eq)]
    pub enum SymbolKind {
        Function,
        Local,
        Parameter,
    }

    pub struct ResolvedProgram {
        pub ast: Program,
        pub symbols: Vec<Symbol>,
        pub names: HashMap<usize, SymbolId>,
        pub globals: HashMap<String, SymbolId>,
    }

    struct Resolver {
        symbols: Vec<Symbol>,
        names: HashMap<usize, SymbolId>,
        globals: HashMap<String, SymbolId>,
        scopes: Vec<HashMap<String, SymbolId>>,
    }

    pub fn resolve(ast: Program) -> Result<ResolvedProgram, Error> {
        let mut r = Resolver {
            symbols: vec![],
            names: HashMap::new(),
            globals: HashMap::new(),
            scopes: vec![],
        };
        for f in &ast.functions {
            if r.globals.contains_key(&f.name) {
                return Err(Error::new(
                    "resolve",
                    f.span,
                    format!("duplicate function `{}`", f.name),
                ));
            }
            let id = r.declare(&f.name, SymbolKind::Function, f.span);
            r.globals.insert(f.name.clone(), id);
        }
        for f in &ast.functions {
            r.scopes.push(HashMap::new());
            for p in &f.params {
                if r.scopes.last().unwrap().contains_key(&p.name) {
                    return Err(Error::new("resolve", p.span, "duplicate parameter"));
                }
                let id = r.declare(&p.name, SymbolKind::Parameter, p.span);
                r.scopes.last_mut().unwrap().insert(p.name.clone(), id);
            }
            r.block(&f.body)?;
            r.scopes.pop();
        }
        Ok(ResolvedProgram {
            ast,
            symbols: r.symbols,
            names: r.names,
            globals: r.globals,
        })
    }

    impl Resolver {
        fn declare(&mut self, name: &str, kind: SymbolKind, span: Span) -> SymbolId {
            let id = SymbolId(self.symbols.len());
            self.symbols.push(Symbol {
                name: name.into(),
                kind,
                span,
            });
            self.names.insert(span.start, id);
            id
        }
        fn lookup_local(&self, name: &str) -> Option<SymbolId> {
            self.scopes
                .iter()
                .rev()
                .find_map(|scope| scope.get(name).copied())
        }
        fn block(&mut self, block: &Block) -> Result<(), Error> {
            self.scopes.push(HashMap::new());
            for stmt in &block.statements {
                match stmt {
                    Stmt::Let { name, span, value } => {
                        self.expr(value)?;
                        let id = self.declare(name, SymbolKind::Local, *span);
                        self.scopes.last_mut().unwrap().insert(name.clone(), id);
                    }
                    Stmt::Return { value, .. } => {
                        if let Some(value) = value {
                            self.expr(value)?;
                        }
                    }
                    Stmt::Expr(value) => self.expr(value)?,
                }
            }
            if let Some(tail) = &block.tail {
                self.expr(tail)?;
            }
            self.scopes.pop();
            Ok(())
        }
        fn expr(&mut self, expr: &Expr) -> Result<(), Error> {
            match &expr.kind {
                ExprKind::Int(_) => {}
                ExprKind::Name(name) => {
                    let id = self.lookup_local(name).ok_or_else(|| {
                        Error::new(
                            "resolve",
                            expr.span,
                            format!("cannot find local `{name}` in this scope"),
                        )
                    })?;
                    self.names.insert(expr.span.start, id);
                }
                ExprKind::Binary { left, right, .. } => {
                    self.expr(left)?;
                    self.expr(right)?;
                }
                ExprKind::Call { callee, args } => {
                    if self.lookup_local(callee).is_some() {
                        return Err(Error::new(
                            "resolve",
                            expr.span,
                            format!("`{callee}` is a local, not a callable function"),
                        ));
                    }
                    let id = self.globals.get(callee).copied().ok_or_else(|| {
                        Error::new(
                            "resolve",
                            expr.span,
                            format!("cannot find function `{callee}`"),
                        )
                    })?;
                    self.names.insert(expr.span.start, id);
                    for arg in args {
                        self.expr(arg)?;
                    }
                }
                ExprKind::Println { value } => self.expr(value)?,
                ExprKind::Block(block) => self.block(block)?,
            }
            Ok(())
        }
    }

    #[derive(Debug, Clone)]
    pub struct FnSig {
        pub params: Vec<Type>,
        pub ret: Type,
    }

    pub struct TypedProgram {
        pub resolved: ResolvedProgram,
        pub signatures: HashMap<SymbolId, FnSig>,
        pub types: HashMap<(usize, usize), Type>,
    }

    pub fn type_check(resolved: ResolvedProgram) -> Result<TypedProgram, Error> {
        let main = resolved
            .ast
            .functions
            .iter()
            .find(|f| f.name == "main")
            .ok_or_else(|| Error::new("type", Span { start: 0, end: 0 }, "missing fn main()"))?;
        if !main.params.is_empty() || main.return_type != Type::Unit {
            return Err(Error::new(
                "type",
                main.span,
                "main must have no parameters and return unit",
            ));
        }
        let mut signatures = HashMap::new();
        for f in &resolved.ast.functions {
            if f.params.len() > 6 {
                return Err(Error::new(
                    "type",
                    f.span,
                    "MicroRust supports at most six parameters",
                ));
            }
            signatures.insert(
                resolved.globals[&f.name],
                FnSig {
                    params: vec![Type::I32; f.params.len()],
                    ret: f.return_type,
                },
            );
        }
        let mut c = Checker {
            resolved: &resolved,
            signatures: &signatures,
            types: HashMap::new(),
            exits: HashSet::new(),
            ret: Type::Unit,
        };
        for f in &resolved.ast.functions {
            c.ret = f.return_type;
            let actual = c.block(&f.body)?;
            expect_type(f.return_type, actual, f.body.span)?;
        }
        let types = c.types;
        Ok(TypedProgram {
            resolved,
            signatures,
            types,
        })
    }

    fn expect_type(expected: Type, actual: Type, span: Span) -> Result<(), Error> {
        if actual == expected || actual == Type::Never {
            Ok(())
        } else {
            Err(Error::new(
                "type",
                span,
                format!("expected {expected:?}, found {actual:?}"),
            ))
        }
    }

    struct Checker<'a> {
        resolved: &'a ResolvedProgram,
        signatures: &'a HashMap<SymbolId, FnSig>,
        types: HashMap<(usize, usize), Type>,
        exits: HashSet<(usize, usize)>,
        ret: Type,
    }

    impl Checker<'_> {
        fn block(&mut self, block: &Block) -> Result<Type, Error> {
            let mut diverges = false;
            for stmt in &block.statements {
                match stmt {
                    Stmt::Let { value, .. } => {
                        let ty = self.expr(value)?;
                        expect_type(Type::I32, ty, value.span)?;
                        diverges |= self.exits.contains(&(value.span.start, value.span.end));
                    }
                    Stmt::Return { value, span } => {
                        let actual = match value {
                            Some(value) => self.expr(value)?,
                            None => Type::Unit,
                        };
                        expect_type(self.ret, actual, *span)?;
                        diverges = true;
                    }
                    Stmt::Expr(value) => {
                        self.expr(value)?;
                        diverges |= self.exits.contains(&(value.span.start, value.span.end));
                    }
                }
            }
            let ty = if let Some(expr) = &block.tail {
                let ty = self.expr(expr)?;
                diverges |= self.exits.contains(&(expr.span.start, expr.span.end));
                ty
            } else if diverges {
                Type::Never
            } else {
                Type::Unit
            };
            if diverges {
                self.exits.insert((block.span.start, block.span.end));
            }
            // An explicit tail still has to type-check, even when unreachable.
            Ok(ty)
        }
        fn expr(&mut self, expr: &Expr) -> Result<Type, Error> {
            let ty = match &expr.kind {
                ExprKind::Int(_) | ExprKind::Name(_) => Type::I32,
                ExprKind::Binary { left, right, .. } => {
                    let a = self.expr(left)?;
                    let b = self.expr(right)?;
                    expect_type(Type::I32, a, left.span)?;
                    expect_type(Type::I32, b, right.span)?;
                    if a == Type::Never || b == Type::Never {
                        Type::Never
                    } else {
                        Type::I32
                    }
                }
                ExprKind::Call { callee, args } => {
                    let sig = self.signatures[&self.resolved.names[&expr.span.start]].clone();
                    if args.len() != sig.params.len() {
                        return Err(Error::new(
                            "type",
                            expr.span,
                            format!(
                                "function `{callee}` expects {} arguments but got {}",
                                sig.params.len(),
                                args.len()
                            ),
                        ));
                    }
                    let mut diverges = false;
                    for (arg, expected) in args.iter().zip(&sig.params) {
                        let actual = self.expr(arg)?;
                        expect_type(*expected, actual, arg.span)?;
                        diverges |= actual == Type::Never;
                    }
                    if diverges {
                        Type::Never
                    } else {
                        sig.ret
                    }
                }
                ExprKind::Println { value } => {
                    let ty = self.expr(value)?;
                    expect_type(Type::I32, ty, value.span)?;
                    if ty == Type::Never {
                        Type::Never
                    } else {
                        Type::Unit
                    }
                }
                ExprKind::Block(block) => self.block(block)?,
            };
            let exits = match &expr.kind {
                ExprKind::Binary { left, right, .. } => {
                    self.exits.contains(&(left.span.start, left.span.end))
                        || self.exits.contains(&(right.span.start, right.span.end))
                }
                ExprKind::Call { args, .. } => args
                    .iter()
                    .any(|arg| self.exits.contains(&(arg.span.start, arg.span.end))),
                ExprKind::Println { value } => {
                    self.exits.contains(&(value.span.start, value.span.end))
                }
                ExprKind::Block(block) => self.exits.contains(&(block.span.start, block.span.end)),
                _ => false,
            };
            if exits {
                self.exits.insert((expr.span.start, expr.span.end));
            }
            self.types.insert((expr.span.start, expr.span.end), ty);
            Ok(ty)
        }
    }
}

#[allow(dead_code)]
mod ir {
    use crate::semantics::{SymbolId, TypedProgram};
    use crate::syntax::{BinOp, Block, Expr, ExprKind, Stmt, Type};
    use crate::{Error, Span};
    use std::collections::{HashMap, HashSet};

    pub type ValueId = usize;
    pub type LocalId = SymbolId;

    #[derive(Debug, Clone)]
    pub struct IrProgram {
        pub functions: Vec<IrFunction>,
    }

    #[derive(Debug, Clone)]
    pub struct IrFunction {
        pub name: String,
        pub params: Vec<LocalId>,
        pub ret: Type,
        pub block: BasicBlock,
        pub value_count: usize,
    }

    #[derive(Debug, Clone)]
    pub struct BasicBlock {
        pub instrs: Vec<SpannedInstr>,
        pub terminator: Terminator,
    }

    #[derive(Debug, Clone)]
    pub struct SpannedInstr {
        pub instr: Instr,
        pub span: Span,
    }

    #[derive(Debug, Clone)]
    pub enum Instr {
        Const {
            dst: ValueId,
            value: i32,
        },
        LoadLocal {
            dst: ValueId,
            local: LocalId,
        },
        StoreLocal {
            local: LocalId,
            src: ValueId,
        },
        Binary {
            dst: ValueId,
            op: BinOp,
            left: ValueId,
            right: ValueId,
        },
        Call {
            dst: Option<ValueId>,
            name: String,
            args: Vec<ValueId>,
        },
        Println {
            value: ValueId,
        },
    }

    #[derive(Debug, Clone, Copy)]
    pub enum Terminator {
        Return(Option<ValueId>),
    }

    impl Instr {
        pub fn destination(&self) -> Option<ValueId> {
            match self {
                Self::Const { dst, .. }
                | Self::LoadLocal { dst, .. }
                | Self::Binary { dst, .. } => Some(*dst),
                Self::Call { dst, .. } => *dst,
                _ => None,
            }
        }
        pub fn inputs(&self) -> Vec<ValueId> {
            match self {
                Self::Binary { left, right, .. } => vec![*left, *right],
                Self::StoreLocal { src, .. } => vec![*src],
                Self::Println { value } => vec![*value],
                Self::Call { args, .. } => args.clone(),
                _ => vec![],
            }
        }
        pub fn has_side_effects(&self) -> bool {
            // Checked arithmetic can trap even if its result is never used.
            !matches!(self, Self::Const { .. } | Self::LoadLocal { .. })
        }
    }

    pub fn lower(typed: &TypedProgram) -> IrProgram {
        let mut functions = Vec::new();
        for f in &typed.resolved.ast.functions {
            let mut l = Lowerer {
                typed,
                instrs: vec![],
                terminator: None,
                next: 0,
            };
            let tail = l.block(&f.body);
            let terminator = l.terminator.unwrap_or(Terminator::Return(tail));
            functions.push(IrFunction {
                name: f.name.clone(),
                params: f
                    .params
                    .iter()
                    .map(|p| typed.resolved.names[&p.span.start])
                    .collect(),
                ret: f.return_type,
                block: BasicBlock {
                    instrs: l.instrs,
                    terminator,
                },
                value_count: l.next,
            });
        }
        IrProgram { functions }
    }

    struct Lowerer<'a> {
        typed: &'a TypedProgram,
        instrs: Vec<SpannedInstr>,
        terminator: Option<Terminator>,
        next: usize,
    }

    impl Lowerer<'_> {
        fn new_value(&mut self) -> ValueId {
            let id = self.next;
            self.next += 1;
            id
        }
        fn emit(&mut self, instr: Instr, span: Span) {
            self.instrs.push(SpannedInstr { instr, span });
        }
        fn block(&mut self, block: &Block) -> Option<ValueId> {
            for stmt in &block.statements {
                if self.terminator.is_some() {
                    return None;
                }
                match stmt {
                    Stmt::Let { span, value, .. } => {
                        if let Some(src) = self.expr(value) {
                            let local = self.typed.resolved.names[&span.start];
                            self.emit(Instr::StoreLocal { local, src }, *span);
                        }
                    }
                    Stmt::Return { value, .. } => {
                        let value = value.as_ref().and_then(|v| self.expr(v));
                        if self.terminator.is_none() {
                            self.terminator = Some(Terminator::Return(value));
                        }
                    }
                    Stmt::Expr(expr) => {
                        self.expr(expr);
                    }
                }
            }
            if self.terminator.is_some() {
                None
            } else {
                block.tail.as_ref().and_then(|e| self.expr(e))
            }
        }
        fn expr(&mut self, expr: &Expr) -> Option<ValueId> {
            if self.terminator.is_some() {
                return None;
            }
            let instr = match &expr.kind {
                ExprKind::Int(value) => Instr::Const {
                    dst: self.new_value(),
                    value: *value,
                },
                ExprKind::Name(_) => Instr::LoadLocal {
                    dst: self.new_value(),
                    local: self.typed.resolved.names[&expr.span.start],
                },
                ExprKind::Binary { op, left, right } => {
                    let left = self.expr(left)?;
                    let right = self.expr(right)?;
                    Instr::Binary {
                        dst: self.new_value(),
                        op: *op,
                        left,
                        right,
                    }
                }
                ExprKind::Call { callee, args } => {
                    let mut values = Vec::new();
                    for arg in args {
                        values.push(self.expr(arg)?);
                    }
                    let sig = &self.typed.signatures[&self.typed.resolved.names[&expr.span.start]];
                    let dst = if sig.ret == Type::I32 {
                        Some(self.new_value())
                    } else {
                        None
                    };
                    Instr::Call {
                        dst,
                        name: callee.clone(),
                        args: values,
                    }
                }
                ExprKind::Println { value } => Instr::Println {
                    value: self.expr(value)?,
                },
                ExprKind::Block(block) => return self.block(block),
            };
            let dst = instr.destination();
            self.emit(instr, expr.span);
            dst
        }
    }

    pub fn fold_constants(block: &mut BasicBlock) {
        let mut constants = HashMap::<ValueId, i32>::new();
        let mut locals = HashMap::<LocalId, i32>::new();
        for item in &mut block.instrs {
            let replacement = match &item.instr {
                Instr::Const { value, .. } => Some(*value),
                Instr::LoadLocal { local, .. } => locals.get(local).copied(),
                Instr::Binary {
                    op, left, right, ..
                } => constants
                    .get(left)
                    .zip(constants.get(right))
                    .and_then(|(a, b)| op.checked(*a, *b)),
                Instr::StoreLocal { local, src } => {
                    if let Some(value) = constants.get(src) {
                        locals.insert(*local, *value);
                    } else {
                        locals.remove(local);
                    }
                    None
                }
                _ => None,
            };
            if let (Some(dst), Some(value)) = (item.instr.destination(), replacement) {
                constants.insert(dst, value);
                item.instr = Instr::Const { dst, value };
            }
        }
    }

    pub fn eliminate_dead_code(block: &mut BasicBlock) {
        let mut used = HashSet::new();
        let mut needed_locals = HashSet::new();
        if let Terminator::Return(Some(value)) = block.terminator {
            used.insert(value);
        }
        let mut kept = Vec::new();
        for item in block.instrs.drain(..).rev() {
            let live = item
                .instr
                .destination()
                .is_some_and(|dst| used.contains(&dst));
            let keep = match &item.instr {
                Instr::StoreLocal { local, .. } => needed_locals.remove(local),
                _ => live || item.instr.has_side_effects(),
            };
            if keep {
                if let Instr::LoadLocal { local, .. } = &item.instr {
                    needed_locals.insert(*local);
                }
                used.extend(item.instr.inputs());
                kept.push(item);
            }
        }
        kept.reverse();
        block.instrs = kept;
    }

    pub fn optimize(program: &mut IrProgram) {
        for function in &mut program.functions {
            fold_constants(&mut function.block);
            eliminate_dead_code(&mut function.block);
        }
    }

    pub fn verify(program: &IrProgram) -> Result<(), Error> {
        let bad = |message| Error::new("IR", Span { start: 0, end: 0 }, message);
        for f in &program.functions {
            let mut values = HashSet::new();
            let mut locals: HashSet<_> = f.params.iter().copied().collect();
            for item in &f.block.instrs {
                let instr = &item.instr;
                if instr.inputs().iter().any(|v| !values.contains(v)) {
                    return Err(bad("value used before definition"));
                }
                if let Some(dst) = instr.destination() {
                    if dst >= f.value_count || !values.insert(dst) {
                        return Err(bad("invalid or duplicate destination"));
                    }
                }
                match instr {
                    Instr::LoadLocal { local, .. } if !locals.contains(local) => {
                        return Err(bad("load from uninitialized local"))
                    }
                    Instr::StoreLocal { local, .. } => {
                        locals.insert(*local);
                    }
                    Instr::Call { name, args, dst } => {
                        let callee = program
                            .functions
                            .iter()
                            .find(|c| c.name == *name)
                            .ok_or_else(|| bad("missing callee"))?;
                        if args.len() != callee.params.len()
                            || dst.is_some() != (callee.ret == Type::I32)
                        {
                            return Err(bad("invalid IR call signature"));
                        }
                    }
                    _ => {}
                }
            }
            let Terminator::Return(value) = f.block.terminator;
            if value.is_some_and(|v| !values.contains(&v))
                || value.is_some() != (f.ret == Type::I32)
            {
                return Err(bad("invalid IR return"));
            }
        }
        Ok(())
    }

    pub fn execute(program: &IrProgram) -> Result<Vec<i32>, Error> {
        verify(program)?;
        let mut runner = Runner {
            program,
            output: vec![],
            budget: 1_000_000,
        };
        runner.call("main", &[], 0)?;
        Ok(runner.output)
    }

    struct Runner<'a> {
        program: &'a IrProgram,
        output: Vec<i32>,
        budget: usize,
    }

    impl Runner<'_> {
        fn call(&mut self, name: &str, args: &[i32], depth: usize) -> Result<Option<i32>, Error> {
            let error = |message| Error::new("runtime", Span { start: 0, end: 0 }, message);
            if depth >= 256 {
                return Err(error("IR runner call-depth limit reached"));
            }
            let f = self
                .program
                .functions
                .iter()
                .find(|f| f.name == name)
                .ok_or_else(|| error("unknown function"))?;
            let mut locals: HashMap<LocalId, i32> =
                f.params.iter().copied().zip(args.iter().copied()).collect();
            let mut values = vec![0; f.value_count];
            for item in &f.block.instrs {
                if self.budget == 0 {
                    return Err(error("IR runner instruction limit reached"));
                }
                self.budget -= 1;
                match &item.instr {
                    Instr::Const { dst, value } => values[*dst] = *value,
                    Instr::LoadLocal { dst, local } => values[*dst] = locals[local],
                    Instr::StoreLocal { local, src } => {
                        locals.insert(*local, values[*src]);
                    }
                    Instr::Binary {
                        dst,
                        op,
                        left,
                        right,
                    } => {
                        values[*dst] =
                            op.checked(values[*left], values[*right]).ok_or_else(|| {
                                Error::new(
                                    "runtime",
                                    item.span,
                                    "integer overflow or division by zero",
                                )
                            })?;
                    }
                    Instr::Call { dst, name, args } => {
                        let args: Vec<_> = args.iter().map(|v| values[*v]).collect();
                        let result = self.call(name, &args, depth + 1)?;
                        if let Some(dst) = dst {
                            values[*dst] = result.expect("verified i32 call");
                        }
                    }
                    Instr::Println { value } => self.output.push(values[*value]),
                }
            }
            let Terminator::Return(value) = f.block.terminator;
            Ok(value.map(|v| values[v]))
        }
    }

    pub fn display(program: &IrProgram) -> String {
        let mut text = String::new();
        for f in &program.functions {
            text.push_str(&format!("fn {}({:?}):\nbb0:\n", f.name, f.params));
            for item in &f.block.instrs {
                let line = match &item.instr {
                    Instr::Const { dst, value } => format!("%{dst} = const {value}"),
                    Instr::LoadLocal { dst, local } => format!("%{dst} = load #{}", local.0),
                    Instr::StoreLocal { local, src } => format!("store #{}, %{src}", local.0),
                    Instr::Binary {
                        dst,
                        op,
                        left,
                        right,
                    } => format!(
                        "%{dst} = {} %{left}, %{right}",
                        format!("{op:?}").to_lowercase()
                    ),
                    Instr::Call { dst, name, args } => format!(
                        "{}call {name}({})",
                        dst.map_or(String::new(), |d| format!("%{d} = ")),
                        args.iter()
                            .map(|a| format!("%{a}"))
                            .collect::<Vec<_>>()
                            .join(", ")
                    ),
                    Instr::Println { value } => format!("println %{value}"),
                };
                text.push_str(&format!("  {line}\n"));
            }
            let Terminator::Return(value) = f.block.terminator;
            text.push_str(&format!(
                "  return{}\n",
                value.map_or(String::new(), |v| format!(" %{v}"))
            ));
        }
        text
    }
}

fn expression(expr: &syntax::Expr) -> String {
    use syntax::ExprKind;
    match &expr.kind {
        ExprKind::Int(value) => value.to_string(),
        ExprKind::Name(name) => name.clone(),
        ExprKind::Binary { op, left, right } => {
            let symbol = match op {
                syntax::BinOp::Add => "+",
                syntax::BinOp::Sub => "-",
                syntax::BinOp::Mul => "*",
                syntax::BinOp::Div => "/",
            };
            format!("({} {symbol} {})", expression(left), expression(right))
        }
        ExprKind::Call { callee, args } => {
            let args = args.iter().map(expression).collect::<Vec<_>>().join(", ");
            format!("{callee}({args})")
        }
        ExprKind::Println { value } => format!("println({})", expression(value)),
        ExprKind::Block(block) => format!("{{ {} }}", statements(block).join(" ")),
    }
}

fn statements(block: &syntax::Block) -> Vec<String> {
    let mut lines = Vec::new();
    for statement in &block.statements {
        let line = match statement {
            syntax::Stmt::Let { name, value, .. } => {
                format!("let {name}: i32 = {};", expression(value))
            }
            syntax::Stmt::Return { value, .. } => match value {
                Some(value) => format!("return {};", expression(value)),
                None => "return;".into(),
            },
            syntax::Stmt::Expr(value) => format!("{};", expression(value)),
        };
        lines.push(line);
    }
    if let Some(tail) = &block.tail {
        lines.push(format!("tail {}", expression(tail)));
    }
    lines
}

fn print_tree(program: &syntax::Program) {
    for function in &program.functions {
        let params = function
            .params
            .iter()
            .map(|param| format!("{}: i32", param.name))
            .collect::<Vec<_>>()
            .join(", ");
        println!(
            "fn {}({params}) -> {:?}",
            function.name, function.return_type
        );
        for line in statements(&function.body) {
            println!("  {line}");
        }
    }
}

fn main() {
    if let Err(error) = run() {
        eprintln!("{error}");
        std::process::exit(1);
    }
}

fn run() -> Result<(), String> {
    const USAGE: &str = "usage: microrustc <tokens|ast|resolve|check|ir|run> FILE [--no-opt]";
    let mut args = std::env::args().skip(1);
    let mode = args.next().ok_or(USAGE)?;
    if !["tokens", "ast", "resolve", "check", "ir", "run"].contains(&mode.as_str()) {
        return Err(USAGE.into());
    }
    let mut path = None;
    let mut optimize = true;
    for arg in args {
        match arg.as_str() {
            "--no-opt" if matches!(mode.as_str(), "ir" | "run") => optimize = false,
            _ if arg.starts_with('-') => return Err(format!("unknown option {arg}")),
            _ if path.is_none() => path = Some(arg),
            _ => return Err("more than one source file".into()),
        }
    }
    let path = path.ok_or(USAGE)?;
    let source = std::fs::read_to_string(&path).map_err(|error| format!("{path}: {error}"))?;
    let render = |error: Error| error.render(&path, &source);
    let tokens = syntax::lex(&source).map_err(render)?;
    if mode == "tokens" {
        print_tokens(&tokens);
        return Ok(());
    }
    let ast = syntax::Parser::new(tokens).parse().map_err(render)?;
    if mode == "ast" {
        print_tree(&ast);
        return Ok(());
    }
    let resolved = semantics::resolve(ast).map_err(render)?;
    if mode == "resolve" {
        print_symbols(&resolved);
        return Ok(());
    }
    let typed = semantics::type_check(resolved).map_err(render)?;
    if mode == "check" {
        println!("checked {path}");
        return Ok(());
    }
    let mut program = ir::lower(&typed);
    ir::verify(&program).map_err(render)?;
    if optimize {
        ir::optimize(&mut program);
        ir::verify(&program).map_err(render)?;
    }
    if mode == "ir" {
        print!("{}", ir::display(&program));
    } else {
        for value in ir::execute(&program).map_err(render)? {
            println!("{value}");
        }
    }
    Ok(())
}

fn print_tokens(tokens: &[syntax::Token]) {
    for token in tokens {
        println!(
            "{:?} {:?} {}..{}",
            token.kind, token.lexeme, token.span.start, token.span.end
        );
    }
}

fn print_symbols(program: &semantics::ResolvedProgram) {
    for (index, symbol) in program.symbols.iter().enumerate() {
        println!("#{index} {:?} {}", symbol.kind, symbol.name);
    }
}

The additions are fold_constants, eliminate_dead_code, and optimize, plus Instr::has_side_effects. Everything before lowering stays in place. The driver uses the optimized program for both ir and run, unless --no-opt is present.

Use the same canonical program in example.mrs:

fn add(a: i32, b: i32) -> i32 {
    return a + b;
}

fn main() {
    let x: i32 = 5 + 3;
    let y: i32 = x * 2;
    let result: i32 = add(x, y);
    println!("{}", result);
}

Rebuild the compiler, inspect its IR, then run that IR. Start in the folder containing microrust, or skip cd if you're already inside it:

cd microrust
rustc main.rs -o microrustc
./microrustc ir example.mrs
./microrustc run example.mrs
./microrustc run example.mrs --no-opt

The output should be:

fn add([SymbolId(2), SymbolId(3)]):
bb0:
  %0 = load #2
  %1 = load #3
  %2 = add %0, %1
  return %2
fn main([]):
bb0:
  %6 = const 8
  %7 = const 16
  %8 = call add(%6, %7)
  store #6, %8
  %9 = load #6
  println %9
  return
24
24

The first dump is optimized IR. main now contains constants 8 and 16, while add still reads its parameters and performs an addition. We haven't implemented inlining. The two 24 lines come from running with and without optimization.

Use ./microrustc ir example.mrs --no-opt to recover the longer dump from chapter 06. Temporary IDs are not renumbered after cleanup, so main now starts at %6. The earlier constants and the stores to x and y have gone.

For a failure case, replace the program temporarily with fn main() { 1 / 0; }. Both ./microrustc run example.mrs and the same command with --no-opt must report an arithmetic failure, even though the division's result is unused. Restore the canonical program afterward. Preserving that failure is part of the optimizer's job.

Many unused bindings

I'd like to try a larger input now. Two thousand unused bindings should make the change easy to see, while leaving the program's output at 24.

Rather than type all those declarations, we can generate them with a short Rust helper. Save this as make_unused.rs beside main.rs. This helper uses ordinary Rust, including a loop and file I/O; those features are not additions to MicroRust:

fn main() -> std::io::Result<()> {
    let mut source = String::from("fn main() {\n");
    for i in 0..2000 {
        source.push_str(&format!("    let unused_{i}: i32 = {i};\n"));
    }
    source.push_str("    println!(\"{}\", 24);\n}\n");
    std::fs::write("unused.mrs", source)
}

It writes unused.mrs, containing declarations from let unused_0: i32 = 0; through let unused_1999: i32 = 1999;, followed by the print statement. None of those bindings is read. After running it, the source files in our folder will be:

microrust/
|-- main.rs
|-- example.mrs
|-- make_unused.rs
`-- unused.mrs

From inside microrust, build the helper and compare the two IR dumps. Leave the compiler and example.mrs as they are:

rustc make_unused.rs -o make_unused
./make_unused
./microrustc ir unused.mrs --no-opt > unused-before.ir
./microrustc ir unused.mrs > unused-after.ir
wc -l unused-before.ir unused-after.ir
cat unused-after.ir

The spacing from wc may differ, but the counts and optimized IR should be:

4005 unused-before.ir
   5 unused-after.ir
4010 total
fn main([]):
bb0:
  %2000 = const 24
  println %2000
  return

The unoptimized dump contains a constant and a store for each binding: 4,000 instructions we don't need. Cleanup removes them, leaving the constant 24 and its print instruction. The line counts also include the function heading, block label, and return. %2000 keeps its original ID; we haven't renumbered the surviving temporaries.

We should check the output too. These commands run both IR versions and compile the same input with rustc. The warning flag only suppresses Rust's expected warnings about the unused bindings:

./microrustc run unused.mrs --no-opt
./microrustc run unused.mrs
rustc -A unused_variables --crate-name unused unused.mrs -o unused-rust
./unused-rust

Each run should print 24:

24
24
24

The IR shows what disappeared, and the runs check the result for this input. The separate initializer tests check printing and arithmetic failures. Keep unused.mrs; in the linking chapter, we'll build both versions and compare the executable sizes.

Our pipeline can now fold constants, propagate some facts through immutable locals, and remove unused stores and pure instructions. The resulting IR still uses abstract values. In code generation, we'll decide where a value such as %2 lives and which machine instructions compute it.