Attempt number [NaN]

λ.py

from dataclasses import dataclass
from itertools import count


type Name = int


name = count()


def fresh() -> Name:
    return next(name)


type Term = Var | Lam | App


@dataclass(frozen=True)
class Var:
    name: Name


@dataclass(frozen=True)
class Lam:
    binder: Name
    body: Term


@dataclass(frozen=True)
class App:
    fun: Term
    arg: Term


def show(term: Term) -> str:
    match term:
        case Var(name):
            return f"{name}"
        case Lam(binder, body):
            return f"λ{binder}.{show(body)}"
        case App(fun, arg):
            return f"({show(fun)} {show(arg)})"


def nf(term: Term) -> Term:
    match whnf(term):
        case Lam(binder, body):
            return Lam(binder, nf(body))
        case App(fun, arg):
            return App(nf(fun), nf(arg))
        case term:
            return term


def whnf(term: Term) -> Term:
    match term:
        case App(fun, arg):
            match whnf(fun):
                case Lam(binder, body):
                    return whnf(sub(binder, arg, body))
                case fun:
                    return App(fun, arg)
        case term:
            return term


def sub(old: Name, new: Term, term: Term) -> Term:
    match term:
        case Var(name) if old == name:
            return new
        case Lam(binder, body):
            return Lam(fresh_binder:=fresh(), sub(old, new, sub(binder, Var(fresh_binder), body)))
        case App(fun, arg):
            return App(sub(old, new, fun), sub(old, new, arg))
        case term:
            return term


def var() -> Var:
    return Var(fresh())


def lam(binder: Var, body: Term) -> Lam:
    return Lam(binder.name, body)


def app(fun: Term, arg: Term) -> App:
    return App(fun, arg)


main = app(lam(x:=var(), app(x, x)), lam(y:=var(), y))
print(show(main))
print(show(nf(main)))