co-lambda 0.5.0__py3-none-any.whl

co_lambda/_pyast.py

@@ -0,0 +1,418 @@
+"""A Scott-encoded Python AST, generated by reflection on the ``ast`` module, with a decoder.
+
+The compiler's target is a Python AST represented in the pure lambda-calculus by Scott encoding: a
+value of an n-constructor data type is ``lambda h0 ... h_{n-1}. h_tag field0 ... field_k`` (it
+applies the handler for its constructor to its fields). Rather than hand-write a constructor and a
+decoder per node type, both are derived by reflection on ``ast``: ``SUPPORTED`` lists the node
+classes, each class's ``_fields`` gives its arity and field order, and a field is encoded with a
+kind tag (a child node, a list, an int, a string as a list of character codes, a bool, or none) so
+the decoder can rebuild a real ``ast`` node generically with ``cls(*fields)``.
+
+Encoding (a real ``ast`` node to a Scott value) and decoding (a Scott value, run in the interpreter,
+back to a real ``ast`` node) are meta-level: the lambda-calculus does not inspect itself, but the
+interpreter that runs it does, exactly as ``render`` reads the behaviour graph. The decoder reflects
+a Scott value to the host by applying it to handlers whose heads are distinct free de Bruijn
+variables (markers); the weak head normal form is then ``marker child0 ... childk``, a spine whose
+head variable names the constructor and whose arguments are the child nodes.
+"""
+
+from __future__ import annotations
+
+import ast
+import contextlib
+import hashlib
+import struct
+
+from co_lambda._ast import App, Lam, Native, Node, Var, make_app, make_lam, make_var
+from co_lambda._dsl import Builder, app, build, lam
+from co_lambda._codec import church
+from co_lambda._prelude import SCOTT_NIL
+from co_lambda._sugar import cons
+
+# The Python AST node classes the encoding supports, in tag order. Extend as needed; encoding an
+# unsupported class raises, so a gap is loud rather than silent.
+SUPPORTED: "tuple[type[ast.AST], ...]" = (
+    ast.Expression, ast.Module, ast.Expr,
+    ast.Lambda, ast.arguments, ast.arg,
+    ast.FunctionDef, ast.Return, ast.Assign, ast.While, ast.If, ast.Pass,
+    ast.Call, ast.Name, ast.Load, ast.Store, ast.Constant,
+    ast.BinOp, ast.Add, ast.Sub, ast.Mult,
+    ast.Compare, ast.Lt, ast.LtE, ast.Gt, ast.GtE, ast.Eq,
+    ast.Nonlocal, ast.Is,
+    ast.Subscript, ast.Tuple,
+    ast.IfExp, ast.Attribute,
+    ast.ClassDef, ast.AnnAssign,
+)
+_TAG: "dict[type[ast.AST], int]" = {cls: tag for tag, cls in enumerate(SUPPORTED)}
+_ARITY = tuple(len(cls._fields) for cls in SUPPORTED)
+
+# Field kind tags.
+_K_NODE, _K_LIST, _K_INT, _K_STR, _K_BOOL, _K_NONE, _K_IDENT, _K_GENSYM = range(8)
+
+# A gensym table mapping a payload node's identity to a content-addressable identifier. The
+# codegen names each entity by an interned payload node; because the recursion is TABLED, the same
+# payload is the SAME interned node, so it decodes to the SAME name (consistent across def and use),
+# while distinct ones get distinct names. The name is derived from a deterministic Merkle hash of
+# the node's de Bruijn structure (not Python's randomized ``hash()``), so the same structure always
+# yields the same name across processes and the output is stable under local source modifications.
+_gensym_ids: "dict[int, str]" = {}
+
+
+def _reset_gensym() -> None:
+    _gensym_ids.clear()
+    _merkle_cache.clear()
+
+
+_merkle_cache: "dict[int, int]" = {}
+
+
+def _merkle_hash(node: Node) -> int:
+    """A deterministic Merkle hash of the node's de Bruijn structure, independent of PYTHONHASHSEED."""
+    cached = _merkle_cache.get(id(node))
+    if cached is not None:
+        return cached
+    match node:
+        case Var(index=index):
+            data = struct.pack(">BQ", 0, index)
+        case Lam(body=body):
+            data = struct.pack(">BQ", 1, _merkle_hash(body))
+        case App(function=function, argument=argument):
+            data = struct.pack(">BQQ", 2, _merkle_hash(function), _merkle_hash(argument))
+        case Native():
+            data = struct.pack(">BQ", 3, id(node) & 0xFFFFFFFFFFFFFFFF)
+        case _:
+            raise ValueError(f"cannot hash {node!r}")
+    result = int.from_bytes(hashlib.sha256(data).digest()[:8], "big")
+    _merkle_cache[id(node)] = result
+    return result
+
+
+def _gensym_name(payload: "Node") -> str:
+    existing = _gensym_ids.get(id(payload))
+    if existing is not None:
+        return existing
+    name = f"vg_{_merkle_hash(payload):016x}"
+    _gensym_ids[id(payload)] = name
+    return name
+
+
+# An optional per-decode memo keyed by interned-node identity. ``decode`` is a tree recursion, so a
+# shared interned sub-graph (the interpreter hash-conses) is otherwise re-decoded once per occurrence.
+# Lambda-lifted call-by-need emits the SAME interned factory node once per occurrence of a shared
+# sub-term (COMPILE shares ~19x); decoding under ``memoized_decode`` forces and decodes each DISTINCT
+# node once (O(distinct) instead of O(occurrences)), the same node-identity memoization ``to_anf_source``
+# relies on via ``_extract``. Off by default so unrelated decodes stay simple.
+_decode_memo: "dict[int, ast.AST] | None" = None
+
+
+@contextlib.contextmanager
+def memoized_decode():
+    """Decode each distinct interned Scott node once, keyed by node identity (does not nest)."""
+    global _decode_memo
+    assert _decode_memo is None, "memoized_decode does not nest"
+    _decode_memo = {}
+    try:
+        yield
+    finally:
+        _decode_memo = None
+
+# Disjoint free-variable bands used as meta markers (far above any real index; Scott values here
+# are closed, so the only free variables in a probed term are these markers).
+_CTOR_BASE = 1_000_000
+_FIELD_BASE = 2_000_000
+_LIST_BASE = 3_000_000
+_CHURCH_SUCC = 4_000_000
+_CHURCH_ZERO = 4_000_001
+
+
+# --- encoding: a real ast node to a Scott Builder ---------------------------------------------
+
+def _ctor(tag: int, fields: "list[Builder]") -> Builder:
+    def collect(handlers: "list[Builder]") -> Builder:
+        if len(handlers) == len(SUPPORTED):
+            applied = handlers[tag]
+            for field in fields:
+                applied = app(applied, field)
+            return applied
+        return lam(lambda handler: collect(handlers + [handler]))
+
+    return collect([])
+
+
+def _scott_list(elements: "list[Builder]") -> Builder:
+    result = SCOTT_NIL
+    for element in reversed(elements):
+        result = cons(element, result)
+    return result
+
+
+def _kind(kind: int, payload: Builder) -> Builder:
+    # <church kind, payload> as a 2-tuple lambda s. s (church kind) payload.
+    return lam(lambda selector: app(app(selector, church(kind)), payload))
+
+
+def _encode_field(value: object) -> Builder:
+    if isinstance(value, ast.AST):
+        return _kind(_K_NODE, encode(value))
+    if isinstance(value, list):
+        return _kind(_K_LIST, _scott_list([_encode_field(item) for item in value]))
+    if isinstance(value, bool):
+        return _kind(_K_BOOL, church(1 if value else 0))
+    if isinstance(value, int):
+        return _kind(_K_INT, church(value))
+    if isinstance(value, str):
+        return _kind(_K_STR, _scott_list([church(ord(character)) for character in value]))
+    if value is None:
+        return _kind(_K_NONE, church(0))
+    raise ValueError(f"cannot encode field value {value!r} of type {type(value).__name__}")
+
+
+def encode(node: ast.AST) -> Builder:
+    """Encode a real Python ``ast`` node into a Scott-encoded Builder."""
+    cls = type(node)
+    if cls not in _TAG:
+        raise ValueError(f"unsupported ast node {cls.__name__}; add it to SUPPORTED")
+    fields = [_encode_field(getattr(node, name, None)) for name in cls._fields]
+    return _ctor(_TAG[cls], fields)
+
+
+# --- decoding: a Scott value (run in the interpreter) back to a real ast node ------------------
+
+def _handler(tag: int, arity: int, base: int) -> Node:
+    body: Node = make_var(base + tag + arity)  # the free marker, shifted past the arity binders
+    for position in range(arity):
+        body = make_app(body, make_var(arity - 1 - position))
+    for _ in range(arity):
+        body = make_lam(body)
+    return body
+
+
+def _spine(node: Node) -> "tuple[int, list[Node]]":
+    arguments: "list[Node]" = []
+    current = node
+    while True:
+        whnf = current.weak_head_normal_form
+        match whnf:
+            case Var(index=index):
+                arguments.reverse()
+                return index, arguments
+            case App(function=function, argument=argument):
+                arguments.append(argument)
+                current = function
+            case _:
+                raise ValueError(f"expected a variable-headed spine, got {whnf!r}")
+
+
+def _extract(node: Node, arities: "tuple[int, ...]", base: int) -> "tuple[int, list[Node]]":
+    applied = node
+    for tag, arity in enumerate(arities):
+        applied = make_app(applied, _handler(tag, arity, base))
+    head, fields = _spine(applied)
+    return head - base, fields
+
+
+def _church_to_int(node: Node) -> int:
+    current = make_app(make_app(node, make_var(_CHURCH_SUCC)), make_var(_CHURCH_ZERO))
+    count = 0
+    while True:
+        whnf = current.weak_head_normal_form
+        match whnf:
+            case Var(index=index):
+                if index != _CHURCH_ZERO:
+                    raise ValueError(f"church spine ended at variable {index}")
+                return count
+            case App(argument=argument):
+                count += 1
+                current = argument
+            case _:
+                raise ValueError(f"church spine hit {whnf!r}")
+
+
+def _decode_scott_list(node: Node) -> "list[Node]":
+    items: "list[Node]" = []
+    current = node
+    while True:
+        head, fields = _extract(current, (2, 0), _LIST_BASE)  # cons of arity 2, nil of arity 0
+        if head == 1:  # nil
+            return items
+        items.append(fields[0])  # cons head
+        current = fields[1]  # cons tail
+
+
+def _decode_field(node: Node) -> object:
+    head, fields = _extract(node, (2,), _FIELD_BASE)  # the <kind, payload> pair
+    kind = _church_to_int(fields[0])
+    payload = fields[1]
+    match kind:
+        case 0:  # NODE
+            return decode(payload)
+        case 1:  # LIST
+            return [_decode_field(item) for item in _decode_scott_list(payload)]
+        case 2:  # INT
+            return _church_to_int(payload)
+        case 3:  # STR
+            return "".join(chr(_church_to_int(code)) for code in _decode_scott_list(payload))
+        case 4:  # BOOL
+            return bool(_church_to_int(payload))
+        case 5:  # NONE
+            return None
+        case 6:  # IDENT: a list of Nats (an AST path) rendered to one underscore-joined identifier
+            return _path_to_identifier(payload)
+        case 7:  # GENSYM: a fresh identifier per distinct (interned) payload node, by node identity
+            return _gensym_name(payload)
+        case _:
+            raise ValueError(f"unknown field kind {kind}")
+
+
+# The single identifier decoder shared by every runtime: a variable's identifier is the list of Nats
+# naming its AST path, rendered ``v`` then ``_<segment>`` per segment (underscore-joined decimal
+# integers, an alphabetic prefix). Distinct paths give distinct names, so uniqueness is by construction;
+# the lambda compiler emits only the path (a list of Nats), never the rendered string.
+def _path_to_identifier(payload: Node) -> str:
+    segments = [_church_to_int(segment) for segment in _decode_scott_list(payload)]
+    return "_".join(["v", *(str(segment) for segment in segments)])
+
+
+def decode(node: Node) -> ast.AST:
+    """Decode a Scott-encoded ast value (run in the interpreter) back to a real ``ast`` node.
+
+    Under ``memoized_decode`` each distinct interned node is decoded once (the result is shared across
+    its occurrences), collapsing a shared sub-graph instead of re-walking every occurrence.
+    """
+    if _decode_memo is not None:
+        cached = _decode_memo.get(id(node))
+        if cached is not None:
+            return cached
+    tag, fields = _extract(node, _ARITY, _CTOR_BASE)
+    cls = SUPPORTED[tag]
+    decoded = cls(*[_decode_field(field) for field in fields])
+    if _decode_memo is not None:
+        _decode_memo[id(node)] = decoded
+    return decoded
+
+
+def to_python_source(node: Node) -> str:
+    """Decode a Scott-encoded Python AST and unparse it to source."""
+    return ast.unparse(ast.fix_missing_locations(decode(node)))
+
+
+def _field_payload(field_node: Node) -> Node:
+    """The payload Scott node of a ``<kind, payload>`` field (without decoding the payload)."""
+    _, parts = _extract(field_node, (2,), _FIELD_BASE)
+    return parts[1]
+
+
+def to_anf_source(node: Node, binding_name: str) -> str:
+    """Serialize a Scott-encoded Python AST to A-normal-form source, sharing sub-expressions by identity.
+
+    A generic, graph-preserving serialization (not program-specific): each distinct ``ast.Call`` node in
+    the Scott value becomes ONE assignment to a fresh temporary, so a sub-graph shared across the term
+    (the interpreter hash-conses, so a repeated sub-term is the SAME node) is emitted once rather than
+    unfolded, and a deeply nested reconstruction stays under CPython's parser nesting cap. Each distinct
+    node is forced once (``_extract`` memoised by node identity), which keeps a compiler-scale graph from
+    blowing up the interner. Non-``Call`` nodes (``Name``/``Constant``/``Lambda``/...) are inlined whole.
+
+    A compiler-emitted helper definition rides along as an ``ast.FunctionDef`` node embedded in the graph
+    (in the function position of the ``ast.Call`` that invokes it). Such a node is HOISTED: emitted once
+    (deduplicated by node identity, exactly like the ``_k`` call-sharing) into a preamble of top-level
+    ``def``s ahead of the assignments, and replaced in place by a ``Name`` referencing the def. This is the
+    general facility that lets the compiler emit its own top-level helpers (e.g. the analysis fixpoint)
+    instead of importing them from the runtime; the def is a program-independent constant, so it interns to
+    one node and is hoisted exactly once.
+
+    The result is a module: the hoisted ``def``s, then the temp assignments, ending ``<binding_name> =
+    <root>`` (a non-``Call`` root just yields that final bare-expression assignment).
+    """
+    memo: "dict[int, str]" = {}
+    hoisted: "list[ast.stmt]" = []
+    statements: "list[ast.stmt]" = []
+    counter = 0
+
+    def hoist_def(scott_node: Node) -> ast.expr:
+        """Hoist an ``ast.FunctionDef`` node to the top-level preamble (once, by node identity) and
+        reference it by its own name. Decoded whole: a helper def is small and fixed, so its body needs
+        no further A-normal-form flattening."""
+        key = id(scott_node)
+        cached = memo.get(key)
+        if cached is not None:
+            return ast.Name(id=cached, ctx=ast.Load())
+        function_def = decode(scott_node)
+        assert isinstance(function_def, ast.FunctionDef), (
+            f"hoist_def expected an ast.FunctionDef, got {type(function_def).__name__}"
+        )
+        hoisted.append(function_def)
+        memo[key] = function_def.name
+        return ast.Name(id=function_def.name, ctx=ast.Load())
+
+    def emit_callee(callee_node: Node) -> ast.expr:
+        """The function position of a Call: a hoisted ``FunctionDef`` becomes a top-level def reference;
+        anything else (a ``Name``, a ``Lambda``) is inlined whole, as before -- NOT routed through
+        ``emit``, so a curried application ``f(a)(b)`` keeps its nested-call shape rather than ANF-splitting."""
+        tag, _ = _extract(callee_node, _ARITY, _CTOR_BASE)
+        if SUPPORTED[tag] is ast.FunctionDef:
+            return hoist_def(callee_node)
+        inlined = decode(callee_node)
+        assert isinstance(inlined, ast.expr), (
+            f"emit_callee expected an ast.expr, got {type(inlined).__name__}"
+        )
+        return inlined
+
+    def emit(scott_node: Node) -> ast.expr:
+        nonlocal counter
+        key = id(scott_node)
+        if key in memo:
+            return ast.Name(id=memo[key], ctx=ast.Load())
+        tag, fields = _extract(scott_node, _ARITY, _CTOR_BASE)
+        cls = SUPPORTED[tag]
+        if cls is ast.FunctionDef:
+            return hoist_def(scott_node)  # a helper def reached in a value position: hoist it
+        if cls is not ast.Call:
+            inlined = decode(scott_node)  # a leaf or opaque expression (Name/Constant/Lambda/...): inline it
+            assert isinstance(inlined, ast.expr), (
+                f"emit expected an ast.expr, got {type(inlined).__name__}"
+            )
+            return inlined
+        function = emit_callee(_field_payload(fields[0]))
+        arguments = [emit(_field_payload(argument)) for argument in _decode_scott_list(_field_payload(fields[1]))]
+        call = ast.Call(func=function, args=arguments, keywords=[])
+        name = f"_k{counter}"
+        counter += 1
+        statements.append(ast.Assign(targets=[ast.Name(id=name, ctx=ast.Store())], value=call))
+        memo[key] = name
+        return ast.Name(id=name, ctx=ast.Load())
+
+    root = emit(node)
+    statements.append(ast.Assign(targets=[ast.Name(id=binding_name, ctx=ast.Store())], value=root))
+    return ast.unparse(ast.fix_missing_locations(ast.Module(body=hoisted + statements, type_ignores=[])))
+
+
+def roundtrip(source: str, *, mode: str = "eval") -> str:
+    """Parse ``source``, encode it to a Scott value, decode it, and unparse: a faithfulness check."""
+    return to_python_source(build(encode(ast.parse(source, mode=mode))))
+
+
+# --- BinNat readouts (the decode side of the _binnat encoding) -------------------------------------
+
+# A free-variable band used as a meta marker when probing a Scott boolean (a bit). Disjoint from the
+# bands above, so a probed bit's only free variables are these markers.
+_BIT_BASE = 8_500_000
+
+
+def _bit_value(node: "Node") -> int:
+    # A Scott boolean applied to two nullary handlers exposes handler 0 for TRUE, handler 1 for FALSE.
+    tag, _ = _extract(node, (0, 0), _BIT_BASE)
+    return 1 if tag == 0 else 0
+
+
+def binnat_to_int(node: "Node") -> int:
+    """Decode a BinNat (an LSB-first Scott list of bits) to a non-negative int."""
+    value = 0
+    for position, bit in enumerate(_decode_scott_list(node)):
+        value += _bit_value(bit) << position
+    return value
+
+
+def binnat_list_to_identifier(node: "Node", prefix: str = "v") -> str:
+    """Decode a Scott list of BinNats to an underscore-joined identifier, e.g. ``v_12_3_567``."""
+    segments = [binnat_to_int(segment) for segment in _decode_scott_list(node)]
+    return "_".join([prefix, *(str(segment) for segment in segments)])