sonolus.py 0.3.1__py3-none-any.whl → 0.3.3__py3-none-any.whl

sonolus/backend/finalize.py

@@ -1,3 +1,5 @@
+from math import isfinite, isinf, isnan
+
 from sonolus.backend.ir import IRConst, IRGet, IRInstr, IRPureInstr, IRSet
 from sonolus.backend.node import ConstantNode, EngineNode, FunctionNode
 from sonolus.backend.ops import Op
@@ -54,10 +56,20 @@ def cfg_to_engine_node(entry: BasicBlock):
 
 def ir_to_engine_node(stmt) -> EngineNode:
     match stmt:
-        case int() | float():
-            return ConstantNode(value=float(stmt))
-        case IRConst(value=int(value) | float(value)):
-            return ConstantNode(value=value)
+        case int(value) | float(value) | IRConst(value=int(value) | float(value)):
+            value = float(value)
+            if value.is_integer():
+                return ConstantNode(value=int(value))
+            elif isfinite(value):
+                return ConstantNode(value=value)
+            elif isinf(value):
+                # Read values from ROM
+                return FunctionNode(Op.Get, args=[ConstantNode(value=3000), ConstantNode(value=1 if value > 0 else 2)])
+            elif isnan(value):
+                # Read value from ROM
+                return FunctionNode(Op.Get, args=[ConstantNode(value=3000), ConstantNode(value=0)])
+            else:
+                raise ValueError(f"Invalid constant 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):

sonolus/backend/node.py

@@ -1,26 +1,34 @@
 import textwrap
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 
 from sonolus.backend.ops import Op
 
 type EngineNode = ConstantNode | FunctionNode
 
 
-@dataclass
+@dataclass(slots=True)
 class ConstantNode:
     value: float
+    _hash: int = field(init=False, repr=False)
+
+    def __post_init__(self):
+        self._hash = hash(self.value)
 
     def __hash__(self):
         return hash(self.value)
 
 
-@dataclass
+@dataclass(slots=True)
 class FunctionNode:
     func: Op
     args: list[EngineNode]
+    _hash: int = field(init=False, repr=False)
+
+    def __post_init__(self):
+        self._hash = hash((self.func, tuple(self.args)))
 
     def __hash__(self):
-        return hash((self.func, tuple(self.args)))
+        return self._hash
 
 
 def format_engine_node(node: EngineNode) -> str:
@@ -34,7 +42,7 @@ def format_engine_node(node: EngineNode) -> str:
                 return f"{node.func.name}({format_engine_node(node.args[0])})"
             case _:
                 return f"{node.func.name}(\n{
-                    textwrap.indent("\n".join(format_engine_node(arg) for arg in node.args), "  ")
+                    textwrap.indent('\n'.join(format_engine_node(arg) for arg in node.args), '  ')
                 }\n)"
     else:
         raise ValueError(f"Invalid engine node: {node}")

sonolus/backend/optimize/allocate.py

@@ -65,7 +65,8 @@ class Allocate(CompilerPass):
     def run(self, entry: BasicBlock):
         mapping = self.get_mapping(entry)
         for block in traverse_cfg_preorder(entry):
-            block.statements = [self.update_stmt(statement, mapping) for statement in block.statements]
+            updated_statements = [self.update_stmt(statement, mapping) for statement in block.statements]
+            block.statements = [stmt for stmt in updated_statements if stmt is not None]
             block.test = self.update_stmt(block.test, mapping)
         return entry
 
@@ -80,7 +81,19 @@ class Allocate(CompilerPass):
             case IRGet(place=place):
                 return IRGet(place=self.update_stmt(place, mapping))
             case IRSet(place=place, value=value):
-                return IRSet(place=self.update_stmt(place, mapping), value=self.update_stmt(value, mapping))
+                # Do some dead code elimination here which is pretty much free since we already have liveness analysis,
+                # and prevents an error from the dead block place being missing from the mapping.
+                live = get_live(stmt)
+                is_live = not (
+                    (isinstance(place, BlockPlace) and isinstance(place.block, TempBlock) and place.block not in live)
+                    or (isinstance(value, IRGet) and place == value.place)
+                )
+                if is_live:
+                    return IRSet(place=self.update_stmt(place, mapping), value=self.update_stmt(value, mapping))
+                elif isinstance(value, IRInstr) and value.op.side_effects:
+                    return self.update_stmt(value, mapping)
+                else:
+                    return None
             case BlockPlace(block=block, index=index, offset=offset):
                 if isinstance(block, TempBlock):
                     if block.size == 0:
@@ -119,8 +132,32 @@ class Allocate(CompilerPass):
     def get_interference(self, entry: BasicBlock) -> dict[TempBlock, set[TempBlock]]:
         result = {}
         for block in traverse_cfg_preorder(entry):
