sonolus.py 0.1.2__py3-none-any.whl → 0.1.4__py3-none-any.whl

sonolus/backend/excepthook.py

@@ -1,37 +1,37 @@
-"""Exception hook to filter out compiler internal frames from tracebacks."""
-
-import sys
-from types import TracebackType
-
-
-def should_filter_traceback(tb: TracebackType | None) -> bool:
-    return tb is not None and (
-        tb.tb_frame.f_globals.get("_filter_traceback_", False) or should_filter_traceback(tb.tb_next)
-    )
-
-
-def is_compiler_internal(tb: TracebackType):
-    return tb.tb_frame.f_locals.get("_compiler_internal_", False) or tb.tb_frame.f_globals.get(
-        "_compiler_internal_", False
-    )
-
-
-def filter_traceback(tb: TracebackType | None) -> TracebackType | None:
-    if tb is None:
-        return None
-    if is_compiler_internal(tb):
-        return filter_traceback(tb.tb_next)
-    tb.tb_next = filter_traceback(tb.tb_next)
-    return tb
-
-
-def excepthook(exc, value, tb):
-    import traceback
-
-    if should_filter_traceback(tb):
-        tb = filter_traceback(tb)
-    traceback.print_exception(exc, value, tb)
-
-
-def install_excepthook():
-    sys.excepthook = excepthook
+"""Exception hook to filter out compiler internal frames from tracebacks."""
+
+import sys
+from types import TracebackType
+
+
+def should_filter_traceback(tb: TracebackType | None) -> bool:
+    return tb is not None and (
+        tb.tb_frame.f_globals.get("_filter_traceback_", False) or should_filter_traceback(tb.tb_next)
+    )
+
+
+def is_compiler_internal(tb: TracebackType):
+    return tb.tb_frame.f_locals.get("_compiler_internal_", False) or tb.tb_frame.f_globals.get(
+        "_compiler_internal_", False
+    )
+
+
+def filter_traceback(tb: TracebackType | None) -> TracebackType | None:
+    if tb is None:
+        return None
+    if is_compiler_internal(tb):
+        return filter_traceback(tb.tb_next)
+    tb.tb_next = filter_traceback(tb.tb_next)
+    return tb
+
+
+def excepthook(exc, value, tb):
+    import traceback
+
+    if should_filter_traceback(tb):
+        tb = filter_traceback(tb)
+    traceback.print_exception(exc, value, tb)
+
+
+def install_excepthook():
+    sys.excepthook = excepthook

sonolus/backend/finalize.py

