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

sonolus/backend/{liveness.py → optimize/liveness.py}

@@ -1,8 +1,8 @@
 from collections import deque
 
-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.optimize.flow import BasicBlock, traverse_cfg_preorder
+from sonolus.backend.optimize.passes import CompilerPass
 from sonolus.backend.place import BlockPlace, SSAPlace, TempBlock
 
 type HasLiveness = SSAPlace | TempBlock
@@ -15,9 +15,26 @@ class LivenessAnalysis(CompilerPass):
     def run(self, entry: BasicBlock) -> BasicBlock:
         self.preprocess(entry)
         self.process(entry)
-        self.process_arrays(entry)
         return entry
 
+    def preprocess(self, entry: BasicBlock):
+        for block in traverse_cfg_preorder(entry):
+            block.live_out = None
+            block.live_in = set()
+            block.live_phi_targets = set()
+            block.array_defs_in = set()
+            block.array_defs_out = None
+            for statement in block.statements:
+                statement.live = set()
+                statement.visited = False
+                statement.uses = self.get_uses(statement, set())
+                statement.defs = self.get_defs(statement)
+                statement.is_array_init = False  # True if this may be the first assignment to an array
+                statement.array_defs = self.get_array_defs(statement)
+            block.test.live = set()
+            block.test.uses = self.get_uses(block.test, set())
+        self.preprocess_arrays(entry)
+
     def process(self, entry: BasicBlock):
         queue = deque(self.get_exits(entry))
         if not queue:
@@ -27,101 +44,88 @@ class LivenessAnalysis(CompilerPass):
             updated_blocks = self.process_block(block)
             queue.extend(updated_blocks)
 
-    def process_arrays(self, entry: BasicBlock):
-        # With arrays, we can't assume that an assignment will render previous assignments dead.
-        # Before this function is run, arrays are treated as live at a statement as long as they are read from
-        # at some future point.
-        # This function will mark arrays as dead if they could not have been assigned to yet.
-        queue = deque([entry])
+    def preprocess_arrays(self, entry: BasicBlock):
+        queue = {entry}
+        visited = set()
         while queue:
-            block = queue.popleft()
-            if block.live_arrays_in is None:
-                block.live_arrays_in = set()
-            live_arrays_in = block.live_arrays_in.copy()
+            block = queue.pop()
+            array_defs = block.array_defs_in.copy()
+            is_first_visit = block not in visited
+            visited.add(block)
             for statement in block.statements:
-                live_arrays_in.update(self.get_array_defs(statement))
-            updated_blocks = []
-            for edge in block.outgoing:
-                if edge.dst.live_arrays_in is None:
-                    prev_size = -1
-                    edge.dst.live_arrays_in = set()
+                if statement.array_defs - array_defs:
+                    statement.is_array_init = True
+                    array_defs.update(statement.array_defs)
                 else:
-                    prev_size = len(edge.dst.live_arrays_in)
-                edge.dst.live_arrays_in.update(live_arrays_in)
-                if len(edge.dst.live_arrays_in) != prev_size:
-                    updated_blocks.append(edge.dst)
-            queue.extend(updated_blocks)
-
-        for block in traverse_cfg_preorder(entry):
-            live_arrays_in = block.live_arrays_in
-            for statement in block.statements:
-                if not self.can_skip(statement, statement.live):
-                    live_arrays_in.update(self.get_array_defs(statement))
-                statement.live = {
-                    place
-                    for place in statement.live
-                    if not (isinstance(place, TempBlock) and place.size != 1 and place not in live_arrays_in)
-                }
-
-    def preprocess(self, entry: BasicBlock):
-        for block in traverse_cfg_preorder(entry):
-            block.live_out = None
-            block.live_in = None
-            block.live_phi_targets = None
-            block.live_arrays_in = None
+                    statement.is_array_init = False
+            if is_first_visit or array_defs != block.array_defs_out:
+                block.array_defs_out = array_defs
+                for edge in block.outgoing:
+                    queue.add(edge.dst)
+                    edge.dst.array_defs_in.update(array_defs)
 
     def process_block(self, block: BasicBlock) -> list[BasicBlock]:
         if block.live_out is None:
             block.live_out = set()
-        live: set[HasLiveness] = block.live_out.copy()
-        block.test.live = live.copy()
-        live.update(self.get_uses(block.test))
+        live: set[HasLiveness] = {
+            place
+            for place in block.live_out
+            if not (isinstance(place, TempBlock) and place.size > 1 and place not in block.array_defs_out)
+        }
+        block.test.live.update(live)
+        live.update(block.test.uses)
         for statement in reversed(block.statements):