-            for stmt in [*block.statements, block.test]:
+            for stmt in block.statements:
+                if not isinstance(stmt, IRSet):
+                    continue
                 live = {p for p in get_live(stmt) if isinstance(p, TempBlock) and p.size > 0}
                 for place in live:
-                    result.setdefault(place, set()).update(live - {place})
+                    if place not in result:
+                        result[place] = set(live)
+                    else:
+                        result[place].update(live)
         return result
+
+
+class AllocateFast(Allocate):
+    """A bit faster than Allocate but a bit less optimal."""
+
+    def get_mapping(self, entry: BasicBlock) -> dict[TempBlock, int]:
+        interference = self.get_interference(entry)
+        offsets: dict[TempBlock, int] = dict.fromkeys(interference, 0)
+        end_offsets: dict[TempBlock, int] = dict.fromkeys(interference, 0)
+
+        for block, others in interference.items():
+            size = block.size
+            offset = max((end_offsets[other] for other in others), default=0)
+            if offset + size > TEMP_SIZE:
+                raise ValueError("Temporary memory limit exceeded")
+            offsets[block] = offset
+            end_offsets[block] = offset + size
+
+        return offsets

sonolus/backend/optimize/flow.py

@@ -91,12 +91,29 @@ def cfg_to_mermaid(entry: BasicBlock):
 
     lines = ["Entry([Entry]) --> 0"]
     for block, index in block_indexes.items():
-        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]))}]")
+        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:
@@ -114,7 +131,7 @@ def cfg_to_mermaid(entry: BasicBlock):
                 lines.append(f"{index} --> {index}_")
                 for cond, target in tgt.items():
                     lines.append(
-                        f"{index}_ --> |{pre(fmt([cond if cond is not None else "default"]))}| {block_indexes[target]}"
+                        f"{index}_ --> |{pre(fmt([cond if cond is not None else 'default']))}| {block_indexes[target]}"
                     )
     lines.append("Exit([Exit])")
 

sonolus/backend/optimize/optimize.py

