sonolus.py 0.1.0__py3-none-any.whl

sonolus/backend/ops.py

@@ -0,0 +1,197 @@
+from enum import StrEnum
+
+
+class Op(StrEnum):
+    def __new__(cls, name: str, side_effects: bool, pure: bool, control_flow: bool):
+        obj = str.__new__(cls, name)
+        obj._value_ = name
+        obj.side_effects = side_effects
+        obj.pure = pure
+        obj.control_flow = control_flow
+        return obj
+
+    Abs = ("Abs", False, True, False)
+    Add = ("Add", False, True, False)
+    And = ("And", False, True, True)
+    Arccos = ("Arccos", False, True, False)
+    Arcsin = ("Arcsin", False, True, False)
+    Arctan = ("Arctan", False, True, False)
+    Arctan2 = ("Arctan2", False, True, False)
+    BeatToBPM = ("BeatToBPM", False, True, False)
+    BeatToStartingBeat = ("BeatToStartingBeat", False, True, False)
+    BeatToStartingTime = ("BeatToStartingTime", False, True, False)
+    BeatToTime = ("BeatToTime", False, True, False)
+    Block = ("Block", False, True, True)
+    Break = ("Break", True, False, True)
+    Ceil = ("Ceil", False, True, False)
+    Clamp = ("Clamp", False, True, False)
+    Copy = ("Copy", True, False, False)
+    Cos = ("Cos", False, True, False)
+    Cosh = ("Cosh", False, True, False)
+    DebugLog = ("DebugLog", True, False, False)
+    DebugPause = ("DebugPause", True, False, False)
+    DecrementPost = ("DecrementPost", True, False, False)
+    DecrementPostPointed = ("DecrementPostPointed", True, False, False)
+    DecrementPostShifted = ("DecrementPostShifted", True, False, False)
+    DecrementPre = ("DecrementPre", True, False, False)
+    DecrementPrePointed = ("DecrementPrePointed", True, False, False)
+    DecrementPreShifted = ("DecrementPreShifted", True, False, False)
+    Degree = ("Degree", False, True, False)
+    DestroyParticleEffect = ("DestroyParticleEffect", True, False, False)
+    Divide = ("Divide", False, True, False)
+    DoWhile = ("DoWhile", False, True, True)
+    Draw = ("Draw", True, False, False)
+    DrawCurvedB = ("DrawCurvedB", True, False, False)
+    DrawCurvedBT = ("DrawCurvedBT", True, False, False)
+    DrawCurvedL = ("DrawCurvedL", True, False, False)
+    DrawCurvedLR = ("DrawCurvedLR", True, False, False)
+    DrawCurvedR = ("DrawCurvedR", True, False, False)
+    DrawCurvedT = ("DrawCurvedT", True, False, False)
+    EaseInBack = ("EaseInBack", False, True, False)
+    EaseInCirc = ("EaseInCirc", False, True, False)
+    EaseInCubic = ("EaseInCubic", False, True, False)
+    EaseInElastic = ("EaseInElastic", False, True, False)
+    EaseInExpo = ("EaseInExpo", False, True, False)
+    EaseInOutBack = ("EaseInOutBack", False, True, False)
+    EaseInOutCirc = ("EaseInOutCirc", False, True, False)
+    EaseInOutCubic = ("EaseInOutCubic", False, True, False)
+    EaseInOutElastic = ("EaseInOutElastic", False, True, False)
+    EaseInOutExpo = ("EaseInOutExpo", False, True, False)
+    EaseInOutQuad = ("EaseInOutQuad", False, True, False)
+    EaseInOutQuart = ("EaseInOutQuart", False, True, False)
+    EaseInOutQuint = ("EaseInOutQuint", False, True, False)
+    EaseInOutSine = ("EaseInOutSine", False, True, False)
+    EaseInQuad = ("EaseInQuad", False, True, False)
+    EaseInQuart = ("EaseInQuart", False, True, False)
+    EaseInQuint = ("EaseInQuint", False, True, False)
+    EaseInSine = ("EaseInSine", False, True, False)
+    EaseOutBack = ("EaseOutBack", False, True, False)
+    EaseOutCirc = ("EaseOutCirc", False, True, False)
+    EaseOutCubic = ("EaseOutCubic", False, True, False)
+    EaseOutElastic = ("EaseOutElastic", False, True, False)
+    EaseOutExpo = ("EaseOutExpo", False, True, False)
+    EaseOutInBack = ("EaseOutInBack", False, True, False)
+    EaseOutInCirc = ("EaseOutInCirc", False, True, False)
+    EaseOutInCubic = ("EaseOutInCubic", False, True, False)
+    EaseOutInElastic = ("EaseOutInElastic", False, True, False)
+    EaseOutInExpo = ("EaseOutInExpo", False, True, False)
+    EaseOutInQuad = ("EaseOutInQuad", False, True, False)
+    EaseOutInQuart = ("EaseOutInQuart", False, True, False)
+    EaseOutInQuint = ("EaseOutInQuint", False, True, False)
+    EaseOutInSine = ("EaseOutInSine", False, True, False)
+    EaseOutQuad = ("EaseOutQuad", False, True, False)
+    EaseOutQuart = ("EaseOutQuart", False, True, False)
+    EaseOutQuint = ("EaseOutQuint", False, True, False)
+    EaseOutSine = ("EaseOutSine", False, True, False)
+    Equal = ("Equal", False, True, False)
+    Execute = ("Execute", False, True, False)
+    Execute0 = ("Execute0", False, True, False)
+    ExportValue = ("ExportValue", True, False, False)
+    Floor = ("Floor", False, True, False)
+    Frac = ("Frac", False, True, False)
+    Get = ("Get", False, False, False)
+    GetPointed = ("GetPointed", False, False, False)
+    GetShifted = ("GetShifted", False, False, False)
+    Greater = ("Greater", False, True, False)
+    GreaterOr = ("GreaterOr", False, True, False)
+    HasEffectClip = ("HasEffectClip", False, True, False)
+    HasParticleEffect = ("HasParticleEffect", False, True, False)
+    HasSkinSprite = ("HasSkinSprite", False, True, False)
+    If = ("If", False, True, True)
+    IncrementPost = ("IncrementPost", True, False, False)
+    IncrementPostPointed = ("IncrementPostPointed", True, False, False)
+    IncrementPostShifted = ("IncrementPostShifted", True, False, False)
+    IncrementPre = ("IncrementPre", True, False, False)
+    IncrementPrePointed = ("IncrementPrePointed", True, False, False)
+    IncrementPreShifted = ("IncrementPreShifted", True, False, False)
+    Judge = ("Judge", False, True, False)
+    JudgeSimple = ("JudgeSimple", False, True, False)
+    JumpLoop = ("JumpLoop", False, True, True)
+    Lerp = ("Lerp", False, True, False)
+    LerpClamped = ("LerpClamped", False, True, False)
+    Less = ("Less", False, True, False)
+    LessOr = ("LessOr", False, True, False)
+    Log = ("Log", False, True, False)
+    Max = ("Max", False, True, False)
+    Min = ("Min", False, True, False)
+    Mod = ("Mod", False, True, False)
+    MoveParticleEffect = ("MoveParticleEffect", True, False, False)
+    Multiply = ("Multiply", False, True, False)
+    Negate = ("Negate", False, True, False)
+    Not = ("Not", False, True, False)
+    NotEqual = ("NotEqual", False, True, False)
+    Or = ("Or", False, True, True)
+    Paint = ("Paint", True, False, False)
+    Play = ("Play", True, False, False)
+    PlayLooped = ("PlayLooped", True, False, False)
+    PlayLoopedScheduled = ("PlayLoopedScheduled", True, False, False)
+    PlayScheduled = ("PlayScheduled", True, False, False)
+    Power = ("Power", False, True, False)
+    Print = ("Print", True, False, False)
+    Radian = ("Radian", False, True, False)
+    Random = ("Random", False, False, False)
+    RandomInteger = ("RandomInteger", False, False, False)
+    Rem = ("Rem", False, True, False)
+    Remap = ("Remap", False, True, False)
+    RemapClamped = ("RemapClamped", False, True, False)
+    Round = ("Round", False, True, False)
+    Set = ("Set", True, False, False)
+    SetAdd = ("SetAdd", True, False, False)
+    SetAddPointed = ("SetAddPointed", True, False, False)
+    SetAddShifted = ("SetAddShifted", True, False, False)
+    SetDivide = ("SetDivide", True, False, False)
+    SetDividePointed = ("SetDividePointed", True, False, False)
+    SetDivideShifted = ("SetDivideShifted", True, False, False)
+    SetMod = ("SetMod", True, False, False)
+    SetModPointed = ("SetModPointed", True, False, False)
+    SetModShifted = ("SetModShifted", True, False, False)
+    SetMultiply = ("SetMultiply", True, False, False)
+    SetMultiplyPointed = ("SetMultiplyPointed", True, False, False)
+    SetMultiplyShifted = ("SetMultiplyShifted", True, False, False)
+    SetPointed = ("SetPointed", True, False, False)
+    SetPower = ("SetPower", True, False, False)
+    SetPowerPointed = ("SetPowerPointed", True, False, False)
+    SetPowerShifted = ("SetPowerShifted", True, False, False)
+    SetRem = ("SetRem", True, False, False)
+    SetRemPointed = ("SetRemPointed", True, False, False)
+    SetRemShifted = ("SetRemShifted", True, False, False)
+    SetShifted = ("SetShifted", True, False, False)
+    SetSubtract = ("SetSubtract", True, False, False)
+    SetSubtractPointed = ("SetSubtractPointed", True, False, False)
+    SetSubtractShifted = ("SetSubtractShifted", True, False, False)
+    Sign = ("Sign", False, True, False)
+    Sin = ("Sin", False, True, False)
+    Sinh = ("Sinh", False, True, False)
+    Spawn = ("Spawn", True, False, False)
+    SpawnParticleEffect = ("SpawnParticleEffect", True, False, False)
+    StackEnter = ("StackEnter", True, False, False)
+    StackGet = ("StackGet", False, False, False)
+    StackGetFrame = ("StackGetFrame", False, False, False)
+    StackGetFramePointer = ("StackGetFramePointer", False, False, False)
+    StackGetPointer = ("StackGetPointer", False, False, False)
+    StackGrow = ("StackGrow", True, False, False)
+    StackInit = ("StackInit", True, False, False)
+    StackLeave = ("StackLeave", True, False, False)
+    StackPop = ("StackPop", True, False, False)
+    StackPush = ("StackPush", True, False, False)
+    StackSet = ("StackSet", True, False, False)
+    StackSetFrame = ("StackSetFrame", True, False, False)
+    StackSetFramePointer = ("StackSetFramePointer", True, False, False)
+    StackSetPointer = ("StackSetPointer", True, False, False)
+    StopLooped = ("StopLooped", True, False, False)
+    StopLoopedScheduled = ("StopLoopedScheduled", True, False, False)
+    Subtract = ("Subtract", False, True, False)
+    Switch = ("Switch", False, True, True)
+    SwitchInteger = ("SwitchInteger", False, True, True)
+    SwitchIntegerWithDefault = ("SwitchIntegerWithDefault", False, True, True)
+    SwitchWithDefault = ("SwitchWithDefault", False, True, True)
+    Tan = ("Tan", False, True, False)
+    Tanh = ("Tanh", False, True, False)
+    TimeToScaledTime = ("TimeToScaledTime", False, True, False)
+    TimeToStartingScaledTime = ("TimeToStartingScaledTime", False, True, False)
+    TimeToStartingTime = ("TimeToStartingTime", False, True, False)
+    TimeToTimeScale = ("TimeToTimeScale", False, True, False)
+    Trunc = ("Trunc", False, True, False)
+    Unlerp = ("Unlerp", False, True, False)
+    UnlerpClamped = ("UnlerpClamped", False, True, False)
+    While = ("While", False, True, True)

