LC fusion showcase

lc.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 | Let | Eager


@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


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


@dataclass(frozen=True)
class Eager:
    body: Term


def show(term: Term) -> str:
    match term:
        case Var(var):
            return f"{var}"
        case Lam(binder, body):
            return f"λ{binder}.{show(body)}"
        case App(fun, arg):
            return f"({show(fun)} {show(arg)})"
        case Let(binder, body, cont):
            return f"{binder} = {show(body)}; {show(cont)}"
        case Eager(body):
            return f"!!{show(body)}"


def rename(term: Term, depth: Name = 0, ctx: dict[Name, Name] | None = None) -> Term:
    if ctx is None:
        ctx = {}
    match term:
        case Var(var):
            return Var(ctx[var])
        case Lam(binder, body):
            return Lam(depth, rename(body, depth + 1, ctx | {binder: depth}))
        case App(fun, arg):
            return App(rename(fun, depth, ctx), rename(arg, depth, ctx))
        case Let(binder, body, cont):
            return Let(depth, rename(body, depth, ctx), rename(cont, depth + 1, ctx | {binder: depth}))
        case Eager(body):
            return Eager(rename(body, depth, ctx))


def nf(term: Term) -> tuple[Term, int]:
    match whnf(term):
        case Lam(binder, body), r:
            new_body, body_r = nf(body)
            return Lam(binder, new_body), r + body_r
        case App(fun, arg), r:
            new_fun, fun_r = nf(fun)
            new_arg, arg_r = nf(arg)
            return App(new_fun, new_arg), r + fun_r + arg_r
        case Let(binder, body, cont), r:
            new_body, body_r = nf(body)
            new_cont, cont_r = nf(cont)
            return Let(binder, new_body, new_cont), r + body_r + cont_r
        case Eager(body), r:
            new_body, body_r = nf(body)
            return Eager(new_body), r + body_r
        case term, r:
            return term, r


def whnf(term: Term) -> tuple[Term, int]:
    match term:
        case App(fun, arg):
            match whnf(fun):
                case Lam(binder, body), r:
                    new_body, body_r = whnf(sub(binder, arg, body))
                    return new_body, r + body_r + 1
                case fun, r:
                    return App(fun, arg), r
        case Let(binder, body, cont):
            body, body_r = whnf(body)
            new_cont, cont_r = whnf(sub(binder, body, cont))
            return new_cont, body_r + cont_r + 1
        case Eager(body):
            new_body, body_r = nf(body)
            return new_body, body_r
        case term:
            return term, 0


def fv(term: Term) -> set[Name]:
    match term:
        case Var(name):
            return {name}
        case Lam(binder, body):
            return fv(body) - {binder}
        case App(fun, arg):
            return fv(fun) | fv(arg)
        case Let(binder, body, cont):
            return fv(body) | (fv(cont) - {binder})
        case Eager(body):
            return fv(body)


def sub(old: Name, new: Term, term: Term, new_fv: set[Name] | None = None) -> Term:
    if new_fv is None:
        new_fv = fv(new)
    match term:
        case Var(name):
            return new if name == old else term
        case Lam(binder, body):
            if binder == old:
                return term
            if binder in new_fv:
                fresh_binder = fresh()
                return Lam(fresh_binder, sub(old, new, sub(binder, Var(fresh_binder), body, {fresh_binder}), new_fv))
            return Lam(binder, sub(old, new, body, new_fv))
        case App(fun, arg):
            return App(sub(old, new, fun, new_fv), sub(old, new, arg, new_fv))
        case Let(binder, body, cont):
            new_body = sub(old, new, body, new_fv)
            if binder == old:
                return Let(binder, new_body, cont)
            if binder in new_fv:
                fresh_binder = fresh()
                return Let(fresh_binder, new_body, sub(old, new, sub(binder, Var(fresh_binder), cont, {fresh_binder}), new_fv))
            return Let(binder, new_body, sub(old, new, cont, new_fv))
        case Eager(body):
            return Eager(sub(old, new, body, new_fv))


def parse_eof(inp: str) -> bool:
    return not inp


def parse_str(str: str, inp: str) -> str | None:
    if inp.startswith(str):
        return inp[len(str):]
    return None


def parse_trivia(inp: str) -> str:
    while inp:
        if inp[0].isspace():
            inp = inp[1:]
        elif inp[0] == '#':
            while inp and inp[0] != '\n':
                inp = inp[1:]
        else:
            break
    return inp


def parse_id(inp: str) -> tuple[str, str] | None:
    id = ""
    while inp and inp[0] in "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789@$%^&*-_+:'\"/":
        id += inp[0]
        inp = inp[1:]
    if id:
        return id, inp
    return None


def parse_var(inp: str, env: dict[str, Name]) -> tuple[Term, dict[str, Name], str] | None:
    match parse_id(inp):
        case None:
            return None
        case var, inp:
            if var in env:
                return Var(env[var]), env, inp
            return Var(name:=fresh()), env | {var: name}, inp


def parse_lam(inp: str, env: dict[str, Name]) -> tuple[Term, dict[str, Name], str] | None:
    match parse_str("λ", inp):
        case None:
            return None
        case inp:
            pass
    match parse_id(inp):
        case None:
            return None
        case binder, inp:
            pass
    match parse_str(".", inp):
        case None:
            return None
        case inp:
            pass
    match parse_trivia(inp):
        case inp:
            pass
    match parse_term(inp, env):
        case None:
            return None
        case body, env, inp:
            pass

    if binder in env:
        return Lam(env[binder], body), env, inp
    return Lam(name:=fresh(), body), env | {binder: name}, inp