-            statement.live = live.copy()
+            statement.live.update(live)
             if self.can_skip(statement, live):
                 continue
-            live.difference_update(self.get_defs(statement))
-            live.update(self.get_uses(statement))
+            live.difference_update(statement.defs)
+            if statement.is_array_init:
+                live.difference_update(statement.array_defs)
+            live.update(statement.uses)
+        prev_sizes_by_block = {
+            edge.src: len(edge.src.live_out) if edge.src.live_out is not None else -1 for edge in block.incoming
+        }
         live_phi_targets = set()
         for target, args in block.phis.items():
             if target not in live:
                 continue
-            live.difference_update({target})
-            live.update(args.values())
+            live.remove(target)
+            for src_block, arg in args.items():
+                if src_block.live_out is None:
+                    src_block.live_out = set()
+                src_block.live_out.add(arg)
             live_phi_targets.add(target)
-        block.live_in = live.copy()
+        block.live_in = live
         block.live_phi_targets = live_phi_targets
         updated_blocks = []
         for edge in block.incoming:
             if edge.src.live_out is None:
-                prev_size = -1
                 edge.src.live_out = set()
-            else:
-                prev_size = len(edge.src.live_out)
             edge.src.live_out.update(live)
-            if len(edge.src.live_out) != prev_size:
+            if len(edge.src.live_out) != prev_sizes_by_block[edge.src]:
                 updated_blocks.append(edge.src)
         return updated_blocks
 
-    def get_uses(self, stmt: IRStmt | BlockPlace | SSAPlace | TempBlock | int) -> set[HasLiveness]:
-        uses = set()
+    def get_uses(
+        self, stmt: IRStmt | BlockPlace | SSAPlace | TempBlock | int, uses: set[HasLiveness]
+    ) -> set[HasLiveness]:
         match stmt:
             case IRPureInstr(op=_, args=args) | IRInstr(op=_, args=args):
                 for arg in args:
-                    uses.update(self.get_uses(arg))
+                    self.get_uses(arg, uses)
             case IRGet(place=place):
-                uses.update(self.get_uses(place))
+                self.get_uses(place, uses)
             case IRSet(place=place, value=value):
                 if isinstance(place, BlockPlace):
                     if not isinstance(place.block, TempBlock):
-                        uses.update(self.get_uses(place.block))
-                    uses.update(self.get_uses(place.index))
-                uses.update(self.get_uses(value))
+                        self.get_uses(place.block, uses)
+                    self.get_uses(place.index, uses)
+                self.get_uses(value, uses)
             case IRConst() | int():
                 pass
             case BlockPlace(block=block, index=index, offset=_):
-                uses.update(self.get_uses(block))
-                uses.update(self.get_uses(index))
+                self.get_uses(block, uses)
+                self.get_uses(index, uses)
             case SSAPlace() | TempBlock():
                 uses.add(stmt)
             case _:
@@ -151,7 +155,7 @@ class LivenessAnalysis(CompilerPass):
             case IRSet(place=_, value=value):
                 if isinstance(value, IRInstr) and value.op.side_effects:
                     return False
-                defs = self.get_defs(stmt) | self.get_array_defs(stmt)
+                defs = stmt.defs | stmt.array_defs
                 return defs and not (defs & live)
         return False
 

sonolus/backend/optimize/optimize.py

@@ -0,0 +1,44 @@
+from sonolus.backend.optimize.allocate import Allocate, AllocateBasic
+from sonolus.backend.optimize.constant_evaluation import SparseConditionalConstantPropagation
+from sonolus.backend.optimize.copy_coalesce import CopyCoalesce
+from sonolus.backend.optimize.dead_code import (
+    AdvancedDeadCodeElimination,
+    DeadCodeElimination,
+    UnreachableCodeElimination,
+)
+from sonolus.backend.optimize.flow import BasicBlock
+from sonolus.backend.optimize.inlining import InlineVars
+from sonolus.backend.optimize.passes import run_passes
+from sonolus.backend.optimize.simplify import CoalesceFlow, NormalizeSwitch, RewriteToSwitch
+from sonolus.backend.optimize.ssa import FromSSA, ToSSA
+
+MINIMAL_PASSES = [
+    CoalesceFlow(),
+    UnreachableCodeElimination(),
+    AllocateBasic(),
+]
+
+STANDARD_PASSES = [
+    CoalesceFlow(),
+    UnreachableCodeElimination(),
+    DeadCodeElimination(),
+    ToSSA(),
+    SparseConditionalConstantPropagation(),
+    UnreachableCodeElimination(),
+    DeadCodeElimination(),
+    CoalesceFlow(),
+    InlineVars(),
+    DeadCodeElimination(),
+    RewriteToSwitch(),
+    FromSSA(),
+    CoalesceFlow(),
+    CopyCoalesce(),
+    AdvancedDeadCodeElimination(),
+    CoalesceFlow(),
+    NormalizeSwitch(),
+    Allocate(),
+]
+
+
+def optimize_and_allocate(cfg: BasicBlock):
+    return run_passes(cfg, STANDARD_PASSES)

