levain-cli 0.1.0__py3-none-any.whl

levain_recipe/parser.py

@@ -0,0 +1,299 @@
+"""Stage 1: Parse PlusCal source with tree-sitter-tlaplus."""
+
+from __future__ import annotations
+
+import warnings
+from collections.abc import Iterator
+from typing import Any
+
+import tree_sitter
+import tree_sitter_tlaplus as tstla
+
+from .models import ValidationError
+
+# Module-level singleton -- built once, reused across calls.
+# tree-sitter-tlaplus returns an int pointer; suppress the
+# deprecation warning until it ships a PyCapsule API.
+with warnings.catch_warnings():
+    warnings.simplefilter("ignore", DeprecationWarning)
+    LANGUAGE = tree_sitter.Language(tstla.language())
+
+
+def _make_parser() -> tree_sitter.Parser:
+    return tree_sitter.Parser(LANGUAGE)
+
+
+def parse(
+    source: str,
+) -> tuple[tree_sitter.Node, str] | list[ValidationError]:
+    """Parse PlusCal source and return the algorithm AST node.
+
+    Parameters
+    ----------
+    source
+        Full TLA+ module text containing a PlusCal algorithm.
+
+    Returns
+    -------
+    tuple[Node, str] | list[ValidationError]
+        On success: (pcal_algorithm node, module_name).
+        On failure: list of parse errors.
+    """
+    parser = _make_parser()
+    tree = parser.parse(bytes(source, "utf-8"))
+    root = tree.root_node
+
+    errors = _collect_errors(root)
+    if errors:
+        return errors
+
+    module_name = _extract_module_name(root)
+    if module_name is None:
+        return [
+            ValidationError(
+                stage="parse",
+                message="No TLA+ module found.",
+                line=1,
+            )
+        ]
+
+    algo = _find_algorithm(root)
+    if algo is None:
+        return [
+            ValidationError(
+                stage="parse",
+                message="No PlusCal algorithm block found.",
+                line=1,
+            )
+        ]
+
+    return (algo, module_name)
+
+
+def _collect_errors(
+    root: tree_sitter.Node,
+) -> list[ValidationError]:
+    """Walk the tree and collect ERROR / MISSING nodes."""
+    errors: list[ValidationError] = []
+
+    def _walk(node: tree_sitter.Node) -> None:
+        if node.type == "ERROR" or node.is_missing:
+            line = node.start_point[0] + 1
+            text = (node.text or b"").decode("utf-8", errors="replace")
+            snippet = text[:60] if text else ""
+            msg = (
+                f"Syntax error near: {snippet!r}"
+                if snippet
+                else "Unexpected syntax"
+            )
+            errors.append(
+                ValidationError(stage="parse", message=msg, line=line)
+            )
+        for child in node.children:
+            _walk(child)
+
+    _walk(root)
+    return errors
+
+
+def _extract_module_name(
+    root: tree_sitter.Node,
+) -> str | None:
+    """Extract the TLA+ module name from the header."""
+    for node in root.children:
+        if node.type == "module":
+            for child in node.children:
+                if child.type == "identifier":
+                    return (child.text or b"").decode("utf-8")
+    return None
+
+
+def _find_algorithm(
+    root: tree_sitter.Node,
+) -> tree_sitter.Node | None:
+    """Find the pcal_algorithm node inside a block_comment."""
+    for node in root.children:
+        if node.type == "module":
+            for child in node.children:
+                if child.type == "block_comment":
+                    for gc in child.children:
+                        if gc.type == "pcal_algorithm":
+                            return gc
+    return None
+
+
+# ── AST helpers used by later stages ──
+
+
+def iter_children(
+    node: tree_sitter.Node,
+    *types: str,
+) -> list[tree_sitter.Node]:
+    """Return direct children matching any of the given types."""
+    return [c for c in node.children if c.type in types]
+
+
+def find_child(
+    node: tree_sitter.Node,
+    node_type: str,
+) -> tree_sitter.Node | None:
+    """Return the first direct child of the given type."""
+    for c in node.children:
+        if c.type == node_type:
+            return c
+    return None
+
+
+def node_text(node: tree_sitter.Node) -> str:
+    """Decode a node's text as UTF-8."""
+    return (node.text or b"").decode("utf-8")
+
+
+def walk_all(
+    node: tree_sitter.Node,
+) -> Iterator[tree_sitter.Node]:
+    """Recursively yield all descendant nodes (pre-order)."""
+    yield node
+    for c in node.children:
+        yield from walk_all(c)
+
+
+def extract_variables(
+    algo: tree_sitter.Node,
+) -> list[dict[str, Any]]:
+    r"""Extract variable declarations from pcal_var_decls.
+
+    Returns a list of dicts: ``{name, type, default, line,
+    comments}``.  ``comments`` is the list of comment texts that
+    appeared between the previous declaration (or the
+    ``variables`` keyword) and this one — used downstream to
+    associate ``\* @reducer ...`` annotations with their variable.
+    """
+    var_decls = find_child(algo, "pcal_var_decls")
+    if var_decls is None:
+        return []
+
+    results: list[dict[str, Any]] = []
+    pending: list[str] = []
+    for child in var_decls.children:
+        if child.type == "comment":
+            pending.append(node_text(child))
+            continue
+        if child.type != "pcal_var_decl":
+            continue
+        name_node = find_child(child, "identifier")
+        if name_node is None:
+            pending = []
+            continue
+        name = node_text(name_node)
+        val_type, default = _infer_var_type(child)
+        results.append(
+            {
+                "name": name,
+                "type": val_type,
+                "default": default,
+                "line": child.start_point[0] + 1,
+                "comments": pending,
+            }
+        )
+        pending = []
+    return results
+
+
+def _infer_var_type(
+    decl: tree_sitter.Node,
+) -> tuple[str, Any]:
+    """Infer the type and default value from a pcal_var_decl."""
+    for child in decl.children:
+        if child.type == "nat_number":
+            return ("int", int(node_text(child)))
+        if child.type == "boolean":
+            val = node_text(child).strip()
+            return ("bool", val == "TRUE")
+        if child.type == "string":
+            raw = node_text(child)
+            return ("string", raw.strip('"'))
+        if child.type == "tuple_literal":
+            # Only empty tuples ``<<>>`` are supported as
+            # defaults; non-empty ones would need typed elements
+            # which the DSL doesn't model today.
+            if _is_empty_tuple(child):
+                return ("list", [])
+    return ("unknown", None)
+
+
+def _is_empty_tuple(node: tree_sitter.Node) -> bool:
+    """True iff a ``tuple_literal`` has no value children.
+
+    ``<<>>`` parses to a ``tuple_literal`` whose only children
+    are the ``langle_bracket`` / ``rangle_bracket`` delimiters;
+    anything else (``<<1, 2>>``) adds value children in between.
+    """
+    for child in node.children:
+        if child.type not in ("langle_bracket", "rangle_bracket"):
+            return False
+    return True
+
+
+def extract_defines(
+    algo: tree_sitter.Node,
+) -> list[str]:
+    """Extract operator names from the define block."""
+    defs = find_child(algo, "pcal_definitions")
+    if defs is None:
+        return []
+    names: list[str] = []
+    for op in iter_children(defs, "operator_definition"):
+        id_node = find_child(op, "identifier")
+        if id_node is not None:
+            names.append(node_text(id_node))
+    return names
+
+
+def extract_constants(
+    algo: tree_sitter.Node,
+) -> dict[str, Any]:
+    """Extract simple constant definitions from the define block.
+
+    Returns a dict mapping names to their resolved values
+    for definitions like ``MaxIterations == 10``.  Complex
+    expressions (operators, set membership) are skipped.
+    """
+    defs = find_child(algo, "pcal_definitions")
+    if defs is None:
+        return {}
+
+    constants: dict[str, Any] = {}
+    for op in iter_children(defs, "operator_definition"):
+        id_node = find_child(op, "identifier")
+        if id_node is None:
+            continue
+        name = node_text(id_node)
+
+        # Find the value node: the child after def_eq
+        # that is a simple literal.
+        children = list(op.children)
+        after_eq = False
+        for child in children:
+            if child.type == "def_eq":
+                after_eq = True
+                continue
+            if not after_eq:
+                continue
+            # Only extract simple literals.
+            if child.type == "nat_number":
+                constants[name] = int(node_text(child))
+            elif child.type == "boolean":
+                constants[name] = node_text(child) == "TRUE"
+            elif child.type == "string":
+                constants[name] = node_text(child).strip('"')
+            break
+
+    return constants
+
+
+def extract_body(
+    algo: tree_sitter.Node,
+) -> tree_sitter.Node | None:
+    """Return the pcal_algorithm_body node."""
+    return find_child(algo, "pcal_algorithm_body")

