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

sonolus/script/debug.py

@@ -1,15 +1,25 @@
 from collections.abc import Callable, Sequence
 from contextvars import ContextVar
-from typing import Any, Never
+from typing import Any, Literal, Never
 
 from sonolus.backend.mode import Mode
 from sonolus.backend.ops import Op
+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 cfg_to_mermaid
+from sonolus.backend.optimize.inlining import InlineVars
 from sonolus.backend.optimize.passes import CompilerPass, run_passes
-from sonolus.backend.optimize.simplify import CoalesceFlow
-from sonolus.script.internal.context import GlobalContextState, ctx, set_ctx
+from sonolus.backend.optimize.simplify import CoalesceFlow, NormalizeSwitch, RewriteToSwitch
+from sonolus.backend.optimize.ssa import FromSSA, ToSSA
+from sonolus.script.internal.context import GlobalContextState, ReadOnlyMemory, ctx, set_ctx
 from sonolus.script.internal.impl import meta_fn, validate_value
 from sonolus.script.internal.native import native_function
+from sonolus.script.internal.simulation_context import SimulationContext
 from sonolus.script.num import Num
 
 debug_log_callback = ContextVar[Callable[[Num], None]]("debug_log_callback")
@@ -93,12 +103,60 @@ def terminate():
         raise RuntimeError("Terminated")
 
 
-def visualize_cfg(fn: Callable[[], Any], passes: Sequence[CompilerPass] | None = None) -> str:
+def visualize_cfg(
+    fn: Callable[[], Any] | Callable[[type], Any],
+    /,
+    *,
+    mode: Mode = Mode.PLAY,
+    archetype: type | None = None,
+    archetypes: list[type] | None,
+    passes: Sequence[CompilerPass] | Literal["minimal", "basic", "standard"] = "basic",
+) -> str:
     from sonolus.build.compile import callback_to_cfg
 
-    if passes is None:
-        passes = [CoalesceFlow()]
-
-    cfg = callback_to_cfg(GlobalContextState(Mode.PLAY), fn, "")
+    match passes:
+        case "minimal":
+            passes = [
+                CoalesceFlow(),
+            ]
+        case "basic":
+            passes = [
+                CoalesceFlow(),
+                UnreachableCodeElimination(),
+                AdvancedDeadCodeElimination(),
+                CoalesceFlow(),
+            ]
+        case "standard":
+            passes = [
+                CoalesceFlow(),
+                UnreachableCodeElimination(),
+                DeadCodeElimination(),
+                ToSSA(),
+                SparseConditionalConstantPropagation(),
+                UnreachableCodeElimination(),
+                DeadCodeElimination(),
+                CoalesceFlow(),
+                InlineVars(),
+                DeadCodeElimination(),
+                RewriteToSwitch(),
+                FromSSA(),
+                CoalesceFlow(),
+                CopyCoalesce(),
+                AdvancedDeadCodeElimination(),
+                CoalesceFlow(),
+                NormalizeSwitch(),
+            ]
+
+    global_state = GlobalContextState(
+        mode,
+        {a: i for i, a in enumerate(archetypes)} if archetypes is not None else None,
+        ReadOnlyMemory(),
+    )
+
+    cfg = callback_to_cfg(global_state, fn, "", archetype=archetype)
     cfg = run_passes(cfg, passes)
     return cfg_to_mermaid(cfg)
+
+
+def simulation_context() -> SimulationContext:
+    return SimulationContext()

sonolus/script/globals.py

@@ -5,6 +5,7 @@ from sonolus.backend.blocks import Block, PlayBlock, PreviewBlock, TutorialBlock
 from sonolus.backend.mode import Mode
 from sonolus.script.internal.descriptor import SonolusDescriptor
 from sonolus.script.internal.generic import validate_concrete_type
+from sonolus.script.internal.simulation_context import sim_ctx
 from sonolus.script.internal.value import Value
 
 
@@ -27,6 +28,11 @@ class _GlobalField(SonolusDescriptor):
         if instance is None:
             return self
 
+        if sim_ctx():
+            from sonolus.script.values import zeros
+
+            return sim_ctx().get_or_put_value((instance, self), lambda: zeros(self.type))
+
         from sonolus.script.internal.context import ctx
 
         info = owner._global_info_
@@ -38,6 +44,12 @@ class _GlobalField(SonolusDescriptor):
     def __set__(self, instance, value):
         from sonolus.script.internal.context import ctx
 
