sonolus.py 0.10.7__py3-none-any.whl → 0.12.5__py3-none-any.whl

sonolus/script/archetype.py

@@ -2,7 +2,7 @@ from __future__ import annotations
 
 import inspect
 from abc import abstractmethod
-from collections.abc import Callable
+from collections.abc import Callable, Sequence
 from dataclasses import dataclass
 from enum import Enum, StrEnum
 from types import FunctionType
@@ -11,7 +11,9 @@ from typing import Annotated, Any, ClassVar, Self, TypedDict, get_origin
 from sonolus.backend.ir import IRConst, IRExpr, IRInstr, IRPureInstr, IRStmt
 from sonolus.backend.mode import Mode
 from sonolus.backend.ops import Op
+from sonolus.backend.visitor import compile_and_call
 from sonolus.script.bucket import Bucket, Judgment
+from sonolus.script.debug import runtime_checks_enabled, static_error
 from sonolus.script.internal.callbacks import PLAY_CALLBACKS, PREVIEW_CALLBACKS, WATCH_ARCHETYPE_CALLBACKS, CallbackInfo
 from sonolus.script.internal.context import ctx
 from sonolus.script.internal.descriptor import SonolusDescriptor
@@ -207,7 +209,7 @@ class _IsScoredDescriptor(SonolusDescriptor):
         if instance is None:
             return self.value
         elif ctx():
-            return ctx().mode_state.is_scored_by_archetype_id[instance.id]
+            return ctx().mode_state.is_scored_by_archetype_id.get_unchecked(instance.id)
         else:
             return self.value
 
@@ -221,7 +223,7 @@ class _IdDescriptor(SonolusDescriptor):
             raise RuntimeError("Archetype id is only available during compilation")
         if instance is None:
             result = ctx().mode_state.archetypes.get(owner)
-            if result is None:
+            if result is None or owner in ctx().mode_state.compile_time_only_archetypes:
                 raise RuntimeError("Archetype is not registered")
             return result
         else:
@@ -237,7 +239,7 @@ class _KeyDescriptor(SonolusDescriptor):
 
     def __get__(self, instance, owner):
         if instance is not None and ctx():
-            return ctx().mode_state.keys_by_archetype_id[instance.id]
+            return ctx().mode_state.keys_by_archetype_id.get_unchecked(instance.id)
         else:
             return self.value
 
@@ -427,7 +429,7 @@ class _BaseArchetype:
 
     _imported_keys_: ClassVar[dict[str, int]]
     _exported_keys_: ClassVar[dict[str, int]]
-    _callbacks_: ClassVar[list[Callable]]
+    _callbacks_: ClassVar[dict[str, Callable]]
     _data_constructor_signature_: ClassVar[inspect.Signature]
     _spawn_signature_: ClassVar[inspect.Signature]
 
@@ -491,17 +493,67 @@ class _BaseArchetype:
 
     @classmethod
     @meta_fn
-    def at(cls, index: int) -> Self:
+    def _compile_time_id(cls):
+        if not ctx():
+            raise RuntimeError("Archetype id is only available during compilation")
+        result = ctx().mode_state.archetypes.get(cls)
+        if result is None:
+            raise RuntimeError("Archetype is not registered")
+        return result
+
+    @classmethod
+    @meta_fn
+    def at(cls, index: int, check: bool = True) -> Self:
+        """Access the entity of this archetype at the given index.
+
+        Args:
+            index: The index of the entity to reference.
+            check: If true, raises an error if the entity at the index is not of this archetype or of a subclass of
+                   this archetype. If false, no validation is performed.
+
+        Returns:
+            The entity at the given index.
+        """
+        if not ctx():
+            raise RuntimeError("Archetype.at is only available during compilation")
+        compile_and_call(cls._check_is_at, index, check=check)
         result = cls._new()
         result._data_ = _ArchetypeReferenceData(index=Num._accept_(index))
         return result
 
     @classmethod