@@ -1,69 +1,69 @@
-from sonolus.backend.flow import BasicBlock, traverse_cfg_preorder
-from sonolus.backend.ir import IRConst, IRGet, IRInstr, IRPureInstr, IRSet
-from sonolus.backend.node import ConstantNode, EngineNode, FunctionNode
-from sonolus.backend.ops import Op
-from sonolus.backend.place import BlockPlace
-
-
-def cfg_to_engine_node(entry: BasicBlock):
-    block_indexes = {block: i for i, block in enumerate(traverse_cfg_preorder(entry))}
-    block_statements = []
-    for block in block_indexes:
-        statements = []
-        statements.extend(ir_to_engine_node(stmt) for stmt in block.statements)
-        outgoing = {edge.cond: edge.dst for edge in block.outgoing}
-        match outgoing:
-            case {**other} if not other:
-                statements.append(ConstantNode(value=len(block_indexes)))
-            case {None: target, **other} if not other:
-                statements.append(ConstantNode(value=block_indexes[target]))
-            case {0: f_branch, None: t_branch, **other} if not other:
-                statements.append(
-                    FunctionNode(
-                        func=Op.If,
-                        args=[
-                            ir_to_engine_node(block.test),
-                            ConstantNode(value=block_indexes[t_branch]),
-                            ConstantNode(value=block_indexes[f_branch]),
-                        ],
-                    )
-                )
-            case dict() as tgt:
-                args = [ir_to_engine_node(block.test)]
-                default = len(block_indexes)
-                for cond, target in tgt.items():
-                    if cond is None:
-                        default = block_indexes[target]
-                    args.append(ConstantNode(value=cond))
-                    args.append(ConstantNode(value=block_indexes[target]))
-                args.append(ConstantNode(value=default))
-                statements.append(FunctionNode(Op.SwitchWithDefault, args))
-        block_statements.append(FunctionNode(Op.Execute, statements))
-    block_statements.append(ConstantNode(value=0))
-    return FunctionNode(Op.Block, [FunctionNode(Op.JumpLoop, block_statements)])
-
-
-def ir_to_engine_node(stmt) -> EngineNode:
-    match stmt:
-        case int() | float():
-            return ConstantNode(value=float(stmt))
-        case IRConst(value=value):
-            return ConstantNode(value=value)
-        case IRPureInstr(op=op, args=args) | IRInstr(op=op, args=args):
-            return FunctionNode(func=op, args=[ir_to_engine_node(arg) for arg in args])
-        case IRGet(place=place):
-            return ir_to_engine_node(place)
-        case BlockPlace() as place:
-            if place.offset == 0:
-                index = ir_to_engine_node(place.index)
-            elif place.index == 0:
-                index = ConstantNode(value=place.offset)
-            else:
-                index = FunctionNode(
-                    func=Op.Add, args=[ir_to_engine_node(place.index), ConstantNode(value=place.offset)]
-                )
-            return FunctionNode(func=Op.Get, args=[ir_to_engine_node(place.block), index])
-        case IRSet(place=place, value=value):
-            return FunctionNode(func=Op.Set, args=[*ir_to_engine_node(place).args, ir_to_engine_node(value)])
-        case _:
-            raise TypeError(f"Unsupported IR statement: {stmt}")
+from sonolus.backend.flow import BasicBlock, traverse_cfg_preorder
+from sonolus.backend.ir import IRConst, IRGet, IRInstr, IRPureInstr, IRSet
+from sonolus.backend.node import ConstantNode, EngineNode, FunctionNode
+from sonolus.backend.ops import Op
+from sonolus.backend.place import BlockPlace
+
+
+def cfg_to_engine_node(entry: BasicBlock):
+    block_indexes = {block: i for i, block in enumerate(traverse_cfg_preorder(entry))}
+    block_statements = []
+    for block in block_indexes:
+        statements = []
+        statements.extend(ir_to_engine_node(stmt) for stmt in block.statements)
+        outgoing = {edge.cond: edge.dst for edge in block.outgoing}
+        match outgoing:
+            case {**other} if not other:
+                statements.append(ConstantNode(value=len(block_indexes)))
+            case {None: target, **other} if not other:
+                statements.append(ConstantNode(value=block_indexes[target]))
+            case {0: f_branch, None: t_branch, **other} if not other:
+                statements.append(
+                    FunctionNode(
+                        func=Op.If,
+                        args=[
+                            ir_to_engine_node(block.test),
+                            ConstantNode(value=block_indexes[t_branch]),
+                            ConstantNode(value=block_indexes[f_branch]),
+                        ],
+                    )
+                )
+            case dict() as tgt:
+                args = [ir_to_engine_node(block.test)]
+                default = len(block_indexes)
+                for cond, target in tgt.items():
+                    if cond is None:
+                        default = block_indexes[target]
+                    args.append(ConstantNode(value=cond))
+                    args.append(ConstantNode(value=block_indexes[target]))
+                args.append(ConstantNode(value=default))
+                statements.append(FunctionNode(Op.SwitchWithDefault, args))
+        block_statements.append(FunctionNode(Op.Execute, statements))
+    block_statements.append(ConstantNode(value=0))
+    return FunctionNode(Op.Block, [FunctionNode(Op.JumpLoop, block_statements)])
+
+
+def ir_to_engine_node(stmt) -> EngineNode:
+    match stmt:
+        case int() | float():
+            return ConstantNode(value=float(stmt))
+        case IRConst(value=value):
+            return ConstantNode(value=value)
+        case IRPureInstr(op=op, args=args) | IRInstr(op=op, args=args):
+            return FunctionNode(func=op, args=[ir_to_engine_node(arg) for arg in args])
+        case IRGet(place=place):
+            return ir_to_engine_node(place)
+        case BlockPlace() as place:
+            if place.offset == 0:
+                index = ir_to_engine_node(place.index)
+            elif place.index == 0:
+                index = ConstantNode(value=place.offset)
+            else:
+                index = FunctionNode(
+                    func=Op.Add, args=[ir_to_engine_node(place.index), ConstantNode(value=place.offset)]
+                )
+            return FunctionNode(func=Op.Get, args=[ir_to_engine_node(place.block), index])
+        case IRSet(place=place, value=value):
+            return FunctionNode(func=Op.Set, args=[*ir_to_engine_node(place).args, ir_to_engine_node(value)])
+        case _:
+            raise TypeError(f"Unsupported IR statement: {stmt}")

