levain-cli 0.1.0__py3-none-any.whl

levain_recipe/graph.py

@@ -0,0 +1,813 @@
+"""Stage 3: Graph topology check."""
+
+from __future__ import annotations
+
+from collections import defaultdict, deque
+from typing import Any
+
+import tree_sitter
+
+from .models import (
+    Condition,
+    Route,
+    RoutingTable,
+    ValidationError,
+)
+from .parser import (
+    extract_body,
+    find_child,
+    iter_children,
+    node_text,
+    walk_all,
+)
+
+# ── Operator mapping ──
+
+_OP_MAP: dict[str, str] = {
+    "eq": "eq",
+    "neq": "neq",
+    "geq": "gte",
+    "leq": "lte",
+    "lt": "lt",
+    "gt": "gt",
+}
+
+_NEGATE_OP: dict[str, str] = {
+    "eq": "neq",
+    "neq": "eq",
+    "gte": "lt",
+    "lt": "gte",
+    "lte": "gt",
+    "gt": "lte",
+}
+
+
+def check_graph(
+    algo: tree_sitter.Node,
+    variables: list[dict[str, object]],
+    send_fields: dict[str, str] | None = None,
+    *,
+    terminal_labels: set[str] | frozenset[str] = frozenset(),
+) -> list[RoutingTable] | list[ValidationError]:
+    """Extract and validate the graph topology.
+
+    Parameters
+    ----------
+    algo
+        The pcal_algorithm AST node.
+    variables
+        Variable declarations (from parser.extract_variables).
+    send_fields
+        Map of label → state field name for Send fan-out.
+        Extracted from ``@send`` annotations by the profile
+        checker.
+    terminal_labels
+        Labels that count as termination for reachability —
+        currently the set of labels whose ``@node`` function is
+        ``interrupt``.  An interrupt-paused turn IS a clean turn
+        end; the next dispatch resumes the graph.  Without this,
+        a chat-shape recipe with ``Respond -> Wait -> Respond``
+        loop would fail "cannot reach Done."
+
+    Returns
+    -------
+    list[RoutingTable] | list[ValidationError]
+        The routing tables on success, or errors.
+    """
+    body = extract_body(algo)
+    if body is None:
+        return [
+            ValidationError(
+                stage="graph",
+                message="No algorithm body found.",
+                line=algo.start_point[0] + 1,
+            )
+        ]
+
+    labels, label_stmts = _extract_labels(body)
+    if not labels:
+        return [
+            ValidationError(
+                stage="graph",
+                message="No labels found in algorithm body.",
+                line=body.start_point[0] + 1,
+            )
+        ]
+
+    routing = _extract_routing(labels, label_stmts, send_fields or {})
+
+    # Flatten routing tables for structural checks.
+    edges = _edges_from_routing(routing)
+
+    errors: list[ValidationError] = []
+    all_labels = set(labels) | {"Done"}
+    for src, tgt in edges:
+        if tgt not in all_labels:
+            errors.append(
+                ValidationError(
+                    stage="graph",
+                    message=(
+                        f"goto '{tgt}' at label '{src}' "
+                        f"targets an undefined label."
+                    ),
+                )
+            )
+    errors.extend(_check_reachability(labels, edges))
+    errors.extend(
+        _check_termination(labels, edges, terminal_labels=terminal_labels)
+    )
+    errors.extend(_check_condition_completeness(label_stmts))
+    errors.extend(_check_parseable_conditions(label_stmts))
+    errors.extend(_check_assignments(label_stmts, variables))
+    errors.extend(_check_with_types(label_stmts, variables))
+
+    if errors:
+        return errors
+    return routing
+
+
+def _edges_from_routing(
+    routing: list[RoutingTable],
+) -> list[tuple[str, str]]:
+    """Flatten routing tables to (source, target) pairs."""
+    edges: list[tuple[str, str]] = []
+    for table in routing:
+        for route in table.routes:
+            edges.append((table.source, route.target))
+    return edges
+
+
+# ── Label extraction ──
+
+
+def _extract_labels(
+    body: tree_sitter.Node,
+) -> tuple[list[str], dict[str, list[tree_sitter.Node]]]:
+    """Extract label names and their associated statements.
+
+    Returns (ordered_labels, {label: [stmt_nodes]}).
+    """
+    labels: list[str] = []
+    label_stmts: dict[str, list[tree_sitter.Node]] = {}
+
+    children = list(body.children)
+    current_label: str | None = None
+    current_stmts: list[tree_sitter.Node] = []
+
+    i = 0
+    while i < len(children):
+        child = children[i]
+
+        if child.type == "identifier":
+            if i + 1 < len(children) and node_text(children[i + 1]) == ":":
+                if current_label is not None:
+                    label_stmts[current_label] = current_stmts
+                current_label = node_text(child)
+                labels.append(current_label)
+                current_stmts = []
+                i += 2
+                continue
+
+        if current_label is not None and child.type not in (
+            "begin",
+            "comment",
+            ";",
+        ):
+            current_stmts.append(child)
+
+        i += 1
+
+    if current_label is not None:
+        label_stmts[current_label] = current_stmts
+
+    return labels, label_stmts
+
+
+# ── Routing extraction ──
+
+
+def _extract_routing(
+    labels: list[str],
+    label_stmts: dict[str, list[tree_sitter.Node]],
+    send_fields: dict[str, str],
+) -> list[RoutingTable]:
+    """Build routing tables from gotos and if/else."""
+    tables: list[RoutingTable] = []
+
+    for idx, label in enumerate(labels):
+        stmts = label_stmts.get(label, [])
+        raw_routes = _routes_from_stmts(stmts)
+        send_field = send_fields.get(label)
+
+        has_while = any(s.type == "pcal_while" for s in stmts)
+        if has_while:
+            # While loop: add self-edge (loop-back) and
+            # fallthrough (exit) to next label.
+            raw_routes.append((label, []))
+            if idx + 1 < len(labels):
+                raw_routes.append((labels[idx + 1], []))
+        elif not raw_routes:
+            # Implicit fallthrough: no routes at all.
+            if idx + 1 < len(labels):
+                raw_routes.append((labels[idx + 1], []))
+
+        routes = [
+            Route(
+                target=target,
+                conditions=conds,
+                send_field=send_field,
+            )
+            for target, conds in raw_routes
+        ]
+
+        if routes:
+            tables.append(RoutingTable(source=label, routes=routes))
+
+    return tables
+
+
+# Type alias for intermediate route data.
+_RawRoute = tuple[str, list[Condition]]
+
+
+def _routes_from_stmts(
+    stmts: list[tree_sitter.Node],
+) -> list[_RawRoute]:
+    """Collect (target, conditions) from goto/if nodes."""
+    results: list[_RawRoute] = []
+
+    for stmt in stmts:
+        if stmt.type == "pcal_goto":
+            target_node = find_child(stmt, "identifier")
+            if target_node:
+                results.append((node_text(target_node), []))
+
+        elif stmt.type == "pcal_if":
+            results.extend(_routes_from_if(stmt))
+
+        elif stmt.type == "pcal_while":
+            results.extend(_routes_from_while(stmt))
+
+        elif stmt.type == "pcal_with":
+            inner = [
+                c
+                for c in stmt.children
+                if c.type
+                not in (
+                    "with",
+                    "do",
+                    "identifier",
+                    "set_in",
+                    "finite_set_literal",
+                    "pcal_end_with",
+                    ";",
+                )
+            ]
+            results.extend(_routes_from_stmts(inner))
+
+    return results
+
+
+_IfBranch = tuple[list[tree_sitter.Node] | None, list[tree_sitter.Node]]
+
+
+def _if_branches(if_node: tree_sitter.Node) -> list[_IfBranch]:
+    """Split a ``pcal_if`` into its branches.
+
+    tree-sitter-tlaplus emits ``elsif`` / ``else`` as *flat sibling
+    tokens* of the ``pcal_if`` node — there is no ``pcal_elsif``
+    subnode — so branch boundaries must be re-derived from the token
+    stream: each conditioned branch is its keyword, the condition
+    parts up to ``then``, and the statements up to the next boundary.
+    Returns ``(condition_parts, stmts)`` per branch, with
+    ``condition_parts=None`` for the else branch.
+
+    The old walk assumed nested ``pcal_elsif`` nodes and swallowed
+    every elsif (keyword, condition, and body) into the first branch's
+    statements — which flattened every elsif chain into "first
+    condition ANDed onto all branches" and left whole subgraphs
+    unreachable at runtime.
+    """
+    children = list(if_node.children)
+    branches: list[_IfBranch] = []
+    i, n = 0, len(children)
+    while i < n:
+        child = children[i]
+        if child.type in ("if", "elsif"):
+            i += 1
+            cond_parts: list[tree_sitter.Node] = []
+            while i < n and children[i].type != "then":
+                cond_parts.append(children[i])
+                i += 1
+            i += 1  # consume 'then'
+            stmts: list[tree_sitter.Node] = []
+            while i < n and children[i].type not in (
+                "elsif",
+                "else",
+                "pcal_end_if",
+            ):
+                stmts.append(children[i])
+                i += 1
+            branches.append((cond_parts, stmts))
+            continue
+        if child.type == "else":
+            i += 1
+            stmts = []
+            while i < n and children[i].type != "pcal_end_if":
+                stmts.append(children[i])
+                i += 1
+            branches.append((None, stmts))
+            continue
+        i += 1
+    return branches
+
+
+def _routes_from_if(
+    if_node: tree_sitter.Node,
+) -> list[_RawRoute]:
+    r"""Extract structured routes from an if/elsif/else chain.
+
+    Branch *k*'s routes carry the negation of every earlier branch
+    condition ANDed with its own (Route conditions are ANDed); the
+    else branch carries all the negations.  A disjunctive condition
+    (``a = TRUE \/ b = TRUE``) expands into one route per disjunct —
+    its negation is the conjunction of the negated disjuncts
+    (De Morgan), which fits the ANDed-conditions model directly.
+
+    A conditioned branch whose condition cannot be parsed contributes
+    no routes — never an unconditional catch-all —
+    ``_check_parseable_conditions`` rejects the spec loudly instead.
+    """
+    results: list[_RawRoute] = []
+    neg_stack: list[Condition] = []
+    for cond_parts, stmts in _if_branches(if_node):
+        inner_routes = _routes_from_stmts(stmts)
+        if cond_parts is None:
+            for target, inner_conds in inner_routes:
+                results.append((target, [*neg_stack, *inner_conds]))
+            continue
+        disjuncts = _parse_condition_list(cond_parts)
+        if disjuncts is None:
+            continue
+        for target, inner_conds in inner_routes:
+            for disjunct in disjuncts:
+                results.append((target, [*neg_stack, disjunct, *inner_conds]))
+        neg_stack.extend(_negate(d) for d in disjuncts)
+    return results
+
+
+def _routes_from_while(
+    while_node: tree_sitter.Node,
+) -> list[_RawRoute]:
+    """Extract routes from a while loop body."""
+    body_stmts: list[tree_sitter.Node] = []
+    in_body = False
+    for child in while_node.children:
+        if child.type == "do":
+            in_body = True
+        elif in_body and child.type != "pcal_end_while":
+            body_stmts.append(child)
+    return _routes_from_stmts(body_stmts)
+
+
+# ── Condition parsing ──
+
+
+def _parse_condition_list(
+    parts: list[tree_sitter.Node],
+) -> list[Condition] | None:
+    r"""Parse a branch condition into its disjuncts.
+
+    A simple comparison yields one Condition; a top-level ``\/``
+    chain yields one per disjunct (each must itself be a simple
+    comparison).  Anything else — conjunctions, bare booleans,
+    prefix negation — is unsupported and yields None so the caller
+    fails loudly instead of guessing.
+    """
+    for part in parts:
+        if part.type == "bound_infix_op":
+            return _parse_disjuncts(part)
+    return None
+
+
+def _parse_disjuncts(
+    node: tree_sitter.Node,
+) -> list[Condition] | None:
+    children = list(node.children)
+    if len(children) == 3 and children[1].type == "lor":
+        lhs, _op, rhs = children
+        if lhs.type != "bound_infix_op" or rhs.type != "bound_infix_op":
+            return None
+        left = _parse_disjuncts(lhs)
+        right = _parse_disjuncts(rhs)
+        if left is None or right is None:
+            return None
+        return [*left, *right]
+    cond = _parse_infix(node)
+    return [cond] if cond is not None else None
+
+
+def _parse_infix(
+    node: tree_sitter.Node,
+) -> Condition | None:
+    """Parse ``field OP value`` from a bound_infix_op."""
+    children = list(node.children)
+    if len(children) != 3:
+        return None
+
+    lhs, op_node, rhs = children
+
+    # Determine field name (LHS).
+    field_name: str | None = None
+    if lhs.type == "identifier_ref":
+        field_name = node_text(lhs)
+    if field_name is None:
+        return None
+
+    # Determine operator.
+    op = _OP_MAP.get(op_node.type)
+    if op is None:
+        return None
+
+    # Determine value (RHS).
+    value = _parse_value(rhs)
+    if value is None:
+        # RHS might be an identifier_ref (e.g., MaxIterations).
+        if rhs.type == "identifier_ref":
+            value = node_text(rhs)
+        else:
+            return None
+
+    return Condition(field=field_name, op=op, value=value)
+
+
+def _parse_value(node: tree_sitter.Node) -> Any:
+    """Extract a literal value from an AST node."""
+    if node.type == "boolean":
+        return node_text(node).strip() == "TRUE"
+    if node.type == "nat_number":
+        return int(node_text(node))
+    if node.type == "string":
+        return node_text(node).strip('"')
+    return None
+
+
+def _negate(cond: Condition) -> Condition:
+    """Return the logical negation of a condition."""
+    neg_op = _NEGATE_OP.get(cond.op)
+    if neg_op is None:
+        raise ValueError(f"Cannot negate operator {cond.op!r}")
+    return Condition(field=cond.field, op=neg_op, value=cond.value)
+
+
+# ── Reachability ──
+
+
+def _check_reachability(
+    labels: list[str],
+    edges: list[tuple[str, str]],
+) -> list[ValidationError]:
+    """BFS from first label: all labels must be reachable."""
+    if not labels:
+        return []
+
+    adj: dict[str, set[str]] = defaultdict(set)
+    for source, target in edges:
+        adj[source].add(target)
+
+    reachable: set[str] = set()
+    queue: deque[str] = deque([labels[0]])
+    while queue:
+        node = queue.popleft()
+        if node in reachable:
+            continue
+        reachable.add(node)
+        for target in adj[node]:
+            if target not in reachable:
+                queue.append(target)
+
+    unreachable = set(labels) - reachable
+
+    errors: list[ValidationError] = []
+    for label in sorted(unreachable):
+        errors.append(
+            ValidationError(
+                stage="graph",
+                message=(
+                    f"Label '{label}' is not reachable from '{labels[0]}'."
+                ),
+            )
+        )
+    return errors
+
+
+def _check_termination(
+    labels: list[str],
+    edges: list[tuple[str, str]],
+    *,
+    terminal_labels: set[str] | frozenset[str] = frozenset(),
+) -> list[ValidationError]:
+    """Reverse BFS from Done (or any terminal label): all labels reach a sink.
+
+    ``terminal_labels`` lists labels that count as termination
+    for reachability purposes — used for ``interrupt`` nodes,
+    which pause the graph indefinitely waiting for the next
+    user message.  An interrupt-paused turn IS a clean turn
+    end; the next turn re-enters the graph from the interrupt.
+    """
+    reverse_adj: dict[str, set[str]] = defaultdict(set)
+    for source, target in edges:
+        reverse_adj[target].add(source)
+
+    can_reach_done: set[str] = set()
+    seeds: list[str] = ["Done", *terminal_labels]
+    queue: deque[str] = deque(seeds)
+    while queue:
+        node = queue.popleft()
+        if node in can_reach_done:
+            continue
+        can_reach_done.add(node)
+        for source in reverse_adj[node]:
+            if source not in can_reach_done:
+                queue.append(source)
+
+    stuck = set(labels) - can_reach_done - {"Done"}
+
+    errors: list[ValidationError] = []
+    targets = {t for _, t in edges}
+    has_done_path = "Done" in targets or "Done" in set(labels)
+    # At least one sink (Done or an interrupt/terminal label) must
+    # exist.  Per-node reachability is enforced by the ``stuck``
+    # set above — this outer check only verifies that a sink
+    # exists at all so the graph has somewhere to terminate.
+    has_terminal_path = bool(terminal_labels)
+    if not has_done_path and not has_terminal_path:
+        errors.append(
+            ValidationError(
+                stage="graph",
+                message=(
+                    "No transition targets 'Done'. "
+                    "The algorithm cannot terminate."
+                ),
+            )
+        )
+
+    for label in sorted(stuck):
+        errors.append(
+            ValidationError(
+                stage="graph",
+                message=(f"Label '{label}' cannot reach 'Done'."),
+            )
+        )
+    return errors
+
+
+# ── Condition completeness ──
+
+
+def _check_condition_completeness(
+    label_stmts: dict[str, list[tree_sitter.Node]],
+) -> list[ValidationError]:
+    """Every if must have an else branch."""
+    errors: list[ValidationError] = []
+
+    for label, stmts in label_stmts.items():
+        for stmt in stmts:
+            for node in walk_all(stmt):
+                if node.type not in ("pcal_if", "pcal_elsif"):
+                    continue
+                has_else = any(
+                    c.type in ("else", "pcal_elsif") for c in node.children
+                )
+                if not has_else:
+                    construct = "elsif" if node.type == "pcal_elsif" else "if"
+                    errors.append(
+                        ValidationError(
+                            stage="graph",
+                            message=(
+                                f"'{construct}' at label "
+                                f"'{label}' has no 'else' "
+                                f"branch. All conditionals "
+                                f"must have a default path."
+                            ),
+                            line=(node.start_point[0] + 1),
+                        )
+                    )
+
+    return errors
+
+
+def _check_parseable_conditions(
+    label_stmts: dict[str, list[tree_sitter.Node]],
+) -> list[ValidationError]:
+    """Reject if/elsif conditions the router cannot parse.
+
+    Bare boolean identifiers (``if flag then``), prefix
+    negation (``if ~flag then``), and compound conjunctions
+    are not supported.  Authors must use explicit comparisons
+    like ``flag = TRUE`` or ``flag = FALSE``.
+    """
+    errors: list[ValidationError] = []
+
+    for label, stmts in label_stmts.items():
+        for stmt in stmts:
+            for node in walk_all(stmt):
+                if node.type != "pcal_if":
+                    continue
+                # Every branch of the chain — the old walk only ever
+                # saw the first ``if`` condition, so unparseable elsif
+                # conditions slipped through while their branches were
+                # silently dropped from routing.
+                for cond_parts, _branch_stmts in _if_branches(node):
+                    if cond_parts is None:
+                        continue
+                    meaningful = [
+                        p for p in cond_parts if p.type not in ("(", ")")
+                    ]
+                    if (
+                        meaningful
+                        and _parse_condition_list(cond_parts) is None
+                    ):
+                        snippet = " ".join(node_text(p) for p in meaningful)
+                        errors.append(
+                            ValidationError(
+                                stage="graph",
+                                message=(
+                                    f"Condition at label '{label}' "
+                                    f"cannot be parsed: {snippet!r}. "
+                                    f"Use explicit comparisons "
+                                    f"(e.g. x = TRUE), optionally "
+                                    f"joined with \\/."
+                                ),
+                                line=node.start_point[0] + 1,
+                            )
+                        )
+
+    return errors
+
+
+# ── State consistency ──
+
+
+def _check_assignments(
+    label_stmts: dict[str, list[tree_sitter.Node]],
+    variables: list[dict[str, object]],
+) -> list[ValidationError]:
+    """Every assignment must target a declared variable."""
+    var_names = {str(v["name"]) for v in variables}
+    errors: list[ValidationError] = []
+
+    for label, stmts in label_stmts.items():
+        for stmt in stmts:
+            for node in walk_all(stmt):
+                if node.type != "pcal_assign":
+                    continue
+                lhs = find_child(node, "pcal_lhs")
+                if lhs is None:
+                    continue
+                ref = find_child(lhs, "identifier_ref")
+                if ref is None:
+                    continue
+                name = node_text(ref)
+                if name not in var_names:
+                    errors.append(
+                        ValidationError(
+                            stage="graph",
+                            message=(
+                                f"Assignment to '{name}' "
+                                f"at label '{label}' — "
+                                "not a declared variable."
+                            ),
+                            line=node.start_point[0] + 1,
+                        )
+                    )
+
+    return errors
+
+
+# ── With-block type validation ──
+
+_TYPE_TO_NODE_TYPES: dict[str, set[str]] = {
+    "bool": {"boolean"},
+    "int": {"nat_number"},
+    "string": {"string"},
+}
+
+_BOOL_FULL = {"TRUE", "FALSE"}
+
+
+def _check_with_types(
+    label_stmts: dict[str, list[tree_sitter.Node]],
+    variables: list[dict[str, object]],
+) -> list[ValidationError]:
+    """Validate with-block value sets against variable types."""
+    var_types = {str(v["name"]): str(v["type"]) for v in variables}
+    errors: list[ValidationError] = []
+
+    for label, stmts in label_stmts.items():
+        for stmt in stmts:
+            for node in walk_all(stmt):
+                if node.type != "pcal_with":
+                    continue
+                errors.extend(_check_one_with(node, label, var_types))
+
+    return errors
+
+
+def _check_one_with(
+    with_node: tree_sitter.Node,
+    label: str,
+    var_types: dict[str, str],
+) -> list[ValidationError]:
+    """Check a single with block."""
+    errors: list[ValidationError] = []
+    line = with_node.start_point[0] + 1
+
+    value_set = find_child(with_node, "finite_set_literal")
+    if value_set is None:
+        return errors
+
+    with_var = find_child(with_node, "identifier")
+    if with_var is None:
+        return errors
+    with_var_name = node_text(with_var)
+
+    assigned_var = _find_with_target(with_node, with_var_name)
+    if assigned_var is None:
+        return errors
+
+    var_type = var_types.get(assigned_var)
+    if var_type is None:
+        return errors
+
+    expected_node_types = _TYPE_TO_NODE_TYPES.get(var_type, set())
+    if not expected_node_types:
+        return errors
+
+    values = iter_children(value_set, "boolean", "nat_number", "string")
+
+    pre_count = len(errors)
+    for val_node in values:
+        if val_node.type not in expected_node_types:
+            errors.append(
+                ValidationError(
+                    stage="graph",
+                    message=(
+                        f"'with' at label '{label}' has "
+                        f"value {node_text(val_node)!r} "
+                        f"but '{assigned_var}' is "
+                        f"declared as {var_type}."
+                    ),
+                    line=line,
+                )
+            )
+
+    if var_type == "bool" and values and len(errors) == pre_count:
+        present = {node_text(v) for v in values if v.type == "boolean"}
+        missing = _BOOL_FULL - present
+        if missing:
+            errors.append(
+                ValidationError(
+                    stage="graph",
+                    message=(
+                        f"'with' at label '{label}' "
+                        f"assigns to bool "
+                        f"'{assigned_var}' but only "
+                        f"covers {present}. Missing: "
+                        f"{sorted(missing)}."
+                    ),
+                    line=line,
+                )
+            )
+
+    return errors
+
+
+def _find_with_target(
+    with_node: tree_sitter.Node,
+    with_var_name: str,
+) -> str | None:
+    """Find the state variable assigned from the with-var."""
+    for node in walk_all(with_node):
+        if node.type != "pcal_assign":
+            continue
+        lhs = find_child(node, "pcal_lhs")
+        if lhs is None:
+            continue
+        lhs_positions = {
+            (d.start_point, d.end_point)
+            for d in walk_all(lhs)
+            if d.type == "identifier_ref"
+        }
+        rhs_refs = [
+            d
+            for d in walk_all(node)
+            if d.type == "identifier_ref"
+            and (d.start_point, d.end_point) not in lhs_positions
+        ]
+        if not any(node_text(r) == with_var_name for r in rhs_refs):
+            continue
+        lhs_ref = find_child(lhs, "identifier_ref")
+        if lhs_ref is not None:
+            return node_text(lhs_ref)
+    return None