sonolus/backend/{passes.py → optimize/passes.py}

@@ -3,7 +3,7 @@ from __future__ import annotations
 from abc import ABC, abstractmethod
 from collections import deque
 
-from sonolus.backend.flow import BasicBlock
+from sonolus.backend.optimize.flow import BasicBlock
 
 
 class CompilerPass(ABC):

sonolus/backend/optimize/simplify.py

@@ -0,0 +1,191 @@
+from sonolus.backend.ir import IRConst, IRGet, IRPureInstr, IRSet
+from sonolus.backend.ops import Op
+from sonolus.backend.optimize.flow import BasicBlock, traverse_cfg_preorder
+from sonolus.backend.optimize.passes import CompilerPass
+
+
+class CoalesceFlow(CompilerPass):
+    def run(self, entry: BasicBlock) -> BasicBlock:
+        queue = [entry]
+        processed = set()
+        while queue:
+            block = queue.pop()
+            if block in processed:
+                continue
+            processed.add(block)
+            for edge in block.outgoing:
+                while True:
+                    dst = edge.dst
+                    if dst.phis or dst.statements or len(dst.outgoing) != 1 or dst is block:
+                        break
+                    next_dst = next(iter(dst.outgoing)).dst
+                    if next_dst.phis:
+                        break
+                    dst.incoming.remove(edge)
+                    if not dst.incoming:
+                        for dst_edge in dst.outgoing:
+                            dst_edge.dst.incoming.remove(dst_edge)
+                        processed.add(dst)
+                    edge.dst = next_dst
+                    next_dst.incoming.add(edge)
+                    if dst is edge.dst:
+                        break
+            default_edge = next((edge for edge in block.outgoing if edge.cond is None), None)
+            if default_edge is not None:
+                for edge in [*block.outgoing]:
+                    if edge is default_edge:
+                        continue
+                    if edge.dst is default_edge.dst:
+                        block.outgoing.remove(edge)
+                        edge.dst.incoming.remove(edge)
+            if len(block.outgoing) != 1:
+                queue.extend(edge.dst for edge in block.outgoing)
+                continue
+            next_block = next(iter(block.outgoing)).dst
+            if len(next_block.incoming) != 1:
+                queue.append(next_block)
+                if not block.statements and not block.phis and not next_block.phis:
+                    for edge in block.incoming:
+                        edge.dst = next_block
+                        next_block.incoming.add(edge)
+                    for edge in block.outgoing:  # There should be exactly one
+                        next_block.incoming.remove(edge)
+                    if block is entry:
+                        entry = next_block
+                continue
+            for p, args in next_block.phis.items():
+                if block not in args:
+                    continue
+                block.statements.append(IRSet(p, IRGet(args[block])))
+            block.statements.extend(next_block.statements)
+            block.test = next_block.test
+            block.outgoing = next_block.outgoing
+            for edge in block.outgoing:
+                edge.src = block
+                dst = edge.dst
+                for args in dst.phis.values():
+                    if next_block in args:
+                        args[block] = args.pop(next_block)
+            processed.add(next_block)
+            queue.extend(edge.dst for edge in block.outgoing)
+            processed.remove(block)
+            queue.append(block)
+        return entry
+
+
+class RewriteToSwitch(CompilerPass):
+    """Rewrite if-else chains to switch statements.
+
+    Note that this needs inlining (and dead code elimination) to be run first to really do anything useful.
+    """
+
+    def run(self, entry: BasicBlock) -> BasicBlock:
+        self.ifs_to_switch(entry)
+        self.combine_blocks(entry)
+        self.remove_unreachable(entry)
+        return entry
+
+    def ifs_to_switch(self, entry: BasicBlock):
+        for block in traverse_cfg_preorder(entry):
+            if len(block.outgoing) != 2 or {edge.cond for edge in block.outgoing} != {None, 0}:
+                continue
+            test = block.test
+            if not isinstance(test, IRPureInstr) or test.op != Op.Equal:
+                continue
+            assert len(test.args) == 2
+            if isinstance(test.args[0], IRConst):
+                const, other = test.args
+            elif isinstance(test.args[1], IRConst):
+                other, const = test.args
+            else:
+                continue
+            block.test = other
+            for edge in block.outgoing:
+                if edge.cond is None:
+                    edge.cond = const.value
+                else:
+                    edge.cond = None
+
+    def combine_blocks(self, entry: BasicBlock):
+        queue = [entry]
+        processed = set()
+        while queue:
+            block = queue.pop()
+            if block in processed:
+                continue
+            processed.add(block)
+            queue.extend(edge.dst for edge in block.outgoing)
+
+            default_edge = next((edge for edge in block.outgoing if edge.cond is None), None)
+            if default_edge is None:
+                continue
+
+            next_block = default_edge.dst
+            if (
+                len(next_block.incoming) > 1
+                or next_block.statements
+                or next_block.phis
+                or block.test != next_block.test
+                or block is next_block
+                or next_block is entry
+            ):
+                continue
+
+            outgoing_by_cond = {edge.cond: edge for edge in block.outgoing}
+            assert len(outgoing_by_cond) == len(block.outgoing)
+            outgoing_by_cond.pop(None)
+            for edge in next_block.outgoing:
+                if edge.cond in outgoing_by_cond:
+                    # This edge is unreachable since an equivalent edge would have been taken
+                    edge.dst.incoming.remove(edge)
+                    continue
+                outgoing_by_cond[edge.cond] = edge
+                edge.src = block
+                for args in edge.dst.phis.values():
+                    if next_block in args:
+                        args[block] = args.pop(next_block)
+            block.outgoing = set(outgoing_by_cond.values())
+            processed.add(next_block)
+            queue.append(block)
+            processed.remove(block)
+
+    def remove_unreachable(self, entry: BasicBlock):
+        reachable = {*traverse_cfg_preorder(entry)}
+        for block in traverse_cfg_preorder(entry):
+            block.incoming = {edge for edge in block.incoming if edge.src in reachable}
+            block.outgoing = {edge for edge in block.outgoing if edge.dst in reachable}
+
+
+class NormalizeSwitch(CompilerPass):
+    """Normalize branches like cond -> case a, case a + b, case a + 2b to ((cond - a) / b) -> case 0, case 1, case 2."""
+
+    def run(self, entry: BasicBlock) -> BasicBlock:
+        for block in traverse_cfg_preorder(entry):
+            cases = {edge.cond for edge in block.outgoing}
+            if len(cases) <= 2:
+                continue
+            assert None in cases, "Non-terminal blocks should always have a default edge"
+            cases.remove(None)
+            offset, stride = self.get_offset_stride(cases)
+            if offset is None or (offset == 0 and stride == 1):
+                continue
+            for edge in block.outgoing:
+                if edge.cond is None:
+                    continue
+                edge.cond = (edge.cond - offset) // stride
+            if offset != 0:
+                block.test = IRPureInstr(Op.Subtract, [block.test, IRConst(offset)])
+            if stride != 1:
+                block.test = IRPureInstr(Op.Divide, [block.test, IRConst(stride)])
+        return entry
+
+    def get_offset_stride(self, cases: set[int]) -> tuple[int | None, int | None]:
+        cases = sorted(cases)
+        offset = cases[0]
+        stride = cases[1] - offset
+        if int(offset) != offset or int(stride) != stride:
+            return None, None
+        for i, case in enumerate(cases[2:], 2):
+            if case != offset + i * stride:
+                return None, None
+        return offset, stride