+        if sim_ctx():
+            from sonolus.script.values import zeros
+
+            sim_ctx().set_or_put_value((instance, self), lambda: zeros(self.type), value)
+            return
+
         info = instance._global_info_
         if not ctx():
             raise RuntimeError("Global field access outside of compilation")
@@ -64,6 +76,11 @@ class _GlobalPlaceholder:
         base = ctx().get_global_base(self)
         return self.type._from_place_(base)
 
+    def _get_sim_replacement_(self):
+        from sonolus.script.values import zeros
+
+        return sim_ctx().get_or_put_value(self, lambda: zeros(self.type))
+
 
 def _create_global(cls: type, blocks: dict[Mode, Block], offset: int | None):
     if issubclass(cls, Value):

sonolus/script/internal/builtin_impls.py

@@ -1,4 +1,3 @@
-from abc import ABC
 from collections.abc import Iterable
 from typing import overload
 
@@ -32,7 +31,7 @@ def _isinstance(value, type_):
         return isinstance(value, TupleImpl)
     if type_ in {_int, _float, _bool}:
         raise TypeError("Instance check against int, float, or bool is not supported, use Num instead")
-    if not (isinstance(type_, type) and issubclass(type_, Value | ABC)):
+    if not (isinstance(type_, type) and (issubclass(type_, Value) or getattr(type_, "_allow_instance_check_", False))):
         raise TypeError(f"Unsupported type: {type_} for isinstance")
     return validate_value(isinstance(value, type_))
 
@@ -87,7 +86,7 @@ def _zip(*iterables):
     if any(isinstance(iterable, TupleImpl) for iterable in iterables):
         if not all(isinstance(iterable, TupleImpl) for iterable in iterables):
             raise TypeError("Cannot mix tuples with other types in zip")
-        return TupleImpl._accept_(tuple(zip(*(iterable._as_py_() for iterable in iterables), strict=False)))
+        return TupleImpl._accept_(tuple(zip(*(iterable.value for iterable in iterables), strict=False)))
     iterators = [compile_and_call(iterable.__iter__) for iterable in iterables]
     if not all(isinstance(iterator, SonolusIterator) for iterator in iterators):
         raise TypeError("Only subclasses of SonolusIterator are supported as iterators")

sonolus/script/internal/context.py

@@ -19,6 +19,25 @@ _compiler_internal_ = True
 context_var = ContextVar("context_var", default=None)
 
 
+@dataclass(frozen=True)
+class DebugConfig:
+    unchecked_reads: bool
+    unchecked_writes: bool
+
+
+_full_debug_config = DebugConfig(
+    unchecked_reads=True,
+    unchecked_writes=True,
+)
+
+_disabled_debug_config = DebugConfig(
+    unchecked_reads=False,
+    unchecked_writes=False,
+)
+
+debug_var = ContextVar("debug_var", default=_disabled_debug_config)
+
+
 class GlobalContextState:
     archetypes: dict[type, int]
     rom: ReadOnlyMemory
@@ -92,12 +111,16 @@ class Context:
         return self.callback_state.used_names
 
     def check_readable(self, place: BlockPlace):
+        if debug_config().unchecked_reads:
+            return
         if not self.callback:
             return
         if isinstance(place.block, BlockData) and self.callback not in self.blocks(place.block).readable:
             raise RuntimeError(f"Block {place.block} is not readable in {self.callback}")
 
     def check_writable(self, place: BlockPlace):
+        if debug_config().unchecked_writes:
+            return
         if not self.callback:
             return
         if isinstance(place.block, BlockData) and self.callback not in self.blocks(place.block).writable:
@@ -238,6 +261,12 @@ class Context:
             context.outgoing[None] = target
         return target
 
+    def register_archetype(self, type_: type) -> int:
+        with self.global_state.lock:
+            if type_ not in self.global_state.archetypes:
+                self.global_state.archetypes[type_] = len(self.global_state.archetypes)
+            return self.global_state.archetypes[type_]
+
 
 def ctx() -> Context | None:
     return context_var.get()
@@ -283,6 +312,27 @@ class ReadOnlyMemory:
         else:
             return PlayBlock.EngineRom
 
+    def get_value(self, index: int) -> float:
+        with self._lock:
+            if index < 0 or index >= len(self.values):
+                raise IndexError(f"Index {index} out of bounds for ReadOnlyMemory")
+            return self.values[index]
+
+
+@contextmanager
+def enable_debug(config: DebugConfig | None = None):
+    if config is None:
+        config = _full_debug_config
+    token = debug_var.set(config)
+    try:
+        yield
+    finally:
+        debug_var.reset(token)
+
+
+def debug_config() -> DebugConfig:
+    return debug_var.get()
+
 
 @dataclass
 class ValueBinding:

