dinet.py

from collections.abc import Callable
from dataclasses import dataclass
from itertools import count
from typing import Any


type Name = int


name = count(0)


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


type Term = Var | Dup | Sup | Data | Func


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


@dataclass(frozen=True)
class Dup:
    name: Name
    body: Term


@dataclass(frozen=True)
class Sup:
    fst: Term
    snd: Term


@dataclass(frozen=True)
class Data:
    data: Any


@dataclass(frozen=True)
class Func:
    func: Callable[..., Any]
    args: tuple[Term, ...]


def whnf(term: Term, ctx: dict[Name, Term] | None = None) -> Term:
    if ctx is None:
        ctx = {}

    match term:
        case Var(name) if name in ctx:
            return whnf(ctx[name], ctx)
        case Dup(name, body):
            match whnf(body, ctx):
                case Sup(fst, snd):
                    ctx[name] = snd
                    return whnf(fst, ctx)
                case Data(data):
                    # copy reference for now
                    ctx[name] = Data(data)
                    return Data(data)
                case body:
                    return Dup(name, body)
        case Func(func, (head, *tail)):
            match whnf(head, ctx):
                case Sup(fst, snd):
                    args1, args2 = [], []
                    for arg in tail:
                        args1.append(Dup(x := fresh(), arg))
                        args2.append(Var(x))
                    return whnf(Sup(Func(func, (fst, *args1)), Func(func, (snd, *args2))), ctx)
                case Data(data):
                    return whnf(Func(lambda *args: func(data, *args), tail), ctx) if tail else Data(func(data))
                case head:
                    return Func(func, (head, *tail))
        case term:
            return term


def compose(term: Term) -> Any:
    match whnf(term):
        case Data(data):
            return data
        case _:
            raise ValueError("composition failed")


data_a = "A"
data_b = "B"
func_c = lambda x: f"C({x})"
func_d = lambda x, y: f"D({x}, {y})"

term_a = Data(data_a)
term_b = Data(data_b)
term_ab = Sup(term_a, term_b)
term_c = Func(func_c, (term_ab,))
term_c1, term_c2 = Dup(lc := fresh(), term_c), Var(lc)
term_d = Func(func_d, (term_c1, term_c2))

print(compose(term_d))