sonolus/backend/optimize.py

@@ -0,0 +1,9 @@
+from sonolus.backend.allocate import AllocateBasic
+from sonolus.backend.flow import BasicBlock
+from sonolus.backend.simplify import CoalesceFlow
+
+
+def optimize_and_allocate(cfg: BasicBlock):
+    cfg = CoalesceFlow().run(cfg)
+    cfg = AllocateBasic().run(cfg)
+    return cfg

sonolus/backend/passes.py

@@ -0,0 +1,6 @@
+from sonolus.backend.flow import BasicBlock
+
+
+class CompilerPass:
+    def run(self, entry: BasicBlock) -> BasicBlock:
+        pass

sonolus/backend/place.py

@@ -0,0 +1,90 @@
+from collections.abc import Iterator
+from typing import Self
+
+from sonolus.backend.blocks import Block
+
+type Place = BlockPlace | SSAPlace
+type BlockValue = Block | int | TempBlock | Place
+type IndexValue = int | Place
+
+
+class TempBlock:
+    name: str
+    size: int
+
+    def __init__(self, name: str, size: int = 1):
+        self.name = name
+        self.size = size
+
+    def __repr__(self):
+        return f"TempBlock(name={self.name!r}, size={self.size!r})"
+
+    def __str__(self):
+        return f"{self.name}"
+
+    def __getitem__(self, item) -> "BlockPlace":
+        return BlockPlace(self, item)
+
+    def __iter__(self) -> "Iterator[BlockPlace]":
+        for i in range(self.size):
+            yield self[i]
+
+    def __eq__(self, other):
+        return isinstance(other, TempBlock) and self.name == other.name and self.size == other.size
+
+    def __hash__(self):
+        return hash((self.name, self.size))
+
+
+class BlockPlace:
+    block: BlockValue
+    index: IndexValue
+    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})"
+
+    def __str__(self):
+        if isinstance(self.block, TempBlock) and self.block.size == 1 and self.index == 0 and self.offset == 0:
+            return f"{self.block}"
+        elif isinstance(self.index, int):
+            return f"{self.block}[{self.index + self.offset}]"
+        elif self.offset == 0:
+            return f"{self.block}[{self.index}]"
+        else:
+            return f"{self.block}[{self.index} + {self.offset}]"
+
+    def __eq__(self, other):
+        return isinstance(other, BlockPlace) and self.block == other.block and self.index == other.index
+
+    def __hash__(self):
+        return hash((self.block, self.index))
+
+    def add_offset(self, offset: int) -> Self:
+        return BlockPlace(self.block, self.index, self.offset + offset)
+
+
+class SSAPlace:
+    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})"
+
+    def __str__(self):
+        return f"{self.name}.{self.num}"
+
+    def __eq__(self, other):
+        return isinstance(other, SSAPlace) and self.name == other.name and self.num == other.num
+
+    def __hash__(self):
+        return hash((self.name, self.num))