sonolus/backend/{ssa.py → optimize/ssa.py}

@@ -1,7 +1,7 @@
-from sonolus.backend.dominance import DominanceFrontiers, get_df, get_dom_children
-from sonolus.backend.flow import BasicBlock, FlowEdge, traverse_cfg_preorder
 from sonolus.backend.ir import IRConst, IRGet, IRInstr, IRPureInstr, IRSet, IRStmt
-from sonolus.backend.passes import CompilerPass
+from sonolus.backend.optimize.dominance import DominanceFrontiers, get_df, get_dom_children
+from sonolus.backend.optimize.flow import BasicBlock, FlowEdge, traverse_cfg_preorder
+from sonolus.backend.optimize.passes import CompilerPass
 from sonolus.backend.place import BlockPlace, SSAPlace, TempBlock
 
 
@@ -23,13 +23,14 @@ class ToSSA(CompilerPass):
         ssa_places: dict[TempBlock, list[SSAPlace]],
         used: dict[str, int],
     ):
-        original_ssa_place_lens = {var: len(ssa_places[var]) for var in defs}
+        to_pop = []
         for var, args in [*block.phis.items()]:
             if isinstance(var, SSAPlace):
                 continue
             ssa_places[var].append(self.get_new_ssa_place(var.name, used))