sonolus/script/internal/simulation_context.py

@@ -0,0 +1,131 @@
+import builtins
+import sys
+from collections.abc import Callable
+from types import MethodType, ModuleType
+from typing import Any, ClassVar, Self
+
+from sonolus.script.internal.impl import validate_value
+
+_MISSING = object()
+
+
+class SimulationContext:
+    """Context manager to simulate additional Sonolus runtime features like level memory."""
+
+    _active_context: ClassVar[Self | None] = None
+    additional_replacements: dict[Any, Any]
+    values: dict[Any, Any]
+    _original_values: dict[str, dict[str, Any]]  # module_name: {var_name: original_value}
+    _original_import: Callable[[str, Any, Any, Any, Any], Any] | None
+
+    def __init__(self, *, additional_replacements: dict[Any, Any] | None = None):
+        if SimulationContext._active_context is not None:
+            raise RuntimeError("SimulationContext is already active")
+        SimulationContext._active_context = self
+        self.additional_replacements = additional_replacements or {}
+        self.values = {}
+        self._original_values = {}
+        self._original_import = None
+
+    def get_or_put_value(self, key: Any, factory: Callable[[], Any]) -> Any:
+        if key not in self.values:
+            self.values[key] = validate_value(factory())
+        return self.values[key]._get_()._as_py_()
+
+    def set_or_put_value(self, key: Any, factory: Callable[[], Any], value: Any):
+        if key not in self.values:
+            self.values[key] = validate_value(factory())
+        existing_value = self.values[key]
+        existing_type = type(existing_value)
+        value = existing_type._accept_(validate_value(value))
+        if existing_type._is_value_type_():
+            existing_value._set_(value)
+        else:
+            existing_value._copy_from_(value)
+
+    def _get_replacement(self, value: Any) -> Any:
+        try:
+            if value in self.additional_replacements:
+                return self.additional_replacements[value]
+        except TypeError:
+            pass
+        if hasattr(value, "_get_sim_replacement_") and isinstance(value._get_sim_replacement_, MethodType):
+            return value._get_sim_replacement_()
+        return _MISSING
+
+    def _substitute_module_variables(self, module) -> None:
+        if not isinstance(module, ModuleType):
+            return
+
+        module_name = module.__name__
+
+        if module_name in self._original_values:
+            return
+
+        original_values = {}
+
+        for var_name, var_value in list(module.__dict__.items()):
+            replacement = self._get_replacement(var_value)
+            if replacement is not _MISSING:
+                original_values[var_name] = var_value
+                setattr(module, var_name, replacement)
+
+        if original_values:
+            self._original_values[module_name] = original_values
+
+    def _update_loaded_modules(self) -> None:
+        loaded_modules = list(sys.modules.values())
+
+        for module in loaded_modules:
+            if module is not None:
+                self._substitute_module_variables(module)
+
+    def _create_import_hook(self):
+        original_import = builtins.__import__
+
+        def hooked_import(name, globals=None, locals=None, fromlist=(), level=0):  # noqa: A002
+            module = original_import(name, globals, locals, fromlist, level)
+            if isinstance(module, ModuleType):
+                self._substitute_module_variables(module)
+
+            if fromlist:
+                for item in fromlist:
+                    if hasattr(module, item):
+                        attr = getattr(module, item)
+                        if isinstance(attr, ModuleType):
+                            self._substitute_module_variables(attr)
+
+            return module
+
+        return hooked_import
+
+    def __enter__(self) -> Self:
+        self._update_loaded_modules()
+
+        self._original_import = builtins.__import__
+        builtins.__import__ = self._create_import_hook()
+
+        return self
+
+    def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any):
+        try:
+            if self._original_import is not None:
+                builtins.__import__ = self._original_import
+                self._original_import = None
+
+            for module_name, original_values in self._original_values.items():
+                if module_name in sys.modules:
+                    module = sys.modules[module_name]
+                    if isinstance(module, ModuleType):
+                        for var_name, original_value in original_values.items():
+                            setattr(module, var_name, original_value)
+
+            self._original_values.clear()
+
+        finally:
+            SimulationContext._active_context = None
+
+
+def sim_ctx() -> SimulationContext | None:
+    """Get the current simulation context, or None if not active."""
+    return SimulationContext._active_context

sonolus/script/internal/tuple_impl.py

@@ -30,7 +30,7 @@ class TupleImpl(TransientValue):
         return len(self.value)
 
     def __eq__(self, other):
