lean-bench/tests/test_align.py

import json

from leanbench.align import align, llm_pairing
from leanbench.llm import ScriptedClient
from leanbench.project import project
from leanbench.srcsimplex import SourceRole, SourceSimplex


def src(*roles, exists=False):
    out = []
    counts: dict[str, int] = {}
    for kind, symbols, desc, keywords, *rest in roles:
        counts[kind] = counts.get(kind, 0) + 1
        out.append(SourceRole(
            kind, symbols, desc, keywords, rest[0] if rest else "given", f"{kind}{counts[kind]}"))
    return SourceSimplex(out, exists)


GROUP = project(
    "theorem t (G : Type*) [Group G] [Fintype G] (g : G) : g ^ Fintype.card G = 1 := by sorry")
GROUP_SRC = src(
    ("object", ["G"], "finite group", ["finite", "group"]),
    ("object", ["g"], "element of G", []),
    ("fact", ["G"], "G is finite", ["finite"]),
    ("goal", ["g", "G"], "g raised to the size of G is the identity", ["order"]),
)


def test_faithful_statement_is_not_flagged():
    result = align(GROUP_SRC, GROUP)
    assert not result.flagged, result.as_dict()
    assert result.score > 0.8
    assert {m.source for m in result.matches} >= {"object1", "object2", "fact1", "goal1"}


def test_missing_hypothesis_is_named():
    lean = project("theorem t (n : ℕ) (a : ℤ) : a ^ n = a ^ n := by rfl")
    source = src(
        ("object", ["n"], "positive integer", ["integer"]),
        ("object", ["a"], "integer", ["integer"]),
        ("fact", ["a", "n"], "a and n are coprime", ["coprime"]),
        ("goal", ["a", "n"], "a to the n", []),
    )
    result = align(source, lean)
    assert "missing_hypothesis" in result.errors
    assert result.flagged


def test_extra_hypothesis_is_named():
    lean = project("theorem t (n : ℤ) (h : 0 < n) : n = n := by rfl")
    source = src(
        ("object", ["n"], "integer", ["integer"]),
        ("goal", ["n"], "n equals itself", []),
    )
    result = align(source, lean)
    assert "extra_hypothesis" in result.errors and result.flagged


def test_changed_number_system_is_a_type_conflict():
    lean = project("theorem t (n : ℕ) : n = n := by rfl")
    source = src(
        ("object", ["n"], "an integer", ["integer"]),
        ("goal", ["n"], "n equals itself", []),
    )
    result = align(source, lean)
    assert [c.kind for c in result.conflicts if c.kind == "type_conflict"] == ["type_conflict"]


def test_exists_turned_into_given_is_a_quantifier_error():
    lean = project("theorem t (m : ℕ) : m = m := by rfl")
    source = src(
        ("object", ["m"], "a natural number", ["natural"], "exists"),
        ("goal", ["m"], "there is such an m", [], "given"),
    )
    result = align(source, lean)
    assert "quantifier" in result.errors


def test_existential_in_goal_is_a_clean_match():
    lean = project("theorem t (n : ℕ) : ∃ m : ℕ, m > n := by sorry")
    source = src(
        ("object", ["n"], "natural number", ["natural"]),
        ("object", ["m"], "a larger natural number", ["natural"], "exists"),
        ("goal", ["m", "n"], "there is a larger number", []),
        exists=True,
    )
    result = align(source, lean)
    assert "quantifier" not in result.errors
    assert "missing_object" not in result.errors


def test_goal_exists_mismatch():
    lean = project("theorem t (n : ℕ) (a : ℕ) : n = a := by sorry")
    source = src(
        ("object", ["n"], "natural number", ["natural"]),
        ("object", ["a"], "natural number", ["natural"]),
        ("goal", ["n", "a"], "there exist", []),
        exists=True,
    )
    assert "quantifier" in align(source, lean).errors


def test_wrong_goal_is_flagged():
    lean = project("theorem t (G : Type*) [Group G] (g : G) : g * g = g := by sorry")
    source = src(
        ("object", ["G"], "group", ["group"]),
        ("object", ["g"], "element", []),
        ("goal", ["x"], "the subgroup is normal", ["normal", "subgroup"]),
    )
    assert "goal_mismatch" in align(source, lean).errors


def test_degenerate_lean_is_reported():
    lean = project("theorem t : 2 + 2 = 4 := by norm_num")
    source = src(("object", ["n"], "integer", ["integer"]), ("goal", ["n"], "n", []))
    assert "degenerate" in align(source, lean).errors


def test_extra_object_is_reported_but_not_flagged_alone():
    lean = project("theorem t (n : ℕ) (k : ℕ) : n = n := by rfl")
    source = src(("object", ["n"], "natural", ["natural"]), ("goal", ["n"], "n = n", []))
    result = align(source, lean)
    assert "extra_object" in result.errors and not result.flagged


def test_pairing_overrides_the_deterministic_matcher():
    lean = project("theorem t (x : ℤ) (hx : 0 < x) : x = x := by rfl")
    source = src(
        ("object", ["z"], "integer", ["integer"]),
        ("fact", [], "it is positive", []),
        ("goal", [], "trivial", []),
    )
    pairing = {"object1": "x", "fact1": "hx"}
    result = align(source, lean, pairing=pairing)
    assert {(m.source, m.lean) for m in result.matches} >= {("object1", "x"), ("fact1", "hx")}
    assert "missing_hypothesis" not in result.errors


def test_llm_pairing_keeps_only_real_keys():
    client = ScriptedClient([json.dumps({"pairs": {"object1": "x", "fact1": "nonsense"}})])
    lean = project("theorem t (x : ℤ) (hx : 0 < x) : x = x := by rfl")
    source = src(("object", ["z"], "integer", []), ("fact", [], "positive", []),
                 ("goal", [], "", []))
    assert llm_pairing(client, source, lean) == {"object1": "x"}


def test_alignment_dict_is_json_serialisable():
    json.dumps(align(GROUP_SRC, GROUP).as_dict())