ce_backtrack/core.py

"""Fast magma-equation core: parse, codegen evaluator, counterexample checks.

Pure Python, no Lean needed. Used to benchmark & develop counterexample search
over the public problem set (ground-truth `answer` field available).
"""
from __future__ import annotations

import itertools
import re


# ── Parsing: equation text -> (var_names, lhs_ast, rhs_ast) ──
# AST node: str (variable name) | ("op", left, right)

DIAMOND = "◇"


def _normalize(text: str) -> str:
    return text.replace("*", DIAMOND)


def parse_variables(text: str) -> list[str]:
    seen, out = set(), []
    for v in re.findall(r"\b([a-z])\b", text):
        if v not in seen:
            seen.add(v)
            out.append(v)
    return out


def _parse_expr(s: str, vars_set: set[str]):
    s = s.strip()
    # strip fully-enclosing parens
    while len(s) >= 2 and s[0] == "(" and s[-1] == ")":
        depth = 0
        matched = True
        for i, c in enumerate(s):
            if c == "(":
                depth += 1
            elif c == ")":
                depth -= 1
            if depth == 0 and i < len(s) - 1:
                matched = False
                break
        if matched:
            s = s[1:-1].strip()
        else:
            break
    # find top-level (last) diamond operator
    depth = 0
    last_op = -1
    for i, c in enumerate(s):
        if c == "(":
            depth += 1
        elif c == ")":
            depth -= 1
        elif c == DIAMOND and depth == 0:
            last_op = i
    if last_op >= 0:
        left = _parse_expr(s[:last_op], vars_set)
        right = _parse_expr(s[last_op + 1:], vars_set)
        return ("op", left, right)
    s = s.strip()
    if len(s) == 1 and s in vars_set:
        return s
    raise ValueError(f"cannot parse: {s!r}")


def parse_equation(text: str):
    text = _normalize(text)
    variables = parse_variables(text)
    lhs_str, rhs_str = text.split("=", 1)
    vs = set(variables)
    return variables, _parse_expr(lhs_str, vs), _parse_expr(rhs_str, vs)


# ── Codegen: compile an equation to fast holds/violated checkers ──

def _expr_code(ast) -> str:
    if isinstance(ast, str):
        return ast  # variable name -> local
    _, l, r = ast
    return f"t[{_expr_code(l)}][{_expr_code(r)}]"


class Equation:
    """Compiled equation: .holds(t,n) and .violated(t,n) run native code."""

    __slots__ = ("variables", "lhs", "rhs", "_holds", "_violated", "text")

    def __init__(self, text: str):
        self.text = text
        self.variables, self.lhs, self.rhs = parse_equation(text)
        k = len(self.variables)
        vs = self.variables
        lhs_c = _expr_code(self.lhs)
        rhs_c = _expr_code(self.rhs)
        if k == 0:
            unpack = "_"
        elif k == 1:
            unpack = vs[0] + ","  # unpack 1-tuple: `for x, in assigns`
        else:
            unpack = ", ".join(vs)
        # holds: True iff lhs==rhs for every assignment
        src_holds = (
            f"def _holds(t, assigns):\n"
            f"    for {unpack} in assigns:\n"
            f"        if {lhs_c} != {rhs_c}:\n"
            f"            return False\n"
            f"    return True\n"
        )
        # violated: True iff some assignment has lhs!=rhs
        src_viol = (
            f"def _violated(t, assigns):\n"
            f"    for {unpack} in assigns:\n"
            f"        if {lhs_c} != {rhs_c}:\n"
            f"            return True\n"
            f"    return False\n"
        )
        ns: dict = {}
        exec(src_holds, ns)
        exec(src_viol, ns)
        self._holds = ns["_holds"]
        self._violated = ns["_violated"]

    def holds(self, t, assigns) -> bool:
        return self._holds(t, assigns)

    def violated(self, t, assigns) -> bool:
        return self._violated(t, assigns)


_ASSIGN_CACHE: dict[tuple[int, int], list] = {}


def assignments(n: int, k: int):
    key = (n, k)
    a = _ASSIGN_CACHE.get(key)
    if a is None:
        a = list(itertools.product(range(n), repeat=k))
        _ASSIGN_CACHE[key] = a
    return a