Type checker for an extremely simple imperative programming language (Playground: https://play.rust-lang.org/?version=stable&mode=debug&edition=2021&gist=5d667459cabb1fc900473f85c541ee0a)

type_check_statements.rs

#![allow(dead_code)]

#[derive(Debug, Copy, Clone, PartialEq, Eq)]
enum Type {                             // t ::=
    Unit,                               //   | ()
    Int,                                //   | Int
    Float,                              //   | Float
}

type Expr = Box<ExprNode>;

enum ExprNode {                         // e ::=
    Var(String),                        //   | x
    Unit,                               //   | ()
    Int(i32),                           //   | ... | -1 | 0 | 1 | ...
    Float(f32),                         //   | -Inf | ... | 0.0 | ... | Inf | NaN
    Add(Expr, Expr),                    //   | e + e
}

enum Statement {                        // s ::=
    Let(String, Option<Type>, Expr),    //   | let x: T = e
    Assign(String, Expr),               //   | x = e
    Return(Expr)                        //   | return e
}

type Block = Vec<Statement>;            // b ::= s0; ...; sn

struct Context {
    locals: Vec<(String, Type)>,
}

impl Context {
    fn new() -> Context {
        Context {
            locals: Vec::new(),
        }
    }

    fn lookup(&self, x: &str) -> Type {
        self.locals.iter()
            .find_map(|(y, ty)| (x == y).then(|| *ty))
            .expect("unbound variable")
    }

    fn check_expr(&mut self, expr: &ExprNode, ty: Type) {
        match expr {
            // Ambiguous expressions that need type annotations go here
            //     ...currently there are none :)

            // Fallback to inference
            expr => assert_eq!(self.infer_expr(expr), ty),
        }
    }

    fn infer_expr(&mut self, expr: &ExprNode) -> Type {
        match expr {
            ExprNode::Var(x) => self.lookup(x),
            ExprNode::Unit => Type::Unit,
            ExprNode::Int(_) => Type::Int,
            ExprNode::Float(_) => Type::Float,
            ExprNode::Add(x, y) => {
                let x_ty = self.infer_expr(x);
                let y_ty = self.infer_expr(y);
                assert_eq!(x_ty, y_ty);
                x_ty
            }
        }
    }

    fn check_block(&mut self, block: &[Statement], return_ty: Type) {
        let mut statements = block.iter().peekable();
        while let Some(statement) = statements.next() {
            match statement {
                Statement::Return(expr) => {
                    self.check_expr(expr, return_ty);
                }
                _ if statements.peek().is_none() && return_ty != Type::Unit => {
                    panic!("missing return statement");
                }
                Statement::Let(x, Some(ty), expr) => {
                    self.check_expr(expr, *ty);
                    self.locals.push((x.clone(), *ty));
                }
                Statement::Let(x, None, expr) => {
                    let ty = self.infer_expr(expr);
                    self.locals.push((x.clone(), ty));
                }
                Statement::Assign(x, expr) => {
                    let ty = self.lookup(x);
                    self.check_expr(expr, ty);
                }
            }
        }
    }
}

fn main() {
    Context::new().check_block(
        &[
            Statement::Let("x".to_owned(), None, Box::new(ExprNode::Float(3.0))),
            Statement::Let("y".to_owned(), None, Box::new(ExprNode::Unit)),
            Statement::Let("z".to_owned(), Some(Type::Int), Box::new(ExprNode::Int(33))),
            Statement::Assign(
                "x".to_owned(), 
                Box::new(ExprNode::Add(
                    Box::new(ExprNode::Var("x".to_owned())),
                    Box::new(ExprNode::Var("x".to_owned())),
                )),
            ),
            Statement::Return(Box::new(ExprNode::Var("x".to_owned()))),
        ],
        Type::Float,
    );
}