sonolus/backend/flow.py

@@ -1,92 +1,121 @@
-import textwrap
-from collections import deque
-from collections.abc import Iterator
-
-from sonolus.backend.ir import IRConst, IRExpr, IRStmt
-
-
-class FlowEdge:
-    src: "BasicBlock"
-    dst: "BasicBlock"
-    cond: IRExpr | None
-
-    def __init__(self, src: "BasicBlock", dst: "BasicBlock", cond: IRExpr | None = None):
-        self.src = src
-        self.dst = dst
-        self.cond = cond
-
-
-class BasicBlock:
-    statements: list[IRStmt]
-    test: IRExpr
-    incoming: set[FlowEdge]
-    outgoing: set[FlowEdge]
-
-    def __init__(
-        self,
-        statements: list[IRStmt] | None = None,
-        test: IRExpr | None = None,
-        incoming: set[FlowEdge] | None = None,
-        outgoing: set[FlowEdge] | None = None,
-    ):
-        self.statements = statements or []
-        self.test = test or IRConst(1)
-        self.incoming = incoming or set()
-        self.outgoing = outgoing or set()
-
-    def connect_to(self, other: "BasicBlock", cond: IRExpr | None = None):
-        edge = FlowEdge(self, other, cond)
-        self.outgoing.add(edge)
-        other.incoming.add(edge)
-
-
-def traverse_cfg_preorder(block: BasicBlock) -> Iterator[BasicBlock]:
-    visited = set()
-    queue = deque([block])
-    while queue:
-        block = queue.popleft()
-        if block in visited:
-            continue
-        visited.add(block)
-        yield block
-        for edge in sorted(block.outgoing, key=lambda e: (e.cond is None, e.cond)):
-            queue.append(edge.dst)
-
-
-def cfg_to_mermaid(entry: BasicBlock):
-    def pre(s: str):
-        return "\"<pre style='text-align: left;'>" + s.replace("\n", "<br/>") + '</pre>"'
-
-    def fmt(nodes):
-        if nodes:
-            return "\n".join(str(n) for n in nodes)
-        else:
-            return "{}"
-
-    block_indexes = {block: i for i, block in enumerate(traverse_cfg_preorder(entry))}
-
-    lines = ["Entry([Entry]) --> 0"]
-    for block in traverse_cfg_preorder(entry):
-        index = block_indexes[block]
-        lines.append(f"{index}[{pre(fmt([f'#{index}', *block.statements]))}]")
-
-        outgoing = {edge.cond: edge.dst for edge in block.outgoing}
-        match outgoing:
-            case {**other} if not other:
-                lines.append(f"{index} --> Exit")
-            case {None: target, **other} if not other:
-                lines.append(f"{index} --> {block_indexes[target]}")
-            case {0: f_branch, None: t_branch, **other} if not other:
-                lines.append(f"{index}_{{{pre(fmt([block.test]))}}}")
-                lines.append(f"{index} --> {index}_")
-                lines.append(f"{index}_ --> |true| {block_indexes[t_branch]}")
-                lines.append(f"{index}_ --> |false| {block_indexes[f_branch]}")
-            case dict() as tgt:
-                lines.append(f"{index}_{{{{{pre(fmt([block.test]))}}}}}")
-                lines.append(f"{index} --> {index}_")
-                for cond, target in tgt.items():
-                    lines.append(f"{index}_ --> |{pre(fmt([cond or "default"]))}| {block_indexes[target]}")
-    lines.append("Exit([Exit])")
-
-    body = textwrap.indent("\n".join(lines), "    ")
-    return f"graph\n{body}"
+import textwrap
+from collections import deque
+from collections.abc import Iterator
+from typing import Self
+
+from sonolus.backend.ir import IRConst, IRExpr, IRStmt
+from sonolus.backend.place import SSAPlace, TempBlock
+
+
+class FlowEdge:
+    src: "BasicBlock"
+    dst: "BasicBlock"
+    cond: float | int | None
+
+    def __init__(self, src: "BasicBlock", dst: "BasicBlock", cond: float | None = None):
+        self.src = src
+        self.dst = dst
+        self.cond = cond
+
+
+class BasicBlock:
+    phis: dict[SSAPlace | TempBlock, dict[Self, SSAPlace]]
+    statements: list[IRStmt]
+    test: IRExpr
+    incoming: set[FlowEdge]
+    outgoing: set[FlowEdge]
+
+    def __init__(
+        self,
+        *,
+        phi: dict[SSAPlace, dict[Self, SSAPlace]] | None = None,
+        statements: list[IRStmt] | None = None,
+        test: IRExpr | None = None,
+        incoming: set[FlowEdge] | None = None,
+        outgoing: set[FlowEdge] | None = None,
+    ):
+        self.phis = phi or {}
+        self.statements = statements or []
+        self.test = test or IRConst(0)
+        self.incoming = incoming or set()
+        self.outgoing = outgoing or set()
+
+    def connect_to(self, other: "BasicBlock", cond: int | float | None = None):
+        edge = FlowEdge(self, other, cond)
+        self.outgoing.add(edge)
+        other.incoming.add(edge)
+
+
+def traverse_cfg_preorder(block: BasicBlock) -> Iterator[BasicBlock]:
+    visited = set()
+    queue = deque([block])
+    while queue:
+        block = queue.popleft()
+        if block in visited:
+            continue
+        visited.add(block)
+        yield block
+        for edge in sorted(block.outgoing, key=lambda e: (e.cond is None, e.cond)):
+            queue.append(edge.dst)
+
+
+def traverse_cfg_postorder(block: BasicBlock) -> Iterator[BasicBlock]:
+    visited = set()
+
+    def dfs(current: BasicBlock):
+        if current in visited:
+            return
+        visited.add(current)
+        for edge in sorted(current.outgoing, key=lambda e: (e.cond is None, e.cond)):
+            yield from dfs(edge.dst)
+        yield current
+
+    yield from dfs(block)
+
+
+def traverse_cfg_reverse_postorder(block: BasicBlock) -> Iterator[BasicBlock]:
+    yield from reversed(list(traverse_cfg_postorder(block)))
+
+
+def cfg_to_mermaid(entry: BasicBlock):
+    def pre(s: str):
+        return "\"<pre style='text-align: left;'>" + s.replace("\n", "<br/>") + '</pre>"'
+
+    def fmt(nodes):
+        if nodes:
+            return "\n".join(str(n) for n in nodes)
+        else:
+            return "{}"
+
+    block_indexes = {block: i for i, block in enumerate(traverse_cfg_preorder(entry))}
+
+    lines = ["Entry([Entry]) --> 0"]
+    for block in traverse_cfg_preorder(entry):
+        index = block_indexes[block]
+        lines.append(f"{index}[{pre(fmt([f'#{index}', *(
+            f"{dst} := phi({", ".join(f"{block_indexes.get(src_block, "<dead>")}: {src_place}"
+                                      for src_block, src_place
+                                      in sorted(phis.items(), key=lambda x: block_indexes.get(x[0])))})"
+            for dst, phis in block.phis.items()
+        ), *block.statements]))}]")
+
+        outgoing = {edge.cond: edge.dst for edge in block.outgoing}
+        match outgoing:
+            case {**other} if not other:
+                lines.append(f"{index} --> Exit")
+            case {None: target, **other} if not other:
+                lines.append(f"{index} --> {block_indexes[target]}")
+            case {0: f_branch, None: t_branch, **other} if not other:
+                lines.append(f"{index}_{{{pre(fmt([block.test]))}}}")
+                lines.append(f"{index} --> {index}_")
+                lines.append(f"{index}_ --> |true| {block_indexes[t_branch]}")
+                lines.append(f"{index}_ --> |false| {block_indexes[f_branch]}")
+            case dict() as tgt:
+                lines.append(f"{index}_{{{{{pre(fmt([block.test]))}}}}}")
+                lines.append(f"{index} --> {index}_")
+                for cond, target in tgt.items():
+                    lines.append(f"{index}_ --> |{pre(fmt([cond or "default"]))}| {block_indexes[target]}")
+    lines.append("Exit([Exit])")
+
+    body = textwrap.indent("\n".join(lines), "    ")
+    return f"graph\n{body}"

