engine.py

"""Eleusis rule engine — extracted verbatim from the eleusis env (v0.1.48).

Safe predicate compilation + behavioral equivalence checking, so human guesses
in the Space are judged by the exact machinery the model was evaluated with.
"""

from __future__ import annotations

import ast
import re
import textwrap
from typing import Any, Callable, Sequence

from cards import Card, deck, parse_card

CompiledRule = Callable[[Card, Sequence[Card]], bool]

_SAFE_BUILTINS: dict[str, Any] = {
    "abs": abs,
    "all": all,
    "any": any,
    "bool": bool,
    "dict": dict,
    "enumerate": enumerate,
    "int": int,
    "len": len,
    "list": list,
    "max": max,
    "min": min,
    "range": range,
    "set": set,
    "sorted": sorted,
    "str": str,
    "sum": sum,
    "tuple": tuple,
    "zip": zip,
}

_ALLOWED_CALL_METHODS = {"count", "endswith", "index", "startswith"}

def _rules_match_on_mainlines(
    guessed_rule: CompiledRule,
    target_rule: Rule,
    mainlines: Sequence[Sequence[Card]],
) -> bool:
    for mainline in mainlines:
        for symbol in deck():
            card = parse_card(symbol)
            try:
                guessed_value = guessed_rule(card, mainline)
            except Exception:
                return False
            if guessed_value != target_rule.evaluate(card, mainline):
                return False
    return True


def compile_python_rule(source: str) -> CompiledRule:
    function_source = _as_function_source(source)
    tree = ast.parse(function_source, mode="exec")
    _validate_rule_ast(tree)
    namespace: dict[str, Any] = {"__builtins__": _SAFE_BUILTINS}
    exec(
        compile(tree, "<eleusis_rule>", "exec"),
        namespace,
        namespace,
    )
    fn = namespace.get("evaluate_rule")
    if not callable(fn):
        raise ValueError("Rule code must define evaluate_rule(card, mainline).")

    def evaluate(card: Card, mainline: Sequence[Card]) -> bool:
        return bool(fn(card, list(mainline)))

    return evaluate


def _as_function_source(source: str) -> str:
    code = textwrap.dedent(str(source or "")).strip()
    if not code:
        raise ValueError("Empty rule code.")
    if len(code) > 2_000:
        raise ValueError("Rule code is too long.")
    if re.match(r"^def\s+evaluate_rule\s*\(", code):
        return code
    if re.match(r"^def\s+is_playable\s*\(", code):
        call = (
            "is_playable(card, mainline)"
            if _defines_two_arg_is_playable(code)
            else "is_playable(card)"
        )
        return "\n\n".join(
            [
                code,
                "def evaluate_rule(card, mainline):\n"
                f"    return {call}",
            ]
        )
    if "\n" in code or re.search(r"\breturn\b", code):
        return "def evaluate_rule(card, mainline):\n" + textwrap.indent(code, "    ")
    return "def evaluate_rule(card, mainline):\n" + textwrap.indent(
        f"return bool({code})", "    "
    )


def _defines_two_arg_is_playable(code: str) -> bool:
    tree = ast.parse(code, mode="exec")
    for node in tree.body:
        if isinstance(node, ast.FunctionDef) and node.name == "is_playable":
            return len(node.args.args) >= 2
    return False


