lean-bench/leanbench/typeindex.py

"""What an item is *about*, in terms the database can index and search on.

An item is indexed by the types it takes in -- or, when it is not a proof, by the
type it is itself. Those types are what make one item relevant to another, so the
index is the thing retrieval later hangs off.

Types are read straight out of the Lean signature (so they reflect the code that
actually elaborated), and merged with whatever the model volunteered.
"""
from __future__ import annotations

import re
from dataclasses import dataclass

PROOF_KEYWORDS = {"theorem", "lemma", "example"}
DECL_KEYWORDS = PROOF_KEYWORDS | {
    "def", "abbrev", "instance", "structure", "class", "inductive", "opaque", "axiom",
}

ROLE_INPUT = "input"   # a type the item consumes (binders, hypotheses)
ROLE_GOAL = "goal"     # the proposition a proof concludes
ROLE_OWN = "own"       # the type of a definition, i.e. what it *is*

_DECL = re.compile(
    r"^[ \t]*(?:@\[[^\]]*\][ \t\n]*)*"
    r"(?:(?:private|protected|noncomputable|partial|unsafe|scoped|local)[ \t]+)*"
    r"(" + "|".join(sorted(DECL_KEYWORDS, key=len, reverse=True)) + r")[ \t]+"
    r"([^\s:(\[{⦃]+)",
    re.MULTILINE,
)

_IDENT = re.compile(r"[A-Za-z_][A-Za-z0-9_'!?₀-₉]*(?:\.[A-Za-z_][A-Za-z0-9_'!?₀-₉]*)*")

UNICODE_TYPES = {
    "ℕ": "Nat", "ℤ": "Int", "ℚ": "Rat", "ℝ": "Real", "ℂ": "Complex",
    "𝔽": "Field", "𝕜": "Field", "Ω": "Omega", "ℍ": "Quaternion",
}

NOISE = {
    "fun", "by", "at", "with", "then", "else", "if", "do", "let", "have", "show",
    "from", "this", "sorry", "in", "forall", "exists", "and", "or", "not", "iff",
    "true", "false", "Type", "Sort", "Prop", "_root_", "match", "deriving", "where",
}


def _strip_comments(code: str) -> str:
    code = re.sub(r"/-.*?-/", " ", code, flags=re.DOTALL)
    return re.sub(r"--[^\n]*", " ", code)


def _split_top_level(text: str) -> list[tuple[str, str]]:
    """Split a signature into binder groups: [(opener, inner text), ...]."""
    groups: list[tuple[str, str]] = []
    pairs = {"(": ")", "{": "}", "[": "]", "⦃": "⦄", "⟨": "⟩"}
    depth = 0
    opener = ""
    start = 0
    for index, char in enumerate(text):
        if char in pairs:
            if depth == 0:
                opener, start = char, index + 1
            depth += 1
        elif char in pairs.values():
            depth -= 1
            if depth == 0:
                groups.append((opener, text[start:index]))
            depth = max(depth, 0)
    return groups


def _top_level_result(signature: str) -> str:
    """Whatever follows the final colon that is not inside a binder group."""
    depth = 0
    last_colon = -1
    pairs = {"(": ")", "{": "}", "[": "]", "⦃": "⦄", "⟨": "⟩"}
    for index, char in enumerate(signature):
        if char in pairs:
            depth += 1
        elif char in pairs.values():
            depth = max(0, depth - 1)
        elif char == ":" and depth == 0:
            if signature[index : index + 2] == ":=":
                continue
            last_colon = index
    return signature[last_colon + 1 :] if last_colon != -1 else ""


SORTS = {"Prop", "Type", "Sort"}


def type_names(fragment: str, *, allow_sorts: bool = False) -> list[str]:
    """Candidate type names inside a type expression.

    ``allow_sorts`` keeps ``Prop``/``Type``: noise in a hypothesis, but the whole
    answer for ``def divides (a b : Int) : Prop``, whose own type *is* Prop.
    """
    found: list[str] = []
    for symbol, name in UNICODE_TYPES.items():
        if symbol in fragment:
            found.append(name)
    for match in _IDENT.finditer(fragment):
        token = match.group(0)
        if token in NOISE and not (allow_sorts and token in SORTS):
            continue
        head = token.split(".")[0]
        if not head[:1].isupper():
            continue
        found.append(token)
    ordered: list[str] = []
    for name in found:
        if name not in ordered:
            ordered.append(name)
    return ordered


@dataclass
class Signature:
    keyword: str
    name: str
    is_proof: bool
    inputs: list[str]
    goal: list[str]
    own: list[str]

    def as_pairs(self) -> list[tuple[str, str]]:
        pairs = [(name, ROLE_INPUT) for name in self.inputs]
        pairs += [(name, ROLE_GOAL) for name in self.goal]
        pairs += [(name, ROLE_OWN) for name in self.own]
        seen: set[tuple[str, str]] = set()
        out: list[tuple[str, str]] = []
        for pair in pairs:
            if pair not in seen:
                seen.add(pair)
                out.append(pair)
        return out


def parse_signature(code: str) -> Signature | None:
    """Parse the first declaration in ``code`` into its indexable types."""
    clean = _strip_comments(code)
    match = _DECL.search(clean)
    if not match:
        return None
    keyword, name = match.group(1), match.group(2)
    rest = clean[match.end() :]

    depth = 0
    cut = len(rest)
    pairs = {"(": ")", "{": "}", "[": "]", "⦃": "⦄", "⟨": "⟩"}
    for index, char in enumerate(rest):
        if char in pairs:
            depth += 1
        elif char in pairs.values():
            depth = max(0, depth - 1)
        elif depth == 0:
            if rest[index : index + 2] == ":=":
                cut = index
                break
            if rest[index : index + 6] == "\nwhere":
                cut = index
                break
    signature = rest[:cut]

    inputs: list[str] = []
    for opener, inner in _split_top_level(signature):
        body = inner.split(":", 1)[1] if ":" in inner and opener != "[" else inner
        inputs.extend(type_names(body))

    is_proof = keyword in PROOF_KEYWORDS
    result_text = _top_level_result(signature)
    result = type_names(result_text, allow_sorts=not is_proof)
    return Signature(
        keyword=keyword,
        name=name,
        is_proof=is_proof,
        inputs=_unique(inputs),
        goal=_unique(result) if is_proof else [],
        own=[] if is_proof else _unique(result),
    )


def _unique(names: list[str]) -> list[str]:
    out: list[str] = []
    for name in names:
        if name not in out:
            out.append(name)
    return out


def index_types(code: str, extra: list[str] | None = None) -> list[tuple[str, str]]:
    """The (type name, role) pairs an item should be indexed under."""
    signature = parse_signature(code)
    pairs = signature.as_pairs() if signature else []
    known = {name for name, _ in pairs}
    role = ROLE_INPUT if (signature is None or signature.is_proof) else ROLE_OWN
    for name in extra or []:
        cleaned = str(name).strip()
        if not cleaned:
            continue
        for candidate in type_names(cleaned) or [cleaned]:
            if candidate not in known and candidate not in NOISE:
                known.add(candidate)
                pairs.append((candidate, role))
    return pairs