STLC-like interpreted language in Lean 4

STLC.lean

import Batteries.Data.RBMap

inductive BinOp where
| add | sub | mul | div
deriving Repr, Ord

-- Abstrct Syntax Tree
inductive Expr where
| num (value : Int)
| var (name : String)
| if_ (cond : Bool) (conseq : Expr) (alt : Expr)
| let_ (name : String) (value : Expr) (body : Expr)
| binop (op : BinOp) (lhs : Expr) (rhs : Expr)
| read
deriving Repr, Ord

abbrev Context := Batteries.RBMap String Expr compare

namespace Expr

partial def interpret (expr : Expr) (ctx : Context) : IO Int :=
  match expr with
  | num value => pure value
  | var name =>
    match ctx.find? name with
    | none => pure 0 -- TODO: fix: print error
    | some value => value.interpret ctx
  | let_ name value body => body.interpret (ctx.insert name value)
  | read => do
    let stdin ← IO.getStdin
    let n ← stdin.getLine
    pure (n.toInt!)
  | if_ cond conseq alt =>
    if cond
    then
      conseq.interpret ctx
    else
      alt.interpret ctx
  | binop op lhs rhs =>
    match op with
    | .add => do
      let lr ← lhs.interpret ctx
      let rr ← rhs.interpret ctx
      pure (lr + rr)
    | .sub => do
      let lr ← lhs.interpret ctx
      let rr ← rhs.interpret ctx
      pure (lr - rr)
    | .mul => do
      let lr ← lhs.interpret ctx
      let rr ← rhs.interpret ctx
      pure (lr * rr)
    | .div => do
      let lr ← lhs.interpret ctx
      let rr ← rhs.interpret ctx
      pure (lr / rr)

open BinOp in
def program : Expr := (binop add (num 41) (num 1))

#eval program.interpret .empty

open BinOp in
def programWithVars : Expr := (let_ "x" (num 1) (binop div (var "x") (num 2)))
/-
let x = 1
x + 2
-/

open BinOp in
def programWithBranches : Expr :=
  let_ "x" (num 2)
    (if_ false (binop add (var "x") (num 1)) (var "x"))

#eval programWithBranches.interpret .empty

end Expr