pythonfaster 1.8.0__py3-none-any.whl

pythonfaster/transform/recpass.py

@@ -0,0 +1,522 @@
+"""Pass 9 + Pass 10: recursive-function cdef lowering and pure-function
+sum-LICM folding.
+
+Pass 9 (cdef-rec): for module-level functions that (a) call themselves
+recursively (directly or mutually), (b) are pure (no I/O, no
+global/nonlocal writes, no dynamic calls), and (c) have solved argument
+types from the type solver, emit a ``cdef`` C-function twin with typed
+parameters + a thin ``def`` wrapper that keeps the original name and
+signature visible to Python code.  Recursive calls then go through the
+C stack instead of the Python call protocol.
+
+Pass 9 also lowers "pure leaf helpers": module functions that are not
+themselves recursive but whose every call target is either a lowered
+candidate or a default-parameter name bound (via ``def f(x, g=cand)``)
+to a lowered candidate.  For helpers, defaulted params are dropped from
+the cdef twin and their call sites in the body are rewritten to the
+candidate's ``__fast__`` twin — safe because default values were bound
+at definition time and the ``def`` wrapper retains them for callers
+that pass overrides.
+
+Pass 10 (sum-licm): fold ``sum(F(args) for k in range(lo, hi))`` into
+``(hi - lo) * __fast__F(args[k:=1])`` when:
+
+- ``F`` is a lowered function (candidate or helper),
+- ``F`` provably returns int (float folding is forbidden: IEEE 754
+  rounding of a*b differs from repeated addition; integer math is exact
+  under CPython's bignum arithmetic, so the fold is an identity), and
+- ``F``'s result is provably independent of the comprehension loop
+  variable ``k`` — checked by name-expansion through F's parameter
+  bindings and local assignments, not by merely grepping the argument
+  list (binary_trees' ``make_check((k, d))`` unpacks ``k`` into a name
+  that is never used, so the fold is still valid).
+"""
+
+from __future__ import annotations
+
+import ast
+import copy
+import re
+
+
+def _is_pure_stmt(stmt: ast.stmt, self_names: set[str]) -> bool:
+    """Statement-level purity: assignments, if/return only."""
+    if isinstance(stmt, ast.Assign):
+        return True
+    if isinstance(stmt, ast.AugAssign):
+        t = stmt.target
+        return isinstance(t, ast.Name) and t.id in self_names
+    if isinstance(stmt, ast.Return):
+        return True
+    if isinstance(stmt, ast.If):
+        return (all(_is_pure_stmt(s, self_names) for s in stmt.body)
+                and all(_is_pure_stmt(s, self_names) for s in stmt.orelse))
+    return False
+
+
+def _is_pure_expr(node: ast.AST, allowed_calls: set[str]) -> bool:
+    """Expression purity: names, constants, arithmetic, comparison,
+    boolean ops, tuple/list building, and calls ONLY to allowed
+    (known-pure) function names.  Operator nodes (ast.Add etc.) are not
+    expressions and are never recursed into."""
+    if isinstance(node, ast.Name):
+        return True
+    if isinstance(node, ast.Constant):
+        return True
+    if isinstance(node, ast.Tuple):
+        return all(_is_pure_expr(e, allowed_calls) for e in node.elts)
+    if isinstance(node, ast.List):
+        return all(_is_pure_expr(e, allowed_calls) for e in node.elts)
+    if isinstance(node, ast.BinOp):
+        return (_is_pure_expr(node.left, allowed_calls)
+                and _is_pure_expr(node.right, allowed_calls))
+    if isinstance(node, ast.UnaryOp):
+        return _is_pure_expr(node.operand, allowed_calls)
+    if isinstance(node, ast.BoolOp):
+        return all(_is_pure_expr(v, allowed_calls) for v in node.values)
+    if isinstance(node, ast.Compare):
+        return (_is_pure_expr(node.left, allowed_calls)
+                and all(_is_pure_expr(c, allowed_calls)
+                        for c in node.comparators))
+    if isinstance(node, ast.IfExp):
+        return (_is_pure_expr(node.test, allowed_calls)
+                and _is_pure_expr(node.body, allowed_calls)
+                and _is_pure_expr(node.orelse, allowed_calls))
+    if isinstance(node, ast.Call):
+        # name-directed call only (no method calls: x.f() unknown)
+        if isinstance(node.func, ast.Name):
+            if node.func.id in allowed_calls and not node.keywords:
+                return all(_is_pure_expr(a, allowed_calls)
+                           for a in node.args)
+        return False
+    return False
+
+
+def _module_funcs(tree: ast.Module) -> dict[str, ast.FunctionDef]:
+    return {n.name: n for n in tree.body if isinstance(n, ast.FunctionDef)}
+
+
+def _recursive_pure_candidates(tree: ast.Module) -> list[str]:
+    """Names of module-level functions that call themselves (directly or
+    through the mutual-call graph) and whose bodies are pure statements
+    over allowed calls = {self} ∪ {other pure functions}."""
+    funcs = _module_funcs(tree)
+    calls: dict[str, set[str]] = {}
+    for name, fn in funcs.items():
+        called = set()
+        for sub in ast.walk(fn):
+            if isinstance(sub, ast.Call) and isinstance(sub.func, ast.Name):
+                if sub.func.id in funcs:
+                    called.add(sub.func.id)
+        calls[name] = called
+
+    pure: set[str] = set()
+    changed = True
+    while changed:
+        changed = False
+        for name, fn in funcs.items():
+            if name in pure:
+                continue
+            allowed = pure | {name}
+            if (all(_is_pure_stmt(s, {a.arg for a in fn.args.args})
+                    for s in fn.body)
+                    and all(_is_pure_expr(e, allowed)
+                            for s in fn.body if isinstance(s, ast.Return)
+                            for e in [s.value] if e is not None)
+                    and calls[name] <= allowed):
+                pure.add(name)
+                changed = True
+
+    def _is_recursive(name: str) -> bool:
+        seen, stack = set(), [name]
+        while stack:
+            cur = stack.pop()
+            for nxt in calls.get(cur, ()):
+                if nxt == name:
+                    return True
+                if nxt not in seen:
+                    seen.add(nxt)
+                    stack.append(nxt)
+        return False
+
+    return [n for n in sorted(pure) if _is_recursive(n)]
+
+
+def _default_bound_params(fn: ast.FunctionDef) -> dict[str, str]:
+    """{param_name: candidate_name} for trailing params whose default is
+    a bare Name referencing a module-level recursive-pure candidate."""
+    args = fn.args.args
+    defaults = fn.args.defaults
+    bound = {}
+    for param, dflt in zip(args[len(args) - len(defaults):], defaults):
+        if isinstance(dflt, ast.Name):
+            bound[param.arg] = dflt.id
+    return bound
+
+
+def _pure_leaf_helpers(tree: ast.Module,
+                       cands: list[str]) -> dict[str, dict[str, str]]:
+    """Module-level non-recursive pure functions whose every in-body
+    call target is either a candidate or a default-param name bound to
+    a candidate.  Returns {func_name: {param: candidate}}."""
+    funcs = _module_funcs(tree)
+    cand_set = set(cands)
+    helpers: dict[str, dict[str, str]] = {}
+    for name, fn in funcs.items():
+        if name in cand_set:
+            continue
+        bound = _default_bound_params(fn)
+        if not all(v in cand_set for v in bound.values()):
+            continue
+        allowed = cand_set | set(bound.keys()) | {name}
+        # every call in the BODY must target an allowed name
+        # (decorators like @cython.boundscheck(False) are ignored)
+        ok = True
+        for stmt in fn.body:
+            for sub in ast.walk(stmt):
+                if isinstance(sub, ast.Call):
+                    if not (isinstance(sub.func, ast.Name)
+                            and sub.func.id in allowed):
+                        ok = False
+                        break
+            if not ok:
+                break
+        if not ok:
+            continue
+        if not (all(_is_pure_stmt(s, {a.arg for a in fn.args.args})
+                    for s in fn.body)
+                and all(_is_pure_expr(e, allowed)
+                        for s in fn.body if isinstance(s, ast.Return)
+                        for e in [s.value] if e is not None)):
+            continue
+        helpers[name] = bound
+    return helpers
+
+
+class _CdefRecLower(ast.NodeTransformer):
+    """Rewrites eligible recursive functions and pure leaf helpers into
+    cdef twins + thin def wrappers."""
+
+    def __init__(self, candidates: list[str], helpers: dict[str, dict],
+                 arg_types: dict, int_returns: set[str]):
+        self.cands = set(candidates)
+        self.helpers = helpers        # {name: {param: cand}}
+        self.arg_types = arg_types
+        self.int_returns = int_returns
+        self.added: list[str] = []
+
+    def visit_FunctionDef(self, node: ast.FunctionDef) -> ast.AST:
+        name = node.name
+        is_cand = name in self.cands
+        is_helper = name in self.helpers
+        if not (is_cand or is_helper):
+            return node
+        fast = f"__fast__{name}"
+        ats = self.arg_types.get(name, {})
+
+        if is_cand:
+            params = ", ".join(
+                f"{ats.get(a.arg, 'object')} {a.arg}"
+                for a in node.args.args)
+        else:
+            # helper: only non-defaulted params go into the cdef twin
+            bound = self.helpers[name]
+            ndef = len(node.args.defaults)
+            live = [a for a in node.args.args if a.arg not in bound]
+            if len(live) + ndef != len(node.args.args):
+                return node  # unexpected shape — skip
+            params = ", ".join(
+                f"{ats.get(a.arg, 'object')} {a.arg}" for a in live)
+
+        ret = "long long" if name in self.int_returns else "object"
+        except_clause = " except? -1" if ret == "long long" else ""
+
+        # textual body with call-site rewriting
+        body_src = []
+        for stmt in node.body:
+            src = ast.unparse(stmt)
+            for cand in sorted(self.cands):
+                src = re.sub(rf"\b{cand}\(", f"__fast__{cand}(", src)
+            if is_helper:
+                for param, cand in self.helpers[name].items():
+                    src = re.sub(rf"\b{param}\(", f"__fast__{cand}(", src)
+            body_src.append("    " + src.replace("\n", "\n    "))
+        header = f"cdef {ret} {fast}({params}){except_clause}:"
+        self.added.append(header + "\n" + "\n".join(body_src))
+
+        # thin wrapper keeps the full original signature
+        args = [a.arg for a in node.args.args]
+        wrapper_src = f"def {name}({ast.unparse(node.args)}):\n"
+        if is_cand:
+            call_args = args
+        else:
+            bound = self.helpers[name]
+            call_args = [a for a in args if a not in bound]
+        wrapper_src += f"    return {fast}({', '.join(call_args)})"
+        wrapper = ast.parse(wrapper_src).body[0]
+        ast.copy_location(wrapper, node)
+        return wrapper
+
+
+def _result_depends_on_loopvar(fn: ast.FunctionDef, args: list[ast.expr],
+                               loopvar: str) -> bool:
+    """True if fn(…args…) may return a different value when loopvar
+    changes.  Tracks parameter bindings and local assignments by name
+    expansion; conservative (returns True) on any complex shape."""
+    params = [a.arg for a in fn.args.args]
+    ndef = len(fn.args.defaults)
+    if len(args) != len(params) - ndef:
+        return True
+    env: dict[str, ast.expr] = dict(zip(params, args))
+
+    locs: dict[str, ast.expr] = {}
+    for stmt in fn.body:
+        if not isinstance(stmt, ast.Assign) or len(stmt.targets) != 1:
+            continue
+        target, value = stmt.targets[0], stmt.value
+        # tuple unpack from a bound parameter:  i, d = itde
+        if (isinstance(target, ast.Tuple) and isinstance(value, ast.Name)
+                and value.id in env and isinstance(env[value.id], ast.Tuple)
+                and len(target.elts) == len(env[value.id].elts)):
+            for t, e in zip(target.elts, env[value.id].elts):
+                if not isinstance(t, ast.Name):
+                    return True
+                locs[t.id] = e
+            continue
+        if isinstance(target, ast.Name):
+            locs[target.id] = value
+        elif isinstance(target, ast.Tuple):
+            if not isinstance(value, ast.Tuple) or len(target.elts) != len(value.elts):
+                return True
+            for t, e in zip(target.elts, value.elts):
+                if not isinstance(t, ast.Name):
+                    return True
+                locs[t.id] = e
+        else:
+            return True
+
+    def mentions_k(node: ast.expr, seen: frozenset) -> bool:
+        for sub in ast.walk(node):
+            if isinstance(sub, ast.Name):
+                nm = sub.id
+                if nm == loopvar:
+                    return True
+                if nm in seen:
+                    continue
+                if nm in locs and mentions_k(
+                        locs[nm], seen | {nm}):
+                    return True
+                if nm in env and any(
+                        isinstance(x, ast.Name) and x.id == loopvar
+                        for x in ast.walk(env[nm])):
+                    return True
+                # env expansion deeper than one level: expand once more
+                if nm in env:
+                    for x in ast.walk(env[nm]):
+                        if isinstance(x, ast.Name) and x.id in locs \
+                                and x.id not in seen \
+                                and mentions_k(locs[x.id], seen | {x.id}):
+                            return True
+        return False
+
+    for stmt in fn.body:
+        if isinstance(stmt, ast.Return) and stmt.value is not None:
+            if mentions_k(stmt.value, frozenset()):
+                return True
+    return False
+
+
+def _replace_loopvar_with_const(node: ast.AST, loopvar: str) -> ast.AST:
+    """Replace Name(loopvar) occurrences with Constant(1) — the value is
+    provably unused, any constant works."""
+    class _R(ast.NodeTransformer):
+        def visit_Name(self, n: ast.Name) -> ast.AST:
+            if n.id == loopvar:
+                return ast.copy_location(ast.Constant(value=1), n)
+            return n
+    return _R().visit(node)
+
+
+def _candidate_int_returns(funcs: dict[str, ast.FunctionDef],
+                          cands: list[str]) -> set[str]:
+    """Fixed-point int-return proof over the recursive-pure candidates:
+    assume self/mutual candidate calls return int, prove arithmetic
+    int-ness of every return expr.  check_tree's ``1 + check_tree(l) +
+    check_tree(r)`` becomes provable once the assumption is in place."""
+    cand_set = set(cands)
+    proven: set[str] = set()
+    changed = True
+    while changed:
+        changed = False
+        for name in cands:
+            if name in proven:
+                continue
+            fn = funcs[name]
+
+            def inty(e: ast.expr) -> bool:
+                if isinstance(e, ast.Constant) and isinstance(e.value, int):
+                    return True
+                if isinstance(e, ast.Name):
+                    return False
+                if isinstance(e, ast.BinOp):
+                    return inty(e.left) and inty(e.right)
+                if isinstance(e, ast.UnaryOp):
+                    return (isinstance(e.op, (ast.UAdd, ast.USub, ast.Invert))
+                            and inty(e.operand))
+                if isinstance(e, ast.Call):
+                    return (isinstance(e.func, ast.Name)
+                            and (e.func.id in proven
+                                 or e.func.id == name))
+                if isinstance(e, ast.IfExp):
+                    return inty(e.body) and inty(e.orelse)
+                return False
+
+            rets = [s.value for s in fn.body
+                    if isinstance(s, ast.Return) and s.value is not None]
+            if not rets:
+                continue
+            if all(inty(r) for r in rets):
+                proven.add(name)
+                changed = True
+    return proven
+
+
+def _helper_int_return(fn: ast.FunctionDef, bound: dict[str, str],
+                      int_returns: set[str]) -> bool:
+    """A helper provably returns int when every return expr is a call
+    to an int-returning candidate (possibly via cand-bound param names)
+    or int arithmetic over such calls."""
+    names_ok = set(int_returns) | set(bound.keys())
+
+    def inty(e: ast.expr) -> bool:
+        if isinstance(e, ast.Constant) and isinstance(e.value, int):
+            return True
+        if isinstance(e, ast.Call) and isinstance(e.func, ast.Name):
+            return e.func.id in names_ok
+        if isinstance(e, ast.BinOp):
+            return inty(e.left) and inty(e.right)
+        if isinstance(e, ast.UnaryOp):
+            return inty(e.operand)
+        return False
+
+    rets = [s.value for s in fn.body
+            if isinstance(s, ast.Return) and s.value is not None]
+    return bool(rets) and all(inty(r) for r in rets)
+
+
+def _sum_licm_fold(tree: ast.Module, lowerable: dict[str, ast.FunctionDef],
+                   int_returns: set[str]) -> list[str]:
+    """Fold sum(F(...) for k in range(lo, hi)) → (hi-lo) * __fast__F(...)
+    when F is lowered, int-returning, and loop-invariant in result."""
+    folded = []
+
+    class _Fold(ast.NodeTransformer):
+        def visit_Assign(self, node: ast.Assign):
+            self.generic_visit(node)
+            if not (len(node.targets) == 1
+                    and isinstance(node.targets[0], ast.Name)):
+                return node
+            val = node.value
+            if not (isinstance(val, ast.Call)
+                    and isinstance(val.func, ast.Name)
+                    and val.func.id == "sum"
+                    and len(val.args) == 1
+                    and isinstance(val.args[0], ast.GeneratorExp)):
+                return node
+            gen = val.args[0]
+            if len(gen.generators) != 1:
+                return node
+            g = gen.generators[0]
+            if g.ifs or g.is_async:
+                return node
+            tgt = g.target
+            if not isinstance(tgt, ast.Name):
+                return node
+            rng = g.iter
+            if not (isinstance(rng, ast.Call)
+                    and isinstance(rng.func, ast.Name)
+                    and rng.func.id == "range" and len(rng.args) == 2
+                    and not rng.keywords):
+                return node
+            lo, hi = rng.args
+            elt = gen.elt
+            if not (isinstance(elt, ast.Call)
+                    and isinstance(elt.func, ast.Name)
+                    and elt.func.id in lowerable
+                    and elt.func.id in int_returns
+                    and not elt.keywords):
+                return node
+            fname = elt.func.id
+            fn = lowerable[fname]
+            if _result_depends_on_loopvar(fn, elt.args, tgt.id):
+                return node
+            folded.append(fname)
+            count = ast.BinOp(left=copy.deepcopy(hi), op=ast.Sub(),
+                              right=copy.deepcopy(lo))
+            once = copy.deepcopy(elt)
+            once.func = ast.copy_location(
+                ast.Name(id=f"__fast__{fname}", ctx=ast.Load()),
+                elt.func)
+            once.args = [_replace_loopvar_with_const(a, tgt.id)
+                         for a in once.args]
+            ast.fix_missing_locations(once)
+            node.value = ast.BinOp(left=count, op=ast.Mult(), right=once)
+            return node
+
+    _Fold().visit(tree)
+    return folded
+
+
+# cdef blocks are not valid Python syntax, so they cannot be inserted as
+# AST nodes.  Pass 9 plants ``__pf_cdef_N = 0`` marker assignments in the
+# tree and registers the real cdef source here; the engine swaps them
+# textually after ast.unparse (same mechanism as Pass 8's markers).
+PENDING_CDEFS: dict[str, str] = {}
+_marker_ctr = [0]
+
+
+def apply_recursive_passes(tree: ast.Module, arg_types: dict,
+                           int_returns: set[str]) -> list[str]:
+    """Entry: run Pass 9 then Pass 10 on the module.  Returns the list
+    of applied strategy tags."""
+    PENDING_CDEFS.clear()
+    applied = []
+    cands = _recursive_pure_candidates(tree)
+    helpers = _pure_leaf_helpers(tree, cands) if cands else {}
+    # snapshot ORIGINAL function defs (pre-lowering) for Pass 10's
+    # dependency analysis — wrappers are useless for that.
+    orig_funcs = _module_funcs(tree)
+    # recursion-aware int-return proof (solver cannot see through
+    # self-calls): fixed point over the candidates, then helpers.
+    int_returns = (set(int_returns)
+                   | _candidate_int_returns(orig_funcs, cands))
+    # helpers inherit int-return proof transitively from candidates
+    for hname, bound in helpers.items():
+        if _helper_int_return(orig_funcs[hname], bound,
+                              int_returns):
+            int_returns = int_returns | {hname}
+    if cands or helpers:
+        lower = _CdefRecLower(cands, helpers, arg_types, int_returns)
+        tree.body = [lower.visit(n) for n in tree.body]
+        insert_at = next((i for i, n in enumerate(tree.body)
+                          if isinstance(n, ast.FunctionDef)), 0)
+        for cdef_src in lower.added:
+            marker = f"__pf_cdef_{_marker_ctr[0]}"
+            _marker_ctr[0] += 1
+            PENDING_CDEFS[marker] = cdef_src
+            placeholder = ast.Assign(
+                targets=[ast.Name(id=marker, ctx=ast.Store())],
+                value=ast.Constant(value=0))
+            tree.body.insert(insert_at, placeholder)
+            insert_at += 1
+        ast.fix_missing_locations(tree)
+        applied.append("cdef-rec")
+
+    lowerable = {n: f for n, f in orig_funcs.items()
+                 if n in set(cands) | set(helpers)}
+    if lowerable:
+        folded = _sum_licm_fold(tree, lowerable, int_returns)
+        if folded:
+            ast.fix_missing_locations(tree)
+            applied.append("sum-licm")
+    return applied

pythonfaster-1.8.0.data/data/pythonfaster.pth

@@ -0,0 +1 @@
+import pythonfaster._bootstrap