def _validate_rule_ast(tree: ast.AST) -> None:
    allowed_nodes = (
        ast.Module,
        ast.FunctionDef,
        ast.arguments,
        ast.arg,
        ast.Return,
        ast.Expr,
        ast.Assign,
        ast.Name,
        ast.Load,
        ast.Store,
        ast.Attribute,
        ast.Constant,
        ast.Compare,
        ast.BoolOp,
        ast.BinOp,
        ast.UnaryOp,
        ast.If,
        ast.IfExp,
        ast.List,
        ast.Tuple,
        ast.Set,
        ast.Dict,
        ast.Subscript,
        ast.Slice,
        ast.Call,
        ast.keyword,
        ast.And,
        ast.Or,
        ast.Not,
        ast.Eq,
        ast.NotEq,
        ast.Lt,
        ast.LtE,
        ast.Gt,
        ast.GtE,
        ast.In,
        ast.NotIn,
        ast.Is,
        ast.IsNot,
        ast.Add,
        ast.Sub,
        ast.Mult,
        ast.Div,
        ast.FloorDiv,
        ast.Mod,
        ast.Pow,
        ast.USub,
        ast.UAdd,
    )
    assigned_names = _assigned_names(tree)
    function_names = _function_names(tree)
    allowed_names = {
        "card",
        "mainline",
        "evaluate_rule",
        "is_playable",
        *assigned_names,
        *function_names,
        *_SAFE_BUILTINS,
    }
    callable_names = {*_SAFE_BUILTINS, *function_names}

    for node in ast.walk(tree):
        if not isinstance(node, allowed_nodes):
            raise ValueError(f"Disallowed Python syntax: {type(node).__name__}")
        if isinstance(node, ast.Attribute) and node.attr.startswith("_"):
            raise ValueError("Private attributes are not allowed.")
        if isinstance(node, ast.Name) and node.id not in allowed_names:
            raise ValueError(f"Unknown name in rule code: {node.id}")
        if isinstance(node, ast.Call):
            if isinstance(node.func, ast.Name):
                if node.func.id not in callable_names:
                    raise ValueError(f"Function is not allowed: {node.func.id}")
            elif isinstance(node.func, ast.Attribute):
                if node.func.attr not in _ALLOWED_CALL_METHODS:
                    raise ValueError(f"Method is not allowed: {node.func.attr}")
            else:
                raise ValueError("Unsupported call target.")


def _representative_mainlines() -> list[list[Card]]:
    cards = {symbol: parse_card(symbol) for symbol in deck()}
    return [
        [],
        [cards["AH"]],
        [cards["2H"]],
        [cards["AC"]],
        [cards["2S"]],
        [cards["5H"]],
        [cards["5C"]],
        [cards["8C"]],
        [cards["QD"]],
        [cards["KS"]],
        [cards["AH"], cards["5C"]],
        [cards["2H"], cards["8C"]],
        [cards["AC"], cards["5H"]],
        [cards["8C"], cards["2D"]],
        [cards["5H"], cards["9C"], cards["KS"]],
        [cards["2H"], cards["8C"], cards["3D"]],
    ]

def guess_matches_rule(guess: str, rule_id: str) -> bool:
    try:
        guessed_rule = compile_python_rule(guess)
        target_rule = get_rule(rule_id)
    except Exception:
        return False

    if _rules_match_on_mainlines(guessed_rule, target_rule, _representative_mainlines()):
        return True

    # In the actual environment God starts with one accepted card before the
    # Scientist acts, so a correct reachable rule may omit an empty-mainline
    # guard like `not mainline or ...`.
    return _rules_match_on_mainlines(
        guessed_rule,
        target_rule,
        [mainline for mainline in _representative_mainlines() if mainline],
    )


def _assigned_names(tree: ast.AST) -> set[str]:
    names: set[str] = set()
    for node in ast.walk(tree):
        if isinstance(node, ast.Assign):
            for target in node.targets:
                names.update(_target_names(target))
    return names


def _function_names(tree: ast.AST) -> set[str]:
    return {node.name for node in ast.walk(tree) if isinstance(node, ast.FunctionDef)}


def _target_names(target: ast.AST) -> set[str]:
    if isinstance(target, ast.Name):
        return {target.id}
    if isinstance(target, (ast.Tuple, ast.List)):
        names: set[str] = set()
        for item in target.elts:
            names.update(_target_names(item))
        return names
    return set()


def get_rule(rule_id: str) -> Rule:
    try:
        return RULES[rule_id]
    except KeyError as exc:
        raise ValueError(f"Unknown rule_id: {rule_id}") from exc


def guess_matches_target(guess: str, target: CompiledRule) -> bool:
    """Behavioral equivalence of a guess against the compiled hidden rule —
    mirrors eleusis.rules.guess_matches_rule (probe battery, with and without
    the empty-mainline guard)."""
    try:
        guessed = compile_python_rule(guess)
    except Exception:
        return False

    class _TargetShim:
        evaluate = staticmethod(target)

    shim = _TargetShim()
    if _rules_match_on_mainlines(guessed, shim, _representative_mainlines()):
        return True
    return _rules_match_on_mainlines(
        guessed,
        shim,
        [m for m in _representative_mainlines() if m],
    )