lean-bench/leanbench/schema.py

"""Roles and simplices: the parameters of a declaration as a small typed schema.

A Lean declaration is a bundle of roles: the objects it is about (``(G : Type*)``),
the facts it assumes (``(h : a ∣ b)``), and the goal it asserts or the thing it
defines. A ``Simplex`` is such a bundle with at least two roles. The same schema can
be read off the prose of a book, so the two can be compared role by role.

Design notes (from the archived ``schema-simplex`` prototype):

* a role may be *constrained* to a closed set of allowed values (``choices``);
* a simplex is itself a role, so simplices nest (a lemma used by another is a role
  of it), which is how a dependency structure appears;
* validation lives in ``__init__``. The prototype put it in a dataclass
  ``__post_init__`` next to a hand-written ``__init__``, where it never ran.
"""
from __future__ import annotations

from collections.abc import Sequence
from typing import Any

KIND_OBJECT = "object"  # something the statement is about: a group G, an integer n
KIND_FACT = "fact"      # something assumed: n > 0, G is finite
KIND_GOAL = "goal"      # what is asserted, or asked to be shown
KIND_OWN = "own"        # what a definition is
KIND_SIMPLEX = "simplex"
KINDS = (KIND_OBJECT, KIND_FACT, KIND_GOAL, KIND_OWN, KIND_SIMPLEX)

MIN_ROLES = 2


class Role:
    """A named slot with a kind, a description and optionally a closed value set."""

    def __init__(
        self,
        name: str,
        kind: str = KIND_OBJECT,
        *,
        type_text: str = "",
        context: str = "",
        choices: Sequence[Any] | None = None,
        attrs: dict[str, Any] | None = None,
    ) -> None:
        if not name:
            raise ValueError("a role needs a name")
        if kind not in KINDS:
            raise ValueError(f"unknown role kind {kind!r}; expected one of {KINDS}")
        self.name = name
        self.kind = kind
        self.type_text = type_text
        self.context = context
        self.choices: list[Any] | None = list(choices) if choices is not None else None
        self.attrs: dict[str, Any] = dict(attrs or {})

    def is_constrained(self) -> bool:
        return self.choices is not None

    def validate_value(self, value: Any) -> bool:
        return self.choices is None or value in self.choices

    def check_value(self, value: Any) -> None:
        if not self.validate_value(value):
            raise ValueError(
                f"{value!r} is not allowed for role {self.name!r}; allowed: {self.choices}"
            )

    def as_dict(self) -> dict[str, Any]:
        out: dict[str, Any] = {"name": self.name, "kind": self.kind}
        if self.type_text:
            out["type"] = self.type_text
        if self.context:
            out["context"] = self.context
        if self.choices is not None:
            out["choices"] = self.choices
        if self.attrs:
            out["attrs"] = self.attrs
        return out

    def __repr__(self) -> str:
        return f"<Role {self.kind}:{self.name}>"


class Simplex(Role):
    """A reusable bundle of at least two roles, keyed by alias. Itself a role."""

    def __init__(self, name: str, roles: dict[str, Role], *, context: str = "") -> None:
        super().__init__(name, KIND_SIMPLEX, context=context)
        if len(roles) < MIN_ROLES:
            raise ValueError(
                f"simplex {name!r} needs at least {MIN_ROLES} roles, got {len(roles)}"
            )
        self.roles = dict(roles)

    def by_kind(self, kind: str) -> list[Role]:
        return [role for role in self.roles.values() if role.kind == kind]

    def walk(self) -> list[Role]:
        """Every non-simplex role, descending into nested simplices."""
        out: list[Role] = []
        for role in self.roles.values():
            if isinstance(role, Simplex):
                out.extend(role.walk())
            else:
                out.append(role)
        return out

    def as_dict(self) -> dict[str, Any]:
        return {
            "name": self.name,
            "roles": {alias: role.as_dict() for alias, role in self.roles.items()},
        }


class State:
    """A role together with a value; unresolved until it has one."""

    def __init__(self, role: Role, value: Any = None) -> None:
        self.role = role
        self.value = value

    def is_resolved(self) -> bool:
        return self.value is not None


class Fact(State):
    """Known at the start."""

    def __init__(self, role: Role, value: Any) -> None:
        if value is None:
            raise ValueError("a Fact must carry a value")
        super().__init__(role, value)


class Goal(State):
    """To be resolved."""