-        if not isinstance(other, tuple):
+        if not self._is_tuple_impl(other):
             return False
         if len(self) != len(other):
             return False
@@ -40,7 +40,7 @@ class TupleImpl(TransientValue):
         return True
 
     def __ne__(self, other):
-        if not isinstance(other, tuple):
+        if not self._is_tuple_impl(other):
             return True
         if len(self) != len(other):
             return True
@@ -50,33 +50,33 @@ class TupleImpl(TransientValue):
         return False
 
     def __lt__(self, other):
-        if not isinstance(other, tuple):
+        if not self._is_tuple_impl(other):
             return NotImplemented
-        for a, b in zip(self.value, other.value):
+        for a, b in zip(self, other):
             if a != b:
                 return a < b
         return len(self.value) < len(other.value)
 
     def __le__(self, other):
-        if not isinstance(other, tuple):
+        if not self._is_tuple_impl(other):
             return NotImplemented
-        for a, b in zip(self.value, other.value):
+        for a, b in zip(self, other):
             if a != b:
                 return a < b
         return len(self.value) <= len(other.value)
 
     def __gt__(self, other):
-        if not isinstance(other, tuple):
+        if not self._is_tuple_impl(other):
             return NotImplemented
-        for a, b in zip(self.value, other.value):
+        for a, b in zip(self, other):
             if a != b:
                 return a > b
         return len(self.value) > len(other.value)
 
     def __ge__(self, other):
-        if not isinstance(other, tuple):
+        if not self._is_tuple_impl(other):
             return NotImplemented
-        for a, b in zip(self.value, other.value):
+        for a, b in zip(self, other):
             if a != b:
                 return a > b
         return len(self.value) >= len(other.value)
@@ -89,6 +89,11 @@ class TupleImpl(TransientValue):
         other = TupleImpl._accept_(other)
         return TupleImpl._accept_(self.value + other.value)
 
+    @staticmethod
+    @meta_fn
+    def _is_tuple_impl(value: Any) -> bool:
+        return isinstance(value, TupleImpl)
+
     @classmethod
     def _accepts_(cls, value: Any) -> bool:
         return isinstance(value, cls | tuple)

sonolus/script/iterator.py

@@ -1,14 +1,13 @@
 from __future__ import annotations
 
 from abc import abstractmethod
-from collections.abc import Iterator
 from typing import Any
 
 from sonolus.script.internal.impl import meta_fn
 from sonolus.script.record import Record
 
 
-class SonolusIterator[T](Iterator[T]):
+class SonolusIterator[T]:
     """Base class for Sonolus iterators.
 
     This class is used to define custom iterators that can be used in Sonolus.py.
@@ -30,6 +29,8 @@ class SonolusIterator[T](Iterator[T]):
         ```
     """
 
+    _allow_instance_check_ = True
+
     def next(self) -> T:
         result = self.get()
         self.advance()

sonolus/script/num.py

@@ -1,4 +1,3 @@
-# ruff: noqa: N801
 from __future__ import annotations
 
 import operator
@@ -81,12 +80,22 @@ class _Num(Value, metaclass=_NumMeta):
             return value
         return cls(value)
 
+    def _is_rom_constant(self) -> bool:
+        return (
+            ctx()
+            and isinstance(self.data, BlockPlace)
+            and self.data.block == ctx().blocks.EngineRom
+            and isinstance(self.data.index, int)
+        )
+
     def _is_py_(self) -> bool:
-        return isinstance(self.data, float | int | bool)
+        return isinstance(self.data, float | int | bool) or self._is_rom_constant()
 
     def _as_py_(self) -> Any:
         if not self._is_py_():
             raise ValueError("Not a compile time constant Num")
+        if self._is_rom_constant():
+            return ctx().rom.get_value(self.data.index + self.data.offset)
         if self.data.is_integer():
             return int(self.data)
         return self.data
@@ -114,10 +123,11 @@ class _Num(Value, metaclass=_NumMeta):
             return Num(self.data)
 
     def _set_(self, value: Self):
+        value = Num._accept_(value)
         if ctx():
             match self.data:
                 case BackingValue():
-                    ctx().add_statements(self.data.write(value))
+                    self.data.write(value.ir())
                 case BlockPlace():
                     ctx().check_writable(self.data)
                     ctx().add_statements(IRSet(self.data, value.ir()))

sonolus/script/options.py

@@ -5,11 +5,13 @@ from typing import Annotated, Any, NewType, dataclass_transform, get_origin
 from sonolus.backend.mode import Mode
 from sonolus.backend.place import BlockPlace
 from sonolus.script.debug import assert_unreachable