def parse_app(inp: str, env: dict[str, Name]) -> tuple[Term, dict[str, Name], str] | None:
    match parse_str("(", inp):
        case None:
            return None
        case inp:
            pass
    match parse_trivia(inp):
        case inp:
            pass
    match parse_term(inp, env):
        case None:
            return None
        case fun, env, inp:
            pass
    match parse_trivia(inp):
        case inp:
            pass
    match parse_term(inp, env):
        case None:
            return None
        case arg, env, inp:
            pass
    match parse_trivia(inp):
        case inp:
            pass
    match parse_str(")", inp):
        case None:
            return None
        case inp:
            pass
    return App(fun, arg), env, inp


def parse_let(inp: str, env: dict[str, Name]) -> tuple[Term, dict[str, Name], str] | None:
    match parse_id(inp):
        case None:
            return None
        case binder, inp:
            pass
    match parse_trivia(inp):
        case inp:
            pass
    match parse_str("=", inp):
        case None:
            return None
        case inp:
            pass
    match parse_trivia(inp):
        case inp:
            pass
    match parse_term(inp, env):
        case None:
            return None
        case body, env, inp:
            pass
    match parse_trivia(inp):
        case inp:
            pass
    match parse_str(";", inp):
        case None:
            return None
        case inp:
            pass
    match parse_trivia(inp):
        case inp:
            pass
    match parse_term(inp, env):
        case None:
            return None
        case cont, env, inp:
            pass
    
    if binder in env:
        return Let(env[binder], body, cont), env, inp
    return Let(name:=fresh(), body, cont), env | {binder: name}, inp


def parse_eager(inp: str, env: dict[str, Name]) -> tuple[Term, dict[str, Name], str] | None:
    match parse_str("!!", inp):
        case None:
            return None
        case inp:
            pass
    match parse_trivia(inp):
        case inp:
            pass
    match parse_term(inp, env):
        case None:
            return None
        case body, env, inp:
            pass
    return Eager(body), env, inp


def parse_term(inp: str, env: dict[str, Name]) -> tuple[Term, dict[str, Name], str] | None:
    match parse_lam(inp, env):
        case None:
            pass
        case lam, env, inp:
            return lam, env, inp
    match parse_eager(inp, env):
        case None:
            pass
        case eager, env, inp:
            return eager, env, inp
    match parse_app(inp, env):
        case None:
            pass
        case app, env, inp:
            return app, env, inp
    match parse_let(inp, env):
        case None:
            pass
        case let, env, inp:
            return let, env, inp
    match parse_var(inp, env):
        case None:
            pass
        case var, env, inp:
            return var, env, inp


def parse_term_trivia(inp: str, env: dict[str, Name]) -> tuple[Term, dict[str, Name], str] | None:
    match parse_trivia(inp):
        case inp:
            pass
    match parse_term(inp, env):
        case None:
            return None
        case term, env, inp:
            pass
    match parse_trivia(inp):
        case inp:
            pass
    return term, env, inp


def parse(inp: str) -> Term:
    match parse_term_trivia(inp, {}):
        case None:
            raise ValueError("Parse error")
        case term, env, inp:
            pass
    match parse_eof(inp):
        case True:
            return term
        case False:
            raise ValueError("Parse error")


text = """
c1 = λf. f;
c2 = λf. f1=!!(c1 f); !!λx.(f1 (f1 x));
c4 = λf. f2=!!(c2 f); !!λx.(f2 (f2 x));
c8 = λf. f4=!!(c4 f); !!λx.(f4 (f4 x));
c16 = λf. f8=!!(c8 f); !!λx.(f8 (f8 x));
c32 = λf. f16=!!(c16 f); !!λx.(f16 (f16 x));
c64 = λf. f32=!!(c32 f); !!λx.(f32 (f32 x));
c128 = λf. f64=!!(c64 f); !!λx.(f64 (f64 x));
c256 = λf. f128=!!(c128 f); !!λx.(f128 (f128 x));
c512 = λf. f256=!!(c256 f); !!λx.(f256 (f256 x));
c1024 = λf. f512=!!(c512 f); !!λx.(f512 (f512 x));
c2048 = λf. f1024=!!(c1024 f); !!λx.(f1024 (f1024 x));
c4096 = λf. f2048=!!(c2048 f); !!λx.(f2048 (f2048 x));
c8192 = λf. f4096=!!(c4096 f); !!λx.(f4096 (f4096 x));
c16384 = λf. f8192=!!(c8192 f); !!λx.(f8192 (f8192 x));
c32768 = λf. f16384=!!(c16384 f); !!λx.(f16384 (f16384 x));
c65536 = λf. f32768=!!(c32768 f); !!λx.(f32768 (f32768 x));
c131072 = λf. f65536=!!(c65536 f); !!λx.(f65536 (f65536 x));
c262144 = λf. f131072=!!(c131072 f); !!λx.(f131072 (f131072 x));
c524288 = λf. f262144=!!(c262144 f); !!λx.(f262144 (f262144 x));
c1048576 = λf. f524288=!!(c524288 f); !!λx.(f524288 (f524288 x));

clr = λb.λt.λf.((b f) f);
inc = λn.λz.λs.((n (s z)) s);

# !! (c64 inc)
!! (c1048576 clr)
"""

term = parse(text)
print("WHNF:")
term, r = whnf(term)
print(" ", show(rename(term)))
print(f"({r} rewrites)")