@@ -1,4 +1,4 @@
-from sonolus.backend.optimize.allocate import Allocate, AllocateBasic
+from sonolus.backend.optimize.allocate import Allocate, AllocateBasic, AllocateFast
 from sonolus.backend.optimize.constant_evaluation import SparseConditionalConstantPropagation
 from sonolus.backend.optimize.copy_coalesce import CopyCoalesce
 from sonolus.backend.optimize.dead_code import (
@@ -6,9 +6,7 @@ from sonolus.backend.optimize.dead_code import (
     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
 
@@ -21,9 +19,8 @@ MINIMAL_PASSES = (
 FAST_PASSES = (
     CoalesceFlow(),
     UnreachableCodeElimination(),
-    AdvancedDeadCodeElimination(),
+    AllocateFast(),  # Does dead code elimination too, so no need for a separate pass
     CoalesceFlow(),
-    Allocate(),
 )
 
 STANDARD_PASSES = (
@@ -46,7 +43,3 @@ STANDARD_PASSES = (
     NormalizeSwitch(),
     Allocate(),
 )
-
-
-def optimize_and_allocate(cfg: BasicBlock):
-    return run_passes(cfg, STANDARD_PASSES)

sonolus/backend/utils.py

@@ -11,10 +11,15 @@ 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"))
+    base_tree = get_tree_from_file(source_file)
     return source_file, find_function(base_tree, start_line)
 
 
+@cache
+def get_tree_from_file(file: str | Path) -> ast.Module:
+    return ast.parse(Path(file).read_text(encoding="utf-8"))
+
+
 class FindFunction(ast.NodeVisitor):
     def __init__(self, line):
         self.line = line

sonolus/backend/visitor.py

@@ -4,6 +4,7 @@ import builtins
 import functools
 import inspect
 from collections.abc import Callable, Sequence
+from inspect import ismethod
 from types import FunctionType, MethodType, MethodWrapperType
 from typing import Any, Never, Self
 
@@ -54,10 +55,22 @@ def eval_fn(fn: Callable, /, *args, **kwargs):
     source_file, node = get_function(fn)
     bound_args = inspect.signature(fn).bind(*args, **kwargs)
     bound_args.apply_defaults()
+    if ismethod(fn):
+        code = fn.__func__.__code__
+        closure = fn.__func__.__closure__
+    else:
+        code = fn.__code__
+        closure = fn.__closure__
+    if closure is None:
+        nonlocal_vars = {}
+    else:
+        nonlocal_vars = {
+            var: cell.cell_contents for var, cell in zip(code.co_freevars, closure, strict=True) if cell is not None
+        }
     global_vars = {
         **builtins.__dict__,
         **fn.__globals__,
-        **inspect.getclosurevars(fn).nonlocals,
+        **nonlocal_vars,
     }
     return Visitor(source_file, bound_args, global_vars).run(node)
 
@@ -346,7 +359,7 @@ class Visitor(ast.NodeVisitor):
         self.loop_head_ctxs.append(header_ctx)
         self.break_ctxs.append([])
         set_ctx(header_ctx)
-        has_next = self.ensure_boolean_num(self.handle_call(node, iterator.has_next))
+        has_next = self.convert_to_boolean_num(node, self.handle_call(node, iterator.has_next))
         if has_next._is_py_() and not has_next._as_py_():
             # The loop will never run, continue after evaluating the condition
             self.loop_head_ctxs.pop()
@@ -387,7 +400,7 @@ class Visitor(ast.NodeVisitor):
         self.loop_head_ctxs.append(header_ctx)
         self.break_ctxs.append([])
         set_ctx(header_ctx)
-        test = self.ensure_boolean_num(self.visit(node.test))
+        test = self.convert_to_boolean_num(node.test, self.visit(node.test))
         if test._is_py_():
             if test._as_py_():
                 # The loop will run until a break / return
@@ -441,7 +454,7 @@ class Visitor(ast.NodeVisitor):
         set_ctx(after_ctx)
 
     def visit_If(self, node):
-        test = self.ensure_boolean_num(self.visit(node.test))
+        test = self.convert_to_boolean_num(node.test, self.visit(node.test))
 
         if test._is_py_():
             if test._as_py_():
@@ -494,7 +507,9 @@ class Visitor(ast.NodeVisitor):
                 set_ctx(false_ctx)
                 continue
             set_ctx(true_ctx)
-            guard = self.ensure_boolean_num(self.visit(case.guard)) if case.guard else validate_value(True)
+            guard = (
+                self.convert_to_boolean_num(case.guard, self.visit(case.guard)) if case.guard else validate_value(True)
+            )
             if guard._is_py_():
                 if guard._as_py_():
                     for stmt in case.body:
@@ -531,7 +546,7 @@ class Visitor(ast.NodeVisitor):
         match pattern:
             case ast.MatchValue(value=value):
                 value = self.visit(value)
-                test = self.ensure_boolean_num(validate_value(subject == value))
+                test = self.convert_to_boolean_num(pattern, validate_value(subject == value))
                 if test._is_py_():
                     if test._as_py_():
                         return ctx(), ctx().into_dead()
@@ -561,7 +576,7 @@ class Visitor(ast.NodeVisitor):
                 target_len = len(patterns)
                 if not (isinstance(subject, Sequence | TupleImpl)):
                     return ctx().into_dead(), ctx()
-                length_test = self.ensure_boolean_num(validate_value(_len(subject) == target_len))
+                length_test = self.convert_to_boolean_num(pattern, validate_value(_len(subject) == target_len))
                 ctx_init = ctx()
                 if not length_test._is_py_():
                     ctx_init.test = length_test.ir()
@@ -726,8 +741,12 @@ class Visitor(ast.NodeVisitor):
     def visit_UnaryOp(self, node):
         operand = self.visit(node.operand)
         if isinstance(node.op, ast.Not):
-            return self.ensure_boolean_num(operand).not_()
+            return self.convert_to_boolean_num(node, operand).not_()
         op = unary_ops[type(node.op)]
+        if operand._is_py_():
+            operand_py = operand._as_py_()
+            if isinstance(operand_py, type) and hasattr(type(operand_py), op):
+                return self.handle_call(node, getattr(type(operand_py), op), operand_py)
         if hasattr(operand, op):
             return self.handle_call(node, getattr(operand, op))
         raise TypeError(f"bad operand type for unary {op_to_symbol[type(node.op)]}: '{type(operand).__name__}'")
@@ -751,7 +770,7 @@ class Visitor(ast.NodeVisitor):
         return validate_value(fn)
 
     def visit_IfExp(self, node):
-        test = self.ensure_boolean_num(self.visit(node.test))
+        test = self.convert_to_boolean_num(node.test, self.visit(node.test))
 
         if test._is_py_():
             if test._as_py_():
@@ -831,10 +850,11 @@ class Visitor(ast.NodeVisitor):
             ):
                 result = self.handle_call(node, getattr(r_val, rcomp_ops[type(op)]), l_val)
             if result is None or self.is_not_implemented(result):
-                if type(op) is ast.Eq:
-                    result = Num._accept_(l_val is r_val)
-                elif type(op) is ast.NotEq:
-                    result = Num._accept_(l_val is not r_val)
+                # Can't defer to the default object.__eq__ or similar since reference equality (is) is not reliable
+                if type(op) is ast.Eq and type(l_val) is not type(r_val):
+                    return Num._accept_(False)
+                elif type(op) is ast.NotEq and type(l_val) is not type(r_val):
+                    return Num._accept_(True)
                 else:
                     raise TypeError(
                         f"'{op_to_symbol[type(op)]}' not supported between instances of '{type(l_val).__name__}' and "
@@ -1058,7 +1078,7 @@ class Visitor(ast.NodeVisitor):
 
     def handle_call[**P, R](
         self, node: ast.stmt | ast.expr, fn: Callable[P, R], /, *args: P.args, **kwargs: P.kwargs
-    ) -> R:
+    ) -> R | Value:
         """Handles a call to the given callable."""
         self.active_ctx = ctx()
         if (
@@ -1109,6 +1129,20 @@ class Visitor(ast.NodeVisitor):
             raise TypeError(f"Invalid type where a bool (Num) was expected: {type(value).__name__}")
         return value
 
+    def convert_to_boolean_num(self, node, value: Value) -> Num:
+        if _is_num(value):
+            return value
+        if hasattr(type(value), "__bool__"):
+            return self.ensure_boolean_num(self.handle_call(node, type(value).__bool__, validate_value(value)))
+        if hasattr(type(value), "__len__"):
+            length = self.handle_call(node, type(value).__len__, validate_value(value))
+            if not _is_num(length):
+                raise TypeError(f"Invalid type for __len__: {type(length).__name__}")
+            if length._is_py_():
+                return Num._accept_(length._as_py_() > 0)
+            return length > Num._accept_(0)
+        raise TypeError(f"Converting {type(value).__name__} to bool is not supported")
+
     def arguments_to_signature(self, arguments: ast.arguments) -> inspect.Signature:
         parameters: list[inspect.Parameter] = []
         pos_only_count = len(arguments.posonlyargs)
@@ -1177,18 +1211,33 @@ class Visitor(ast.NodeVisitor):
         self, node: ast.stmt | ast.expr, fn: Callable[P, R], /, *args: P.args, **kwargs: P.kwargs
     ) -> R:
         """Executes the given function at the given node for a better traceback."""
+        location_args = {
+            "lineno": node.lineno,
+            "col_offset": node.col_offset,
+            "end_lineno": node.end_lineno,
+            "end_col_offset": node.end_col_offset,
+        }
+
         expr = ast.Expression(
             body=ast.Call(
-                func=ast.Name(id="fn", ctx=ast.Load()),
-                args=[ast.Starred(value=ast.Name(id="args", ctx=ast.Load()), ctx=ast.Load())],
-                keywords=[ast.keyword(value=ast.Name(id="kwargs", ctx=ast.Load()), arg=None)],
-                lineno=node.lineno,
-                col_offset=node.col_offset,
-                end_lineno=node.end_lineno,
-                end_col_offset=node.end_col_offset,
-            ),
+                func=ast.Name(id="fn", ctx=ast.Load(), **location_args),
+                args=[
+                    ast.Starred(
+                        value=ast.Name(id="args", ctx=ast.Load(), **location_args),
+                        ctx=ast.Load(),
+                        **location_args,
+                    )
+                ],
+                keywords=[
+                    ast.keyword(
+                        value=ast.Name(id="kwargs", ctx=ast.Load(), **location_args),
+                        arg=None,
+                        **location_args,
+                    )
+                ],
+                **location_args,
+            )
         )
-        expr = ast.fix_missing_locations(expr)
         return eval(
             compile(expr, filename=self.source_file, mode="eval"),
             {"fn": fn, "args": args, "kwargs": kwargs, "_filter_traceback_": True},

sonolus/build/cli.py

@@ -11,7 +11,7 @@ from pathlib import Path
 from time import perf_counter
 
 from sonolus.backend.optimize.optimize import FAST_PASSES, MINIMAL_PASSES, STANDARD_PASSES
-from sonolus.build.engine import package_engine
+from sonolus.build.engine import no_gil, package_engine
 from sonolus.build.level import package_level_data
 from sonolus.build.project import build_project_to_collection, get_project_schema
 from sonolus.script.project import BuildConfig, Project
@@ -215,6 +215,11 @@ def main():
         else:
             parser.error("Module argument is required when multiple or no modules are found")
 
+    if no_gil():
+        print("Multithreading is enabled")
+    if hasattr(sys, "_jit") and sys._jit.is_enabled():
+        print("Python JIT is enabled")
+
     project = import_project(args.module)
     if project is None:
         sys.exit(1)

sonolus/build/engine.py

@@ -32,7 +32,7 @@ from sonolus.script.project import BuildConfig
 from sonolus.script.sprite import Skin
 from sonolus.script.ui import UiConfig
 
-type JsonValue = None | bool | int | float | str | list[JsonValue] | dict[str, JsonValue]
+type JsonValue = bool | int | float | str | list[JsonValue] | dict[str, JsonValue] | None
 
 
 @dataclass