-from sonolus.script.internal.context import ctx
+from sonolus.script.internal.context import ctx, debug_config
 from sonolus.script.internal.descriptor import SonolusDescriptor
 from sonolus.script.internal.generic import validate_concrete_type
 from sonolus.script.internal.introspection import get_field_specifiers
+from sonolus.script.internal.simulation_context import sim_ctx
 from sonolus.script.num import Num
+from sonolus.script.values import copy
 
 
 @dataclass
@@ -186,6 +188,8 @@ class _OptionField(SonolusDescriptor):
         self.index = index
 
     def __get__(self, instance, owner):
+        if sim_ctx():
+            return sim_ctx().get_or_put_value((instance, self), lambda: copy(self.info.default))
         if ctx():
             match ctx().global_state.mode:
                 case Mode.PLAY:
@@ -202,8 +206,27 @@ class _OptionField(SonolusDescriptor):
                 return Num._from_place_(BlockPlace(block, self.index))
             else:
                 return Num._accept_(self.info.default)
+        raise RuntimeError("Options can only be accessed in a context")
 
     def __set__(self, instance, value):
+        if sim_ctx():
+            return sim_ctx().set_or_put_value((instance, self), lambda: copy(self.info.default), value)
+        if ctx() and debug_config().unchecked_writes:
+            match ctx().global_state.mode:
+                case Mode.PLAY:
+                    block = ctx().blocks.LevelOption
+                case Mode.WATCH:
+                    block = ctx().blocks.LevelOption
+                case Mode.PREVIEW:
+                    block = ctx().blocks.PreviewOption
+                case Mode.TUTORIAL:
+                    block = None
+                case _:
+                    assert_unreachable()
+            if block is not None:
+                Num._from_place_(BlockPlace(block, self.index))._set_(Num._accept_(value))
+            else:
+                raise RuntimeError("Options in the current mode cannot be set and use the default value")
         raise AttributeError("Options are read-only")
 
 

sonolus/script/quad.py

@@ -147,6 +147,8 @@ class Quad(Record):
     def contains_point(self, point: Vec2, /) -> bool:
         """Check if the quad contains the given point.
 
+        It is not guaranteed whether points on the edges of the quad are considered inside or outside.
+
         Args:
             point: The point to check.
 

sonolus/script/record.py

@@ -20,8 +20,14 @@ from sonolus.script.internal.value import BackingSource, DataValue, Value
 from sonolus.script.num import Num
 
 
+class RecordMeta(type):
+    @meta_fn
+    def __pos__[T](cls: type[T]) -> T:
+        return cls._zero_()
+
+
 @dataclass_transform(eq_default=True)
-class Record(GenericValue):
+class Record(GenericValue, metaclass=RecordMeta):
     """Base class for user-defined data structures.
 
     Usage:
@@ -38,6 +44,15 @@ class Record(GenericValue):
             field1: T
             field2: U
         ```
+
+        Creating an instance:
+        ```python
+        record = MyRecord(field1=42, field2=True)
+        record_2 = MyGenericRecord[int, int](field1=42, field2=100)
+        record_3 = MyGenericRecord(field1=42, field2=100)  # Type arguments can be inferred
+        record_4 = +MyRecord  # Create a zero-initialized record
+        record_5 = +MyGenericRecord[int, int]
+        ```
     """
 
     _value: dict[str, Value]
@@ -239,12 +254,12 @@ class Record(GenericValue):
 
     def __str__(self):
         return (
-            f"{self.__class__.__name__}({", ".join(f"{field.name}={field.__get__(self)}" for field in self._fields)})"
+            f"{self.__class__.__name__}({', '.join(f'{field.name}={field.__get__(self)}' for field in self._fields)})"
         )
 
     def __repr__(self):
         return (
-            f"{self.__class__.__name__}({", ".join(f"{field.name}={field.__get__(self)!r}" for field in self._fields)})"
+            f"{self.__class__.__name__}({', '.join(f'{field.name}={field.__get__(self)!r}' for field in self._fields)})"
         )
 
     @meta_fn
@@ -268,6 +283,10 @@ class Record(GenericValue):
     def __hash__(self):
         return hash(tuple(field.__get__(self) for field in self._fields))
 
+    def __pos__(self) -> Self:
+        """Return a copy of the record."""
+        return self._copy_()
+
     @classmethod
     @meta_fn
     def type_var_value(cls, var: TypeVar, /) -> Any: