rose-ash/shared/sx/resolver.py

"""
Async resolver — walks an s-expression tree, fetches I/O in parallel,
substitutes results, and renders to HTML.

This is the DAG execution engine applied to page rendering.  The strategy:

1. **Walk** the parsed tree and identify I/O nodes (``frag``, ``query``,
   ``action``, ``current-user``, ``htmx-request?``).
2. **Group** independent fetches.
3. **Dispatch** via ``asyncio.gather()`` for maximum parallelism.
4. **Substitute** resolved values back into the tree.
5. **Render** the fully-resolved tree to HTML via the HTML renderer.

Usage::

    from shared.sx import parse
    from shared.sx.resolver import resolve, RequestContext

    expr = parse('''
        (div :class "page"
          (h1 "Blog")
          (raw! (frag "blog" "link-card" :slug "apple")))
    ''')
    ctx = RequestContext(user=current_user, is_htmx=is_htmx_request())
    html = await resolve(expr, ctx=ctx, env={})
"""

from __future__ import annotations

import asyncio
from typing import Any

from .types import Component, Keyword, Lambda, NIL, Symbol
from .evaluator import _eval as _raw_eval, _trampoline

def _eval(expr, env):
    """Evaluate and unwrap thunks — all resolver.py _eval calls are non-tail."""
    return _trampoline(_raw_eval(expr, env))
from .html import render as html_render, _RawHTML
from .primitives_io import (
    IO_PRIMITIVES,
    RequestContext,
    execute_io,
)


# ---------------------------------------------------------------------------
# Public API
# ---------------------------------------------------------------------------

async def resolve(
    expr: Any,
    *,
    ctx: RequestContext | None = None,
    env: dict[str, Any] | None = None,
) -> str:
    """Resolve an s-expression tree and render to HTML.

    1. Collect all I/O nodes from the tree.
    2. Execute them in parallel.
    3. Substitute results.
    4. Render to HTML.
    """
    if ctx is None:
        ctx = RequestContext()
    if env is None:
        env = {}

    # Resolve I/O nodes (may require multiple passes if I/O results
    # contain further I/O references, though typically one pass suffices).
    resolved = await _resolve_tree(expr, env, ctx)

    # Render the fully-resolved tree to HTML
    return html_render(resolved, env)


# ---------------------------------------------------------------------------
# Tree walker — collect, fetch, substitute
# ---------------------------------------------------------------------------

async def _resolve_tree(
    expr: Any,
    env: dict[str, Any],
    ctx: RequestContext,
    max_depth: int = 5,
) -> Any:
    """Resolve I/O nodes in the tree.  Loops up to *max_depth* passes
    in case resolved values introduce new I/O nodes."""
    resolved = expr
    for _ in range(max_depth):
        # Collect I/O nodes
        io_nodes: list[_IONode] = []
        _collect_io(resolved, env, io_nodes)

        if not io_nodes:
            break  # nothing to fetch

        # Execute all I/O in parallel
        results = await asyncio.gather(
            *[_execute_node(node, ctx) for node in io_nodes],
            return_exceptions=True,
        )

        # Build substitution map (node id → result)
        for node, result in zip(io_nodes, results):
            if isinstance(result, BaseException):
                # On error, substitute empty string (graceful degradation)
                node.result = ""
            else:
                node.result = result

        # Substitute results back into tree
        resolved = _substitute(resolved, env, {id(n.expr): n for n in io_nodes})

    return resolved


class _IONode:
    """A collected I/O node from the tree."""
    __slots__ = ("name", "args", "kwargs", "expr", "result")

    def __init__(self, name: str, args: list[Any], kwargs: dict[str, Any], expr: list):
        self.name = name
        self.args = args
        self.kwargs = kwargs
        self.expr = expr  # original list reference for identity-based substitution
        self.result: Any = None


def _collect_io(
    expr: Any,
    env: dict[str, Any],
    out: list[_IONode],
) -> None:
    """Walk the tree and collect I/O nodes into *out*."""
    if not isinstance(expr, list) or not expr:
        return

    head = expr[0]

    if isinstance(head, Symbol) and head.name in IO_PRIMITIVES:
        # Parse args and kwargs from the rest of the expression
        args, kwargs = _parse_io_args(expr[1:], env)
        out.append(_IONode(head.name, args, kwargs, expr))
        return  # don't recurse into I/O node children

    # Recurse into children
    for child in expr:
        if isinstance(child, list):
            _collect_io(child, env, out)


def _parse_io_args(
    exprs: list[Any],
    env: dict[str, Any],
) -> tuple[list[Any], dict[str, Any]]:
    """Split I/O node arguments into positional args and keyword kwargs.

    Evaluates each argument value so variables/expressions are resolved
    before the I/O call.
    """
    args: list[Any] = []
    kwargs: dict[str, Any] = {}
    i = 0
    while i < len(exprs):
        item = exprs[i]
        if isinstance(item, Keyword) and i + 1 < len(exprs):
            kwargs[item.name] = _eval(exprs[i + 1], env)
            i += 2
        else:
            args.append(_eval(item, env))
            i += 1
    return args, kwargs


async def _execute_node(node: _IONode, ctx: RequestContext) -> Any:
    """Execute a single I/O node."""
    return await execute_io(node.name, node.args, node.kwargs, ctx)


def _substitute(
    expr: Any,
    env: dict[str, Any],
    node_map: dict[int, _IONode],
) -> Any:
    """Replace I/O nodes in the tree with their resolved results."""
    if not isinstance(expr, list) or not expr:
        return expr

    # Check if this exact list is an I/O node
    node = node_map.get(id(expr))
    if node is not None:
        result = node.result
        # Fragment results are HTML strings — wrap as _RawHTML to prevent escaping
        if node.name == "frag" and isinstance(result, str):
            return _RawHTML(result)
        return result

    # Recurse into children
    return [_substitute(child, env, node_map) for child in expr]