sonolus/backend/simplify.py

@@ -0,0 +1,30 @@
+from sonolus.backend.flow import BasicBlock
+from sonolus.backend.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)
+            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)
+                continue
+            block.statements.extend(next_block.statements)
+            block.test = next_block.test
+            for edge in next_block.outgoing:
+                edge.src = block
+            block.outgoing = next_block.outgoing
+            processed.add(next_block)
+            queue.extend(edge.dst for edge in block.outgoing)
+            processed.remove(block)
+            queue.append(block)
+        return entry

sonolus/backend/utils.py

@@ -0,0 +1,48 @@
+# ruff: noqa: N802
+import ast
+import inspect
+from collections.abc import Callable
+from pathlib import Path
+
+
+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)
+    _, start_line = inspect.getsourcelines(fn)
+    base_tree = ast.parse(Path(source_file).read_text(encoding="utf-8"))
+    return source_file, find_function(base_tree, start_line)
+
+
+class FindFunction(ast.NodeVisitor):
+    def __init__(self, line):
+        self.line = line
+        self.node: ast.FunctionDef | None = None
+
+    def visit_FunctionDef(self, node: ast.FunctionDef):
+        if node.lineno == self.line or (
+            node.decorator_list and (node.decorator_list[-1].end_lineno <= self.line <= node.lineno)
+        ):
+            self.node = node
+        else:
+            self.generic_visit(node)
+
+
+def find_function(tree: ast.Module, line: int):
+    visitor = FindFunction(line)
+    visitor.visit(tree)
+    return visitor.node
+
+
+class ScanWrites(ast.NodeVisitor):
+    def __init__(self):
+        self.writes = []
+
+    def visit_Name(self, node):
+        if isinstance(node.ctx, ast.Store | ast.Delete):
+            self.writes.append(node.id)
+
+
+def scan_writes(node: ast.AST) -> set[str]:
+    visitor = ScanWrites()
+    visitor.visit(node)
+    return set(visitor.writes)