sonolus.py 0.3.0__py3-none-any.whl → 0.3.2__py3-none-any.whl

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)
 
@@ -728,6 +741,10 @@ class Visitor(ast.NodeVisitor):
         if isinstance(node.op, ast.Not):
             return self.ensure_boolean_num(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__}'")
@@ -831,10 +848,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 "
@@ -1177,18 +1195,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

sonolus/script/archetype.py

@@ -185,14 +185,14 @@ class _ArchetypeField(SonolusDescriptor):
 def imported(*, name: str | None = None) -> Any:
     """Declare a field as imported.
 
-    Imported fields may be loaded from the level data.
+    Imported fields may be loaded from the level.
 
     In watch mode, data may also be loaded from a corresponding exported field in play mode.
 
     Imported fields may only be updated in the `preprocess` callback, and are read-only in other callbacks.
 
     Usage:
-        ```
+        ```python
         class MyArchetype(PlayArchetype):
             field: int = imported()
             field_with_explicit_name: int = imported(name="field_name")
@@ -207,10 +207,10 @@ def entity_data() -> Any:
     Entity data is accessible from other entities, but may only be updated in the `preprocess` callback
     and is read-only in other callbacks.
 
-    It functions like `imported`, except that it is not loaded from the level data.
+    It functions like `imported` and shares the same underlying storage, except that it is not loaded from a level.
 
     Usage:
-        ```
+        ```python
         class MyArchetype(PlayArchetype):
             field: int = entity_data()
         ```
@@ -226,7 +226,7 @@ def exported(*, name: str | None = None) -> Any:
     Exported fields are write-only.
 
     Usage:
-        ```
+        ```python
         class MyArchetype(PlayArchetype):
             field: int = exported()
             field_with_explicit_name: int = exported(name="#FIELD")
@@ -238,12 +238,13 @@ def exported(*, name: str | None = None) -> Any:
 def entity_memory() -> Any:
     """Declare a field as entity memory.
 
-    Entity memory is private to the entity and is not accessible from other entities.
+    Entity memory is private to the entity and is not accessible from other entities. It may be read or updated in any
+    callback associated with the entity.
 
     Entity memory fields may also be set when an entity is spawned using the `spawn()` method.
 
     Usage:
-        ```
+        ```python
         class MyArchetype(PlayArchetype):
             field: int = entity_memory()
 
@@ -257,10 +258,11 @@ def shared_memory() -> Any:
 
     Shared memory is accessible from other entities.
 
-    Shared memory may only be updated by sequential callbacks such as `preprocess`, `update_sequential`, and `touch`.
+    Shared memory may be read in any callback, but may only be updated by sequential callbacks
+    (`preprocess`, `update_sequential`, and `touch`).
 
     Usage:
-        ```
+        ```python
         class MyArchetype(PlayArchetype):
             field: int = shared_memory()
         ```
@@ -280,7 +282,7 @@ class StandardImport:
     """Standard import annotations for Archetype fields.
 
     Usage:
-        ```
+        ```python
         class MyArchetype(WatchArchetype):
             judgment: StandardImport.JUDGMENT
         ```
@@ -313,7 +315,7 @@ def callback[T: Callable](*, order: int = 0) -> Callable[[T], T]:
     Callbacks are execute from lowest to highest order. By default, callbacks have an order of 0.
 
     Usage:
-        ```
+        ```python
         class MyArchetype(PlayArchetype):
             @callback(order=1)
             def update_sequential(self):
@@ -360,6 +362,8 @@ class ArchetypeSchema(TypedDict):
 class _BaseArchetype:
     _is_comptime_value_ = True
 
+    _removable_prefix: ClassVar[str] = ""
+
     _supported_callbacks_: ClassVar[dict[str, CallbackInfo]]
     _default_callbacks_: ClassVar[set[Callable]]
 
@@ -440,7 +444,7 @@ class _BaseArchetype:
         """Spawn an entity of this archetype, injecting the given values into entity memory.
 
         Usage:
-            ```
+            ```python
             class MyArchetype(PlayArchetype):
                 field: int = entity_memory()
 
@@ -492,7 +496,7 @@ class _BaseArchetype:
             cls._default_callbacks_ = {getattr(cls, cb_info.py_name) for cb_info in cls._supported_callbacks_.values()}
             return
         if cls.name is None or cls.name in {getattr(mro_entry, "name", None) for mro_entry in cls.mro()[1:]}:
-            cls.name = cls.__name__
+            cls.name = cls.__name__.removeprefix(cls._removable_prefix)
         cls._callbacks_ = []
         for name in cls._supported_callbacks_:
             cb = getattr(cls, name)
@@ -642,7 +646,7 @@ class PlayArchetype(_BaseArchetype):
     """Base class for play mode archetypes.
 
     Usage:
-        ```
+        ```python
         class MyArchetype(PlayArchetype):
             # Set to True if the entity is a note and contributes to combo and score
             # Default is False
@@ -659,6 +663,8 @@ class PlayArchetype(_BaseArchetype):
         ```
     """
 
+    _removable_prefix: ClassVar[str] = "Play"
+
     _supported_callbacks_ = PLAY_CALLBACKS
 
     is_scored: ClassVar[bool] = False
@@ -808,7 +814,7 @@ class WatchArchetype(_BaseArchetype):
     """Base class for watch mode archetypes.
 
     Usage:
-        ```
+        ```python
         class MyArchetype(WatchArchetype):
             imported_field: int = imported()
             entity_memory_field: int = entity_memory()
@@ -820,6 +826,8 @@ class WatchArchetype(_BaseArchetype):
         ```
     """
 
+    _removable_prefix: ClassVar[str] = "Watch"
+
     _supported_callbacks_ = WATCH_ARCHETYPE_CALLBACKS
 
     def preprocess(self):
@@ -923,7 +931,7 @@ class PreviewArchetype(_BaseArchetype):
     """Base class for preview mode archetypes.
 
     Usage:
-        ```
+        ```python
         class MyArchetype(PreviewArchetype):
             imported_field: int = imported()
             entity_memory_field: int = entity_memory()
@@ -935,6 +943,8 @@ class PreviewArchetype(_BaseArchetype):
         ```
     """
 
+    _removable_prefix: ClassVar[str] = "Preview"
+
     _supported_callbacks_ = PREVIEW_CALLBACKS
 
     def preprocess(self):
@@ -1078,7 +1088,7 @@ class EntityRef[A: _BaseArchetype](Record):
     May be used with `Any` to reference an unknown archetype.
 
     Usage:
-        ```
+        ```python
         class MyArchetype(PlayArchetype):
             ref_1: EntityRef[OtherArchetype] = imported()
             ref_2: EntityRef[Any] = imported()

sonolus/script/array.py

@@ -1,4 +1,3 @@
-# ruff: noqa: A005
 from __future__ import annotations
 
 from collections.abc import Iterable
@@ -16,14 +15,21 @@ from sonolus.script.internal.value import BackingSource, DataValue, Value
 from sonolus.script.num import Num
 
 
+class ArrayMeta(type):
+    @meta_fn
+    def __pos__[T](cls: type[T]) -> T:
+        return cls._zero_()
+
+
 @final
-class Array[T, Size](GenericValue, ArrayLike[T]):
+class Array[T, Size](GenericValue, ArrayLike[T], metaclass=ArrayMeta):
     """A fixed size array of values.
 
     Usage:
         ```python
         array_1 = Array(1, 2, 3)
         array_2 = Array[int, 0]()
+        array_3 = +Array[int, 3]  # Create a zero-initialized array
         ```
     """
 
@@ -66,7 +72,7 @@ class Array[T, Size](GenericValue, ArrayLike[T]):
         else:
             return parameterized_cls._with_value([value._copy_() for value in values])
 
-    def __init__(self, *_args: T):
+    def __init__(self, *args: T):
         super().__init__()
 
     @classmethod
@@ -289,10 +295,14 @@ class Array[T, Size](GenericValue, ArrayLike[T]):
         if isinstance(self._value, BlockPlace) or callable(self._value):
             return f"{type(self).__name__}({self._value}...)"
         else:
-            return f"{type(self).__name__}({", ".join(str(self[i]) for i in range(self.size()))})"
+            return f"{type(self).__name__}({', '.join(str(self[i]) for i in range(self.size()))})"
 
     def __repr__(self):
         if isinstance(self._value, BlockPlace) or callable(self._value):
             return f"{type(self).__name__}({self._value}...)"
         else:
-            return f"{type(self).__name__}({", ".join(repr(self[i]) for i in range(self.size()))})"
+            return f"{type(self).__name__}({', '.join(repr(self[i]) for i in range(self.size()))})"
+
+    def __pos__(self) -> Self:
+        """Return a copy of the array."""
+        return self._copy_()

sonolus/script/array_like.py

@@ -1,8 +1,8 @@
 from __future__ import annotations
 
 import random
-from abc import ABC, abstractmethod
-from collections.abc import Callable, Sequence
+from abc import abstractmethod
+from collections.abc import Callable
 from typing import Any
 
 from sonolus.script.iterator import SonolusIterator
@@ -13,7 +13,7 @@ from sonolus.script.values import copy
 # Note: we don't use Range in this file because Range itself inherits from ArrayLike
 
 
-class ArrayLike[T](Sequence, ABC):
+class ArrayLike[T]:
     """Mixin for array-like objects.
 
     Inheritors must implement `__len__`, `__getitem__`, and `__setitem__`.
@@ -32,6 +32,8 @@ class ArrayLike[T](Sequence, ABC):
         ```
     """
 
+    _allow_instance_check_ = True
+
     @abstractmethod
     def __len__(self) -> int:
         """Return the length of the array."""

sonolus/script/containers.py

@@ -10,7 +10,7 @@ from sonolus.script.iterator import SonolusIterator
 from sonolus.script.num import Num
 from sonolus.script.pointer import _deref
 from sonolus.script.record import Record
-from sonolus.script.values import alloc, copy
+from sonolus.script.values import copy, zeros
 
 
 class Box[T](Record):
@@ -104,7 +104,7 @@ class VarArray[T, Capacity](Record, ArrayLike[T]):
         """Create a new empty array."""
         element_type = cls.type_var_value(T)
         capacity = cls.type_var_value(Capacity)
-        return cls(0, alloc(Array[element_type, capacity]))
+        return cls(0, zeros(Array[element_type, capacity]))
 
     def __len__(self) -> int:
         """Return the number of elements in the array."""
@@ -457,7 +457,7 @@ class ArrayMap[K, V, Capacity](Record):
         key_type = cls.type_var_value(K)
         value_type = cls.type_var_value(V)
         capacity = cls.type_var_value(Capacity)
-        return cls(0, alloc(Array[_ArrayMapEntry[key_type, value_type], capacity]))
+        return cls(0, zeros(Array[_ArrayMapEntry[key_type, value_type], capacity]))
 
     def __len__(self) -> int:
         """Return the number of key-value pairs in the map."""