levain_recipe/profile.py

@@ -0,0 +1,389 @@
+"""Stage 2: Profile check -- reject disallowed PlusCal constructs."""
+
+from __future__ import annotations
+
+import json
+import re
+from typing import Any
+
+import tree_sitter
+
+from .known_recipes import CALLABLE_RECIPES
+from .models import NodeAnnotation, ValidationError
+from .parser import (
+    LANGUAGE,
+    extract_body,
+    node_text,
+)
+
+# ── Rejected constructs (tree-sitter query) ──
+
+# Each capture name maps to an error message.  The query
+# matches all disallowed AST node types in a single pass.
+_REJECTED_MESSAGES: dict[str, str] = {
+    "process": "Concurrent processes are not allowed.",
+    "procedure": "Procedures are not allowed.",
+    "macro": "Macros are not allowed.",
+    "return_": "'return' is not allowed (no procedures).",
+    "macro_call": "Macro calls are not allowed.",
+    "await_": ("'await'/'when' synchronization is not allowed."),
+    "set_comp": (
+        r"Set comprehensions {x \in S : P(x)} "
+        "are not allowed."
+    ),
+    "func_lit": (
+        r"Function definitions [x \in S |-> e] "
+        "are not allowed."
+    ),
+    "choose": "'CHOOSE' is not allowed.",
+    "quantifier": (r"Quantifiers (\A, \E) are not allowed."),
+    "either": "'either/or' nondeterminism is not allowed.",
+}
+
+_REJECTED_QUERY = tree_sitter.Query(
+    LANGUAGE,
+    """
+    (pcal_process) @process
+    (pcal_procedure) @procedure
+    (pcal_macro) @macro
+    (pcal_return) @return_
+    (pcal_macro_call) @macro_call
+    (pcal_await) @await_
+    (set_filter) @set_comp
+    (function_literal) @func_lit
+    (choose) @choose
+    (bounded_quantification) @quantifier
+    (unbounded_quantification) @quantifier
+    (pcal_either) @either
+    """,
+)
+
+# ── EXTENDS check (tree-sitter query) ──
+
+ALLOWED_EXTENDS: frozenset[str] = frozenset({"Naturals"})
+
+_EXTENDS_QUERY = tree_sitter.Query(
+    LANGUAGE,
+    "(extends (identifier_ref) @ext_module)",
+)
+
+# ── Annotation regex patterns ──
+
+_RE_NODE = re.compile(r"\\[*]\s*@node\s+({.*})\s*$")
+_RE_PROMPT = re.compile(r'\\[*]\s*@prompt\s+"((?:[^"\\]|\\.)*)"\s*$')
+_RE_TOOLS = re.compile(r"\\[*]\s*@allowed_tools\s+(\[.*\])\s*$")
+_RE_MCP = re.compile(r"\\[*]\s*@mcp_servers\s+(\[.*\])\s*$")
+_RE_PERM = re.compile(r'\\[*]\s*@permission_mode\s+"(\w+)"\s*$')
+
+# Valid function reference: bare name (built-in),
+# module:function, or dotted.module:function (user code).
+_RE_FUNC_REF = re.compile(
+    r"^[a-zA-Z_]\w*(?:\.[a-zA-Z_]\w*)*(?::[a-zA-Z_]\w*)?$"
+)
+
+# @send annotation: fan-out over a list state field.
+_RE_SEND = re.compile(r'\\[*]\s*@send\s+"(\w+)"\s*$')
+
+
+def check_profile(
+    algo: tree_sitter.Node,
+    module_root: tree_sitter.Node,
+) -> list[ValidationError] | tuple[list[NodeAnnotation], dict[str, str]]:
+    """Validate the spec against the restricted PlusCal profile.
+
+    Parameters
+    ----------
+    algo
+        The pcal_algorithm AST node.
+    module_root
+        The top-level source_file node (for EXTENDS check).
+
+    Returns
+    -------
+    list[ValidationError] | tuple[list[NodeAnnotation], dict[str, str]]
+        Errors if violations found, otherwise a tuple of
+        (annotations, send_fields) where send_fields maps
+        label → state field name for Send fan-out.
+    """
+    errors: list[ValidationError] = []
+
+    errors.extend(_check_extends(module_root))
+    errors.extend(_check_rejected(algo))
+
+    if errors:
+        return errors
+
+    # Extract and validate annotations.
+    body = extract_body(algo)
+    if body is None:
+        return [
+            ValidationError(
+                stage="profile",
+                message="No algorithm body (begin..end) found.",
+                line=algo.start_point[0] + 1,
+            )
+        ]
+
+    return _extract_annotations(body)
+
+
+def _check_extends(
+    root: tree_sitter.Node,
+) -> list[ValidationError]:
+    """Verify only allowed modules are extended."""
+    errors: list[ValidationError] = []
+    cursor = tree_sitter.QueryCursor(_EXTENDS_QUERY)
+    for _, captures in cursor.matches(root):
+        for node in captures.get("ext_module", []):
+            name = node_text(node)
+            if name not in ALLOWED_EXTENDS:
+                errors.append(
+                    ValidationError(
+                        stage="profile",
+                        message=(
+                            f"EXTENDS {name} is not allowed."
+                            f" Only {ALLOWED_EXTENDS}"
+                            " permitted."
+                        ),
+                        line=node.start_point[0] + 1,
+                    )
+                )
+    return errors
+
+
+def _check_rejected(
+    algo: tree_sitter.Node,
+) -> list[ValidationError]:
+    """Run the rejected-construct query in one pass."""
+    errors: list[ValidationError] = []
+    cursor = tree_sitter.QueryCursor(_REJECTED_QUERY)
+    for _, captures in cursor.matches(algo):
+        for capture_name, nodes in captures.items():
+            msg = _REJECTED_MESSAGES.get(capture_name, "")
+            for node in nodes:
+                errors.append(
+                    ValidationError(
+                        stage="profile",
+                        message=msg,
+                        line=node.start_point[0] + 1,
+                    )
+                )
+    return errors
+
+
+def _extract_annotations(
+    body: tree_sitter.Node,
+) -> list[ValidationError] | tuple[list[NodeAnnotation], dict[str, str]]:
+    """Extract @node and @send annotations from comments.
+
+    Returns (annotations, send_fields) on success, or
+    a list of errors.
+    """
+    annotations: list[NodeAnnotation] = []
+    send_fields: dict[str, str] = {}
+    errors: list[ValidationError] = []
+
+    children = list(body.children)
+    pending_comments: list[tuple[int, str]] = []
+
+    i = 0
+    while i < len(children):
+        child = children[i]
+
+        if child.type == "comment":
+            line = child.start_point[0] + 1
+            text = node_text(child)
+            pending_comments.append((line, text))
+            i += 1
+            continue
+
+        if child.type == "identifier":
+            if i + 1 < len(children) and node_text(children[i + 1]) == ":":
+                label = node_text(child)
+                line = child.start_point[0] + 1
+
+                if label == "Done":
+                    pending_comments.clear()
+                    i += 2
+                    continue
+
+                result = _parse_annotation_block(label, line, pending_comments)
+                pending_comments.clear()
+
+                if isinstance(result, ValidationError):
+                    errors.append(result)
+                else:
+                    ann, send = result
+                    annotations.append(ann)
+                    if send:
+                        send_fields[label] = send
+
+                i += 2
+                continue
+
+        if child.type not in ("begin", ";"):
+            pending_comments.clear()
+        i += 1
+
+    if errors:
+        return errors
+    return annotations, send_fields
+
+
+def _parse_annotation_block(  # noqa: PLR0911, PLR0915
+    label: str,
+    line: int,
+    comments: list[tuple[int, str]],
+) -> ValidationError | tuple[NodeAnnotation, str | None]:
+    """Parse @-comments into a NodeAnnotation + optional send field."""
+    node_json: dict[str, Any] | None = None
+    prompt: str = ""
+    tools: list[str] = []
+    mcp: list[str] = []
+    perm: str | None = None
+    send: str | None = None
+    node_line: int = line
+
+    for cline, text in comments:
+        m = _RE_NODE.match(text)
+        if m:
+            try:
+                node_json = json.loads(m.group(1))
+            except json.JSONDecodeError as e:
+                return ValidationError(
+                    stage="profile",
+                    message=(f"Invalid @node JSON at label '{label}': {e}"),
+                    line=cline,
+                )
+            node_line = cline
+            continue
+
+        m = _RE_PROMPT.match(text)
+        if m:
+            prompt = m.group(1).replace("\\n", "\n")
+            continue
+
+        m = _RE_TOOLS.match(text)
+        if m:
+            try:
+                tools = json.loads(m.group(1))
+            except json.JSONDecodeError as e:
+                return ValidationError(
+                    stage="profile",
+                    message=(
+                        f"Invalid @allowed_tools JSON at label '{label}': {e}"
+                    ),
+                    line=cline,
+                )
+            continue
+
+        m = _RE_MCP.match(text)
+        if m:
+            try:
+                mcp = json.loads(m.group(1))
+            except json.JSONDecodeError as e:
+                return ValidationError(
+                    stage="profile",
+                    message=(
+                        f"Invalid @mcp_servers JSON at label '{label}': {e}"
+                    ),
+                    line=cline,
+                )
+            if not isinstance(mcp, list) or not all(
+                isinstance(s, str) for s in mcp
+            ):
+                return ValidationError(
+                    stage="profile",
+                    message=(
+                        f"@mcp_servers must be a JSON list of"
+                        f" strings at '{label}'"
+                    ),
+                    line=cline,
+                )
+            continue
+
+        m = _RE_PERM.match(text)
+        if m:
+            perm = m.group(1)
+            continue
+
+        m = _RE_SEND.match(text)
+        if m:
+            send = m.group(1)
+            continue
+
+        if re.search(r"\\[*]\s*@\w+", text):
+            return ValidationError(
+                stage="profile",
+                message=(
+                    f"Unrecognized annotation in comment "
+                    f"near label '{label}': {text!r}"
+                ),
+                line=cline,
+            )
+
+    if node_json is None:
+        return ValidationError(
+            stage="profile",
+            message=(f"Label '{label}' is missing a @node annotation."),
+            line=line,
+        )
+
+    func = node_json.pop("function", None)
+    if not func:
+        return ValidationError(
+            stage="profile",
+            message=(f"@node at label '{label}' is missing a 'function' key."),
+            line=node_line,
+        )
+
+    if not _RE_FUNC_REF.match(func):
+        return ValidationError(
+            stage="profile",
+            message=(
+                f"Invalid function reference '{func}' "
+                f"at label '{label}'. Use 'name' for "
+                "built-ins or 'module:function' for "
+                "user code."
+            ),
+            line=node_line,
+        )
+
+    # Validate "recipe" built-in: requires "name" arg that
+    # references a known recipe.
+    if func == "recipe":
+        name = node_json.get("name") or node_json.get("args", {}).get("name")
+        if not name:
+            return ValidationError(
+                stage="profile",
+                message=(
+                    f"@node at label '{label}' uses function "
+                    f"'recipe' but is missing a 'name' key."
+                ),
+                line=node_line,
+            )
+
+        if name not in CALLABLE_RECIPES:
+            return ValidationError(
+                stage="profile",
+                message=(
+                    f"@node at label '{label}' references "
+                    f"unknown recipe '{name}'. "
+                    f"Known: {', '.join(sorted(CALLABLE_RECIPES))}."
+                ),
+                line=node_line,
+            )
+
+    args: dict[str, Any] = node_json.pop("args", {})
+    args.update(node_json)
+    if prompt:
+        args["prompt"] = prompt
+    if tools:
+        args["allowed_tools"] = tools
+    if mcp:
+        args["mcp_servers"] = mcp
+    if perm:
+        args["permission_mode"] = perm
+
+    annotation = NodeAnnotation(label=label, function=func, args=args)
+    return annotation, send