+            to_pop.append(var)
             block.phis[ssa_places[var][-1]] = args
-        block.statements = [self.rename_stmt(stmt, ssa_places, used) for stmt in block.statements]
+        block.statements = [self.rename_stmt(stmt, ssa_places, used, to_pop) for stmt in block.statements]
         for edge in block.outgoing:
             dst = edge.dst
             for var, args in dst.phis.items():
@@ -37,35 +38,47 @@ class ToSSA(CompilerPass):
                     continue
                 if ssa_places[var]:
                     args[block] = ssa_places[var][-1]
-        block.test = self.rename_stmt(block.test, ssa_places, used)
+        block.test = self.rename_stmt(block.test, ssa_places, used, to_pop)
         for dom_child in get_dom_children(block):
             self.rename(dom_child, defs, ssa_places, used)
-        for var, length in original_ssa_place_lens.items():
-            ssa_places[var] = ssa_places[var][:length]
+        for var in to_pop:
+            ssa_places[var].pop()
 
     def remove_placeholder_phis(self, entry: BasicBlock):
         for block in traverse_cfg_preorder(entry):
             block.phis = {var: args for var, args in block.phis.items() if isinstance(var, SSAPlace)}
 
-    def rename_stmt(self, stmt: IRStmt, ssa_places: dict[TempBlock, list[SSAPlace]], used: dict[str, int]):
+    def rename_stmt(
+        self, stmt: IRStmt, ssa_places: dict[TempBlock, list[SSAPlace]], used: dict[str, int], to_pop: list[SSAPlace]
+    ):
         match stmt:
             case IRConst():
                 return stmt
             case IRPureInstr(op=op, args=args):
-                return IRPureInstr(op=op, args=[self.rename_stmt(arg, ssa_places, used) for arg in args])
+                return IRPureInstr(op=op, args=[self.rename_stmt(arg, ssa_places, used, to_pop) for arg in args])
             case IRInstr(op=op, args=args):
-                return IRInstr(op=op, args=[self.rename_stmt(arg, ssa_places, used) for arg in args])
+                return IRInstr(op=op, args=[self.rename_stmt(arg, ssa_places, used, to_pop) for arg in args])
             case IRGet(place=place):
-                return IRGet(place=self.rename_stmt(place, ssa_places, used))
+                return IRGet(place=self.rename_stmt(place, ssa_places, used, to_pop))
             case IRSet(place=place, value=value):
-                value = self.rename_stmt(value, ssa_places, used)
+                value = self.rename_stmt(value, ssa_places, used, to_pop)
                 if isinstance(place, BlockPlace) and isinstance(place.block, TempBlock) and place.block.size == 1:
                     ssa_places[place.block].append(self.get_new_ssa_place(place.block.name, used))
-                place = self.rename_stmt(place, ssa_places, used)
+                    to_pop.append(place.block)
+                place = self.rename_stmt(place, ssa_places, used, to_pop)
                 return IRSet(place=place, value=value)
             case SSAPlace():
                 return stmt
             case TempBlock() if stmt.size == 1:
+                if stmt not in ssa_places or not ssa_places[stmt]:
+                    # This is an access to a definitely undefined variable
+                    # But it might not be reachable in reality, so we should allow it
+                    # Maybe there should be an error if this still happens after optimization,
+                    # but recovering the location of the error in the original code is hard.
+                    # This can happen in places like matching a VarArray[Num, 1] which was just created.
+                    # IR generation won't immediately fold a check that size > 0 to false, so here we
+                    # might see an access to uninitialized memory even though it's not reachable in reality.
+                    return SSAPlace("err", 0)
                 return ssa_places[stmt][-1]
             case TempBlock():
                 return stmt
