guest/plm.py

"""Small whole-program pass. Unsupported programs keep their original AST.

Only adjacent tool assignments are reordered: prepare after argument definitions,
resolve at the original statement. Other statements and control regions are barriers.
"""
import ast
import copy


def transform(source):
    tree = ast.parse(source, filename="<spine>")
    names = {n.id for n in ast.walk(tree) if isinstance(n, ast.Name)}
    names |= {n.name for n in ast.walk(tree) if isinstance(n, ast.ExceptHandler) and n.name}
    data_names = {"inputs"} | {n.id for n in ast.walk(tree) if isinstance(n, ast.Name) and isinstance(n.ctx, ast.Store)}

    def data(node):
        if isinstance(node, ast.Constant):
            return isinstance(node.value, (str, int, float, bool, type(None)))
        if isinstance(node, ast.Name):
            return node.id in data_names
        if isinstance(node, (ast.List, ast.Tuple)):
            return all(data(n) for n in node.elts)
        if isinstance(node, ast.Dict):
            return all(k is not None and data(k) and data(v) for k, v in zip(node.keys, node.values))
        if isinstance(node, ast.Subscript):
            return data(node.value) and data(node.slice)
        if isinstance(node, ast.BinOp):
            return data(node.left) and data(node.right)
        if isinstance(node, ast.UnaryOp):
            return data(node.operand)
        if isinstance(node, ast.BoolOp):
            return all(data(n) for n in node.values)
        if isinstance(node, ast.Compare):
            return data(node.left) and all(data(n) for n in node.comparators)
        return False

    def argument(node):
        if isinstance(node, (ast.Constant, ast.Name)):
            return data(node)
        return (isinstance(node, ast.Subscript) and argument(node.value)
                and isinstance(node.slice, ast.Constant) and isinstance(node.slice.value, (str, int)))

    def tool_call(node):
        return (isinstance(node, ast.Call) and isinstance(node.func, ast.Name)
                and node.func.id == "tool" and len(node.args) == 1
                and isinstance(node.args[0], ast.Constant) and isinstance(node.args[0].value, str)
                and all(k.arg is not None and argument(k.value) for k in node.keywords))

    def block_ok(body):
        for s in body:
            if isinstance(s, ast.Assign):
                if (len(s.targets) != 1 or not isinstance(s.targets[0], ast.Name)
                        or s.targets[0].id in ("inputs", "tool")
                        or not (tool_call(s.value) or data(s.value))):
                    return False
            elif isinstance(s, ast.If):
                if not (data(s.test) and block_ok(s.body) and block_ok(s.orelse)):
                    return False
            elif isinstance(s, ast.Try):
                if not all(block_ok(b) for b in (s.body, s.orelse, s.finalbody)):
                    return False
                for h in s.handlers:
                    if h.name in ("inputs", "tool") or not block_ok(h.body):
                        return False
                    if h.type is not None and not isinstance(h.type, ast.Name):
                        return False
            elif not (isinstance(s, ast.Pass) or isinstance(s, ast.Expr) and isinstance(s.value, ast.Constant)):
                return False
        return True

    # No arbitrary calls, alias mutation, loops or user objects in an optimized program.
    if not block_ok(tree.body):
        return tree, None

    def fresh(base):
        while base in names:
            base += "_"
        names.add(base)
        return base

    prepare, resolve = fresh("_spine_prepare"), fresh("_spine_resolve")
    changed = False

    def rewrite(body):
        nonlocal changed
        result, i = [], 0
        while i < len(body):
            s = body[i]
            if isinstance(s, ast.Assign) and tool_call(s.value):
                group = []
                while i < len(body) and isinstance(body[i], ast.Assign) and tool_call(body[i].value):
                    group.append(body[i])
                    i += 1
                before = [[] for _ in group]
                definitions = {}
                replacements = []
                for index, statement in enumerate(group):
                    call = statement.value
                    loaded = {n.id for k in call.keywords for n in ast.walk(k.value) if isinstance(n, ast.Name)}
                    start = 1 + max((definitions.get(n, -1) for n in loaded), default=-1)
                    slot = fresh("_spine_future")
                    args = ast.Dict(keys=[ast.Constant(k.arg) for k in call.keywords], values=copy.deepcopy([k.value for k in call.keywords]))
                    thunk = ast.Lambda(args=ast.arguments(posonlyargs=[], args=[], kwonlyargs=[], kw_defaults=[], defaults=[]),
                                       body=ast.Tuple(elts=[copy.deepcopy(call.args[0]), args], ctx=ast.Load()))
                    early = ast.Assign(targets=[ast.Name(slot, ast.Store())],
                                       value=ast.Call(ast.Name(prepare, ast.Load()), [thunk], []))
                    before[start].append(ast.copy_location(early, statement))
                    statement.value = ast.copy_location(ast.Call(ast.Name(resolve, ast.Load()),
                        [ast.Name(slot, ast.Load()), *call.args], call.keywords), call)
                    replacements.append(statement)
                    definitions[statement.targets[0].id] = index
                for early, statement in zip(before, replacements):
                    result.extend(early)
                    result.append(statement)
                changed = True
                continue
            if isinstance(s, ast.If):
                s.body, s.orelse = rewrite(s.body), rewrite(s.orelse)
            elif isinstance(s, ast.Try):
                s.body, s.orelse, s.finalbody = rewrite(s.body), rewrite(s.orelse), rewrite(s.finalbody)
                for h in s.handlers:
                    h.body = rewrite(h.body)
            result.append(s)
            i += 1
        return result

    tree.body = rewrite(tree.body)
    return ast.fix_missing_locations(tree), (prepare, resolve) if changed else None