sonolus/backend/inlining.py

@@ -0,0 +1,150 @@
+from sonolus.backend.flow import BasicBlock, traverse_cfg_preorder
+from sonolus.backend.ir import IRConst, IRGet, IRInstr, IRPureInstr, IRSet, IRStmt
+from sonolus.backend.passes import CompilerPass
+from sonolus.backend.place import BlockPlace, SSAPlace, TempBlock
+
+
+class InlineVars(CompilerPass):
+    def run(self, entry: BasicBlock) -> BasicBlock:
+        use_counts: dict[SSAPlace, int] = {}
+        definitions: dict[SSAPlace, IRStmt] = {}
+
+        for block in traverse_cfg_preorder(entry):
+            for stmt in block.statements:
+                if (
+                    isinstance(stmt, IRSet)
+                    and isinstance(stmt.place, SSAPlace)
+                    and isinstance(stmt.value, IRGet)
+                    and isinstance(stmt.value.place, SSAPlace)
+                ):
+                    # This is effectively an alias
+                    pass
+                else:
+                    self.count_uses(stmt, use_counts)
+                if isinstance(stmt, IRSet) and isinstance(stmt.place, SSAPlace):
+                    definitions[stmt.place] = stmt.value
+            self.count_uses(block.test, use_counts)
+
+        for p, defn in definitions.items():
+            while True:
+                if isinstance(defn, IRGet) and isinstance(defn.place, SSAPlace) and defn.place in definitions:
+                    inside_defn = definitions[defn.place]
+                    if not self.is_inlinable(inside_defn):
+                        break
+                    defn = inside_defn
+                    continue
+                inlinable_uses = self.get_inlinable_uses(defn, set())
+                subs = {}
+                for inside_p in inlinable_uses:
+                    if inside_p not in definitions:
+                        continue
+                    inside_defn = definitions[inside_p]
+                    if not self.is_inlinable(inside_defn):
+                        continue
+                    if (isinstance(inside_defn, IRGet) and isinstance(inside_defn.place, SSAPlace)) or use_counts[
+                        inside_p
+                    ] == 1:
+                        subs[inside_p] = inside_defn
+                if not subs:
+                    break
+                defn = self.substitute(defn, subs)
+            definitions[p] = defn
+
+        valid = {
+            p
+            for p, count in use_counts.items()
+            if (p in definitions and count == 1 and self.is_inlinable(definitions[p]))
+        }
+
+        for block in traverse_cfg_preorder(entry):
+            new_statements = []
+            for stmt in [*block.statements, block.test]:
+                inlinable_uses = self.get_inlinable_uses(stmt, set())
+                subs = {}
+                for p in inlinable_uses:
+                    if p not in valid:
+                        continue
+                    definition = definitions[p]
+                    subs[p] = definition
+
+                if subs:
+                    new_statements.append(self.substitute(stmt, subs))
+                else:
+                    new_statements.append(stmt)
+
+            block.statements = new_statements[:-1]
+            block.test = new_statements[-1]
+
+        return entry
+
+    def substitute(self, stmt, subs):
+        match stmt:
+            case IRConst():
+                return stmt
+            case IRInstr(op=op, args=args):
+                return IRInstr(op=op, args=[self.substitute(arg, subs) for arg in args])
+            case IRPureInstr(op=op, args=args):
+                return IRPureInstr(op=op, args=[self.substitute(arg, subs) for arg in args])
+            case IRGet(place=place):
+                if place in subs:
+                    return subs[place]
+                return stmt
+            case IRSet(place=place, value=value):
+                return IRSet(place=place, value=self.substitute(value, subs))
+            case _:
+                raise TypeError(f"Unexpected statement: {stmt}")
+
+    def count_uses(self, stmt, use_counts):
+        match stmt:
+            case IRConst():
+                pass
+            case IRInstr(op=_, args=args) | IRPureInstr(op=_, args=args):
+                for arg in args:
+                    self.count_uses(arg, use_counts)
+            case IRGet(place=place):
+                self.count_uses(place, use_counts)
+            case IRSet(place=place, value=value):
+                if not isinstance(place, SSAPlace):  # We don't want to count the definition itself
+                    self.count_uses(place, use_counts)
+                self.count_uses(value, use_counts)
+            case SSAPlace():
+                use_counts[stmt] = use_counts.get(stmt, 0) + 1
+            case BlockPlace(block=block, index=index, offset=_):
+                self.count_uses(block, use_counts)
+                self.count_uses(index, use_counts)
+            case int() | float():
+                pass
+            case TempBlock():
+                pass
+            case _:
+                raise TypeError(f"Unexpected statement: {stmt}")
+        return use_counts
+
+    def get_inlinable_uses(self, stmt, uses):
+        match stmt:
+            case IRConst():
+                pass
+            case IRInstr(op=_, args=args) | IRPureInstr(op=_, args=args):
+                for arg in args:
+                    self.get_inlinable_uses(arg, uses)
+            case IRGet(place=place):
+                if isinstance(place, SSAPlace):
+                    uses.add(place)
+            case IRSet(place=_, value=value):
+                self.get_inlinable_uses(value, uses)
+            case _:
+                raise TypeError(f"Unexpected statement: {stmt}")
+        return uses
+
+    def is_inlinable(self, stmt):
+        match stmt:
+            case IRConst():
+                return True
+            case IRInstr(op=op, args=args) | IRPureInstr(op=op, args=args):
+                return not op.side_effects and op.pure and all(self.is_inlinable(arg) for arg in args)
+            case IRGet():
+                return isinstance(stmt.place, SSAPlace)
+            case IRSet():
+                return False
+            case _:
+                raise TypeError(f"Unexpected statement: {stmt}")

sonolus/backend/ir.py

@@ -6,9 +6,11 @@ type IRStmt = IRExpr | IRInstr | IRSet
 
 
 class IRConst:
-    value: float
+    value: float | int
 
     def __init__(self, value: float):
+        if isinstance(value, bool):
+            value = int(value)
         self.value = value
 
     def __repr__(self):
@@ -70,9 +72,9 @@ class IRGet:
 
 class IRSet:
     place: Place
-    value: IRStmt
+    value: IRExpr | IRInstr
 
-    def __init__(self, place: Place, value: IRStmt):
+    def __init__(self, place: Place, value: IRExpr | IRInstr):
         self.place = place
         self.value = value