@@ -73,11 +86,11 @@ class ToSSA(CompilerPass):
                 return stmt
             case BlockPlace(block=block, index=index, offset=offset):
                 if isinstance(block, TempBlock) and block.size == 1:
-                    return self.rename_stmt(block, ssa_places, used)
+                    return self.rename_stmt(block, ssa_places, used, to_pop)
                 return BlockPlace(
-                    block=self.rename_stmt(block, ssa_places, used),
-                    index=self.rename_stmt(index, ssa_places, used),
-                    offset=self.rename_stmt(offset, ssa_places, used),
+                    block=self.rename_stmt(block, ssa_places, used, to_pop),
+                    index=self.rename_stmt(index, ssa_places, used, to_pop),
+                    offset=self.rename_stmt(offset, ssa_places, used, to_pop),
                 )
             case _:
                 raise TypeError(f"Unexpected statement: {stmt}")

sonolus/backend/place.py

@@ -1,5 +1,5 @@
 from collections.abc import Iterator
-from typing import Self
+from typing import NamedTuple, Self
 
 from sonolus.backend.blocks import Block
 
@@ -8,13 +8,9 @@ type BlockValue = Block | int | TempBlock | Place
 type IndexValue = int | Place
 
 
-class TempBlock:
+class TempBlock(NamedTuple):
     name: str
-    size: int
-
-    def __init__(self, name: str, size: int = 1):
-        self.name = name
-        self.size = size
+    size: int = 1
 
     def __repr__(self):
         return f"TempBlock(name={self.name!r}, size={self.size!r})"
@@ -33,19 +29,14 @@ class TempBlock:
         return isinstance(other, TempBlock) and self.name == other.name and self.size == other.size
 
     def __hash__(self):
-        return hash((self.name, self.size))
+        return hash(self.name)  # Typically will be unique by name alone
 
 
-class BlockPlace:
+class BlockPlace(NamedTuple):
     block: BlockValue
-    index: IndexValue
+    index: IndexValue = 0
     offset: int = 0
 
-    def __init__(self, block: BlockValue, index: IndexValue = 0, offset: int = 0):
-        self.block = block
-        self.index = index
-        self.offset = offset
-
     def __repr__(self):
         return f"BlockPlace(block={self.block!r}, index={self.index!r}, offset={self.offset!r})"
 
@@ -59,24 +50,25 @@ class BlockPlace:
         else:
             return f"{self.block}[{self.index} + {self.offset}]"
 
+    def add_offset(self, offset: int) -> Self:
+        return BlockPlace(self.block, self.index, self.offset + offset)
+
     def __eq__(self, other):
-        return isinstance(other, BlockPlace) and self.block == other.block and self.index == other.index
+        return (
+            isinstance(other, BlockPlace)
+            and self.block == other.block
+            and self.index == other.index
+            and self.offset == other.offset
+        )
 
     def __hash__(self):
-        return hash((self.block, self.index))
+        return hash((self.block, self.index, self.offset))
 
-    def add_offset(self, offset: int) -> Self:
-        return BlockPlace(self.block, self.index, self.offset + offset)
 
-
-class SSAPlace:
+class SSAPlace(NamedTuple):
     name: str
     num: int
 
-    def __init__(self, name: str, num: int):
-        self.name = name
-        self.num = num
-
     def __repr__(self):
         return f"SSAPlace(name={self.name!r}, num={self.num!r})"
 

sonolus/backend/utils.py

@@ -2,9 +2,11 @@
 import ast
 import inspect
 from collections.abc import Callable
+from functools import cache
 from pathlib import Path
 
 
+@cache
 def get_function(fn: Callable) -> tuple[str, ast.FunctionDef]:
     # This preserves both line number and column number in the returned node
     source_file = inspect.getsourcefile(fn)
@@ -26,6 +28,14 @@ class FindFunction(ast.NodeVisitor):
         else:
             self.generic_visit(node)
 
+    def visit_Lambda(self, node: ast.Lambda):
+        if node.lineno == self.line:
+            if self.node is not None:
+                raise ValueError("Multiple functions defined on the same line are not supported")
+            self.node = node
+        else:
+            self.generic_visit(node)
+
 
 def find_function(tree: ast.Module, line: int):
     visitor = FindFunction(line)