+    def is_at(cls, index: int, strict: bool = False) -> bool:
+        """Return whether the entity at the given index is of this archetype.
+
+        Args:
+            index: The index of the entity to check.
+            strict: If true, only returns true if the entity is exactly of this archetype. If false, also returns true
+                    if the entity is of a subclass of this archetype.
+
+        Returns:
+            Whether the entity at the given index is of this archetype.
+        """
+        if strict:
+            return index >= 0 and cls._compile_time_id() == entity_info_at(index).archetype_id
+        else:
+            mro_ids = cls._get_mro_id_array(entity_info_at(index).archetype_id)
+            return index >= 0 and cls._compile_time_id() in mro_ids
+
+    @classmethod
+    def _check_is_at(cls, index: int, check: bool):
+        if not check or not runtime_checks_enabled():
+            return
+        assert index >= 0, "Entity index must be non-negative"
+        assert cls._compile_time_id() in cls._get_mro_id_array(entity_info_at(index).archetype_id), (
+            "Entity at index is not of the expected archetype"
+        )
+
+    @staticmethod
     @meta_fn
-    def is_at(cls, index: int) -> bool:
+    def _get_mro_id_array(archetype_id: int) -> Sequence[int]:
         if not ctx():
-            raise RuntimeError("is_at is only available during compilation")
-        return entity_info_at(index).archetype_id == cls.id
+            raise RuntimeError("Archetype._get_mro_id_array is only available during compilation")
+        return ctx().get_archetype_mro_id_array(archetype_id)
 
     @classmethod
     @meta_fn
@@ -562,13 +614,16 @@ class _BaseArchetype:
             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__.removeprefix(cls._removable_prefix)
-        cls._callbacks_ = []
+        cls._callbacks_ = {}
         for name in cls._supported_callbacks_:
-            cb = getattr(cls, name)
-            if cb in cls._default_callbacks_:
-                continue
-            cls._callbacks_.append(cb)
+            for mro_entry in cls.mro():
+                if name in mro_entry.__dict__:
+                    cb = mro_entry.__dict__[name]
+                    if cb not in cls._default_callbacks_:
+                        cls._callbacks_[name] = cb
+                        break
         cls._field_init_done = False
+        cls._is_concrete_archetype_ = True
         cls.id = _IdDescriptor()
         cls._key_ = cls.key
         cls.key = _KeyDescriptor(cls.key)
@@ -581,27 +636,39 @@ class _BaseArchetype:
     def _init_fields(cls):
         if cls._field_init_done:
             return
-        cls._field_init_done = True
         for mro_entry in cls.mro()[1:]:
             if hasattr(mro_entry, "_field_init_done"):
                 mro_entry._init_fields()
-        field_specifiers = get_field_specifiers(
-            cls,
-            skip={
-                "id",
-                "key",
-                "name",
-                "life",
-                "is_scored",
-                "_key_",
-                "_is_scored_",
-                "_derived_base_",
-                "_is_derived_",
-                "_default_callbacks_",
-                "_callbacks_",
-                "_field_init_done",
-            },
-        ).items()
+        if sum(issubclass(base, _BaseArchetype) for base in cls.__bases__) > 1:
+            raise TypeError("Multiple inheritance of Archetypes is not supported")
+        mro_from_archetype_parents = set(
+            type("Dummy", tuple(base for base in cls.__bases__ if issubclass(base, _BaseArchetype)), {}).mro()
+        )
+        # Archetype parents would have already initialized relevant fields, so only consider the current class
+        # and mixins that were not already included via an archetype parent
+        mro_excluding_archetype_parents = [entry for entry in cls.mro() if entry not in mro_from_archetype_parents]
+        try:
+            field_specifiers = get_field_specifiers(
+                cls,
+                skip={
+                    "id",
+                    "key",
+                    "name",
+                    "life",
+                    "is_scored",
+                    "_key_",
+                    "_is_scored_",
+                    "_derived_base_",
+                    "_is_derived_",
+                    "_default_callbacks_",
+                    "_callbacks_",
+                    "_field_init_done",
+                    "_is_concrete_archetype_",
+                },
+                included_classes=mro_excluding_archetype_parents,
+            ).items()
+        except Exception as e:
+            raise TypeError(f"Error while processing fields of {cls.__name__}: {e}") from e
         if not hasattr(cls, "_imported_fields_"):
             cls._imported_fields_ = {}
         else:
@@ -643,17 +710,27 @@ class _BaseArchetype:
                             pass
                         else:
                             raise TypeError(
-                                f"Unexpected multiple field annotations for '{name}', "
+                                f"Unexpected multiple annotations for field '{name}' of {cls.__name__}, "
                                 f"expected exactly one of imported, exported, entity_memory, or shared_memory"
                             )
                     else:
                         field_info = metadata
             if field_info is None:
                 raise TypeError(
-                    f"Missing field annotation for '{name}', "
+                    f"Missing annotation for '{name}' of {cls.__name__}, "
                     f"expected exactly one of imported, exported, entity_memory, or shared_memory"
                 )
-            field_type = validate_concrete_type(value.__args__[0])
+            if (
+                name in cls._imported_fields_
+                or name in cls._exported_fields_
+                or name in cls._memory_fields_
+                or name in cls._shared_memory_fields_
+            ):
+                raise ValueError(f"Field '{name}' is already defined in a superclass")
+            try:
+                field_type = validate_concrete_type(value.__args__[0])
+            except Exception as e:
+                raise TypeError(f"Error in field '{name}' of {cls.__name__}: {e}") from e
             match field_info.storage:
                 case _StorageType.IMPORTED:
                     cls._imported_fields_[name] = _ArchetypeField(
@@ -706,6 +783,7 @@ class _BaseArchetype:
             [inspect.Parameter(name, inspect.Parameter.POSITIONAL_OR_KEYWORD) for name in cls._memory_fields_]
         )
         cls._post_init_fields()
+        cls._field_init_done = True
 
     @property
     @abstractmethod
@@ -770,6 +848,18 @@ class _BaseArchetype:
         new_cls = type(name, (cls,), cls_dict)
         return new_cls
 
+    @meta_fn
+    def _delegate(self, name: str, default: int | float | None = None):
+        name = validate_value(name)._as_py_()
+        if hasattr(super(), name):
+            fn = getattr(super(), name)
+            if ctx():
+                return compile_and_call(fn)
+            else:
+                return fn()
+        else:
+            return default
+
 
 class PlayArchetype(_BaseArchetype):
     """Base class for play mode archetypes.
@@ -807,26 +897,28 @@ class PlayArchetype(_BaseArchetype):
 
         Runs first when the level is loaded.
         """
+        self._delegate("preprocess")
 
     def spawn_order(self) -> float:
         """Return the spawn order of the entity.
 
         Runs when the level is loaded after [`preprocess`][sonolus.script.archetype.PlayArchetype.preprocess].
         """
-        return 0.0
+        return self._delegate("spawn_order", default=0.0)
 
     def should_spawn(self) -> bool:
         """Return whether the entity should be spawned.
 
         Runs each frame while the entity is the first entity in the spawn queue.
         """
-        return True
+        return self._delegate("should_spawn", default=True)
 
     def initialize(self):
         """Initialize this entity.
 
         Runs when this entity is spawned.
         """
+        self._delegate("initialize")
 
     def update_sequential(self):
         """Perform non-parallel actions for this frame.
@@ -838,6 +930,7 @@ class PlayArchetype(_BaseArchetype):
         [`update_parallel`][sonolus.script.archetype.PlayArchetype.update_parallel]
         for better performance.
         """
+        self._delegate("update_sequential")
 
     def update_parallel(self):
         """Perform parallel actions for this frame.
@@ -846,18 +939,21 @@ class PlayArchetype(_BaseArchetype):
 
         This is where most gameplay logic should be placed.
         """
+        self._delegate("update_parallel")
 
     def touch(self):
         """Handle user input.
 
         Runs after [`update_sequential`][sonolus.script.archetype.PlayArchetype.update_sequential] each frame.
         """
+        self._delegate("touch")
 
     def terminate(self):
         """Finalize before despawning.
 
         Runs when the entity is despawned.
         """
+        self._delegate("terminate")
 
     @property
     @meta_fn
@@ -965,20 +1061,22 @@ class WatchArchetype(_BaseArchetype):
 
         Runs first when the level is loaded.
         """
+        self._delegate("preprocess")
 
     def spawn_time(self) -> float:
         """Return the spawn time of the entity."""
-        return 0.0
+        return self._delegate("spawn_time", default=0.0)
 
     def despawn_time(self) -> float:
         """Return the despawn time of the entity."""
-        return 0.0
+        return self._delegate("despawn_time", default=0.0)
 
     def initialize(self):
         """Initialize this entity.
 
         Runs when this entity is spawned.
         """
+        self._delegate("initialize")
 
     def update_sequential(self):
         """Perform non-parallel actions for this frame.
@@ -989,6 +1087,7 @@ class WatchArchetype(_BaseArchetype):
         Other logic should typically be placed in
         [`update_parallel`][sonolus.script.archetype.PlayArchetype.update_parallel] for better performance.
         """
+        self._delegate("update_sequential")
 
     def update_parallel(self):
         """Parallel update callback.
@@ -997,12 +1096,14 @@ class WatchArchetype(_BaseArchetype):
 
         This is where most gameplay logic should be placed.
         """
+        self._delegate("update_parallel")
 
     def terminate(self):
         """Finalize before despawning.
 
         Runs when the entity is despawned.
         """
+        self._delegate("terminate")
 
     @property
     @meta_fn
@@ -1073,12 +1174,14 @@ class PreviewArchetype(_BaseArchetype):
 
         Runs first when the level is loaded.
         """
+        self._delegate("preprocess")
 
     def render(self):
         """Render the entity.
 
         Runs after `preprocess`.
         """
+        self._delegate("render")
 
     @property
     @meta_fn
@@ -1210,6 +1313,8 @@ class EntityRef[A: _BaseArchetype](Record):
 
     Usage:
         ```python
+        ref = EntityRef[MyArchetype](index=123)
+
         class MyArchetype(PlayArchetype):
             ref_1: EntityRef[OtherArchetype] = imported()
             ref_2: EntityRef[Any] = imported()
@@ -1239,11 +1344,28 @@ class EntityRef[A: _BaseArchetype](Record):
         return super().__hash__()
 
     @meta_fn
-    def get(self) -> A:
-        """Get the entity."""
+    def __bool__(self):
+        if ctx():
+            static_error("EntityRef cannot be used in a boolean context. Check index directly instead.")
+        return True
+
+    @meta_fn
+    def get(self, *, check: bool = True) -> A:
+        """Get the entity this reference points to.
+
+        Args:
+            check: If true, raises an error if the referenced entity is not of this archetype or of a subclass of
+                   this archetype. If false, no validation is performed.
+
+        Returns:
+            The entity this reference points to.
+        """
+        assert self.archetype() != Any, (
+            "Cannot get entity of unknown (Any) archetype. Use with_archetype() first or use get_as()."
+        )
         if ref := getattr(self, "_ref_", None):
             return ref
-        return self.archetype().at(self.index)
+        return self.archetype().at(self.index, check=check)
 
     @meta_fn
     def get_as(self, archetype: type[_BaseArchetype]) -> _BaseArchetype:
@@ -1252,9 +1374,20 @@ class EntityRef[A: _BaseArchetype](Record):
             raise TypeError("Using get_as in level data is not supported.")
         return self.with_archetype(archetype).get()
 
-    def archetype_matches(self) -> bool:
-        """Check if entity at the index is precisely of the archetype."""
-        return self.index >= 0 and self.archetype().is_at(self.index)
+    def archetype_matches(self, strict: bool = False) -> bool:
+        """Check if entity at the index is of this archetype.
+
+        Args:
+            strict: If true, only returns true if the entity is exactly of this archetype. If false, also returns true
+                    if the entity is of a subclass of this archetype.
+
+        Returns:
+            Whether the entity at the given index is of this archetype.
+        """
+        assert self.archetype() != Any, (
+            "Cannot use archetype_matches with unknown (Any) archetype. Use with_archetype() first."
+        )
+        return self.archetype().is_at(self.index, strict=strict)
 
     def _to_list_(self, level_refs: dict[Any, str] | None = None) -> list[DataValue | str]:
         ref = getattr(self, "_ref_", None)
@@ -1289,6 +1422,11 @@ class EntityRef[A: _BaseArchetype](Record):
             result._ref_ = value._ref_
         return result
 
+    @classmethod
+    def _validate_parameterized_(cls):
+        if not issubclass(cls.archetype(), _BaseArchetype) and cls.archetype() is not Any:
+            raise TypeError("EntityRef type parameter must be an Archetype or Any")
+
 
 class StandardArchetypeName(StrEnum):
     """Standard archetype names."""

sonolus/script/bucket.py

@@ -75,7 +75,7 @@ class JudgmentWindow(Record):
     @perf_meta_fn
     def __mul__(self, other: float | int) -> JudgmentWindow:
         """Multiply the intervals by a scalar."""
-        return JudgmentWindow._quick_construct(
+        return JudgmentWindow._unchecked(
             perfect=self.perfect * other,
             great=self.great * other,
             good=self.good * other,
@@ -84,7 +84,7 @@ class JudgmentWindow(Record):
     @perf_meta_fn
     def __add__(self, other: float | int) -> JudgmentWindow:
         """Add a scalar to the intervals."""
-        return JudgmentWindow._quick_construct(
+        return JudgmentWindow._unchecked(
             perfect=self.perfect + other,
             great=self.great + other,
             good=self.good + other,

sonolus/script/containers.py

@@ -3,6 +3,7 @@ from __future__ import annotations
 from collections.abc import Callable
 from typing import Any, Protocol, Self
 
+from sonolus.backend.visitor import compile_and_call
 from sonolus.script.archetype import AnyArchetype, EntityRef
 from sonolus.script.array import Array
 from sonolus.script.array_like import ArrayLike, get_positive_index
@@ -441,6 +442,56 @@ class ArraySet[T, Capacity](Record):
         self._values.clear()
 
 
+class FrozenNumSet[Size](Record):
+    _values: Array[Num, Size]
+
+    @classmethod
+    @meta_fn
+    def of(cls, *values: Num) -> Self:
+        if ctx():
+            try:
+                num_values = [Num._accept_(v) for v in values]
+            except TypeError:
+                raise TypeError("Only sets of numeric values are supported") from None
+            if all(v._is_py_() for v in num_values):
+                const_values = [v._as_py_() for v in num_values]
+                arr = Array[Num, len(const_values)]._with_value([Num(v) for v in sorted(const_values)])
+                return cls(arr)
+            else:
+                arr = Array[Num, len(values)](*values)
+                compile_and_call(arr.sort)
+        else:
+            arr = Array[Num, len(values)](*sorted(values))
+        return cls(arr)
+
+    def __len__(self) -> int:
+        return len(self._values)
+
+    def __contains__(self, value: Num) -> bool:
+        if len(self) < 8:
+            return value in self._as_tuple()
+        else:
+            left = 0
+            right = len(self) - 1
+            while left <= right:
+                mid = (left + right) // 2
+                mid_value = self._values.get_unchecked(mid)
+                if mid_value == value:
+                    return True
+                elif mid_value < value:
+                    left = mid + 1
+                else:
+                    right = mid - 1
+            return False
+
+    def __iter__(self) -> SonolusIterator[Num]:
+        return self._values.__iter__()
+
+    @meta_fn
+    def _as_tuple(self) -> tuple[Num, ...]:
+        return tuple(self._values.get_unchecked(i) for i in range(Num._accept_(len(self))._as_py_()))
+
+
 class _ArrayMapEntry[K, V](Record):
     key: K
     value: V

sonolus/script/debug.py

@@ -102,6 +102,14 @@ def notify(message: str):
         print(f"[NOTIFY] {message}")
 
 
+@meta_fn
+def runtime_checks_enabled() -> bool:
+    if ctx():
+        return ctx().project_state.runtime_checks != RuntimeChecks.NONE
+    else:
+        return True
+
+
 @meta_fn
 def require(value: int | float | bool, message: str | None = None):
     """Require a condition to be true, or raise an error.
@@ -138,13 +146,22 @@ def require(value: int | float | bool, message: str | None = None):
 
 @meta_fn
 def assert_true(value: int | float | bool, message: str | None = None):
-    if ctx() and ctx().project_state.runtime_checks == RuntimeChecks.NONE:
+    value = validate_value(value)
+    if (
+        ctx()
+        and ctx().project_state.runtime_checks == RuntimeChecks.NONE
+        and not (validate_value(value)._is_py_() and validate_value(value)._as_py_() == 0)
+    ):
+        # Don't do anything if runtime checks are disabled, unless the value is statically known to be false.
+        # `assert False` is a somewhat common pattern to indicate unreachable code, and stripping that out can
+        # cause compilation errors, so we retain those checks.
         return
     require(value, message)
 
 
+@meta_fn
 def assert_false(value: int | float | bool, message: str | None = None):
-    assert_true(not value, message)
+    assert_true(value == 0, message)
 
 
 def static_assert(value: int | float | bool, message: str | None = None):
@@ -211,7 +228,7 @@ def visualize_cfg(
     project_state = ProjectContextState()
     mode_state = ModeContextState(
         mode,
-        {a: i for i, a in enumerate(archetypes)} if archetypes is not None else None,
+        archetypes,
     )
 
     cfg = callback_to_cfg(project_state, mode_state, fn, callback, archetype=archetype)  # type: ignore

sonolus/script/engine.py

@@ -155,6 +155,36 @@ def default_callback() -> Any:
     return 0.0
 
 
+def check_skin(skin: Any):
+    if not hasattr(skin, "_sprites_"):
+        raise ValueError(f"Invalid skin: {skin}. Missing an @skin decorator?")
+
+
+def check_effects(effects: Any):
+    if not hasattr(effects, "_effects_"):
+        raise ValueError(f"Invalid effects: {effects}. Missing an @effects decorator?")
+
+
+def check_particles(particles: Any):
+    if not hasattr(particles, "_particles_"):
+        raise ValueError(f"Invalid particles: {particles}. Missing an @particles decorator?")
+
+
+def check_buckets(buckets: Any):
+    if not hasattr(buckets, "_buckets_"):
+        raise ValueError(f"Invalid buckets: {buckets}. Missing an @buckets decorator?")
+
+
+def check_instructions(instructions: Any):
+    if not hasattr(instructions, "_instructions_"):
+        raise ValueError(f"Invalid instructions: {instructions}. Missing an @instructions decorator?")
+
+
+def check_instruction_icons(instruction_icons: Any):
+    if not hasattr(instruction_icons, "_instruction_icons_"):
+        raise ValueError(f"Invalid instruction icons: {instruction_icons}. Missing an @instruction_icons decorator?")
+
+
 class PlayMode:
     """A play mode definition.
 
@@ -185,6 +215,11 @@ class PlayMode:
             if not issubclass(archetype, PlayArchetype):
                 raise ValueError(f"archetype {archetype} is not a PlayArchetype")
 
+        check_skin(skin)
+        check_effects(effects)
+        check_particles(particles)
+        check_buckets(buckets)
+
 
 class WatchMode:
     """A watch mode definition.
@@ -219,6 +254,11 @@ class WatchMode:
             if not issubclass(archetype, WatchArchetype):
                 raise ValueError(f"archetype {archetype} is not a WatchArchetype")
 
+        check_skin(skin)
+        check_effects(effects)
+        check_particles(particles)
+        check_buckets(buckets)
+
 
 class PreviewMode:
     """A preview mode definition.
@@ -241,6 +281,8 @@ class PreviewMode:
             if not issubclass(archetype, PreviewArchetype):
                 raise ValueError(f"archetype {archetype} is not a PreviewArchetype")
 
+        check_skin(skin)
+
 
 class TutorialMode:
     """A tutorial mode definition.
@@ -277,6 +319,12 @@ class TutorialMode:
         self.navigate = navigate
         self.update = update
 
+        check_skin(skin)
+        check_effects(effects)
+        check_particles(particles)
+        check_instructions(instructions)
+        check_instruction_icons(instruction_icons)
+
 
 def empty_play_mode() -> PlayMode:
     """Create an empty play mode."""