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

sonolus/backend/ops.py

@@ -180,6 +180,12 @@ class Op(StrEnum):
     StackSetPointer = ("StackSetPointer", True, False, False)
     StopLooped = ("StopLooped", True, False, False)
     StopLoopedScheduled = ("StopLoopedScheduled", True, False, False)
+    # Streams are immutable when they're readable, so we can treat read operations as pure.
+    StreamGetNextKey = ("StreamGetNextKey", False, True, False)
+    StreamGetPreviousKey = ("StreamGetPreviousKey", False, True, False)
+    StreamGetValue = ("StreamGetValue", False, True, False)
+    StreamHas = ("StreamHas", False, True, False)
+    StreamSet = ("StreamSet", True, False, False)
     Subtract = ("Subtract", False, True, False)
     Switch = ("Switch", False, True, True)
     SwitchInteger = ("SwitchInteger", False, True, True)

sonolus/script/archetype.py

@@ -8,10 +8,9 @@ from enum import Enum, StrEnum
 from types import FunctionType
 from typing import Annotated, Any, ClassVar, Self, TypedDict, get_origin
 
-from sonolus.backend.ir import IRConst, IRInstr
+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.place import BlockPlace
 from sonolus.script.bucket import Bucket, Judgment
 from sonolus.script.internal.callbacks import PLAY_CALLBACKS, PREVIEW_CALLBACKS, WATCH_ARCHETYPE_CALLBACKS, CallbackInfo
 from sonolus.script.internal.context import ctx
@@ -20,9 +19,9 @@ from sonolus.script.internal.generic import validate_concrete_type
 from sonolus.script.internal.impl import meta_fn, validate_value
 from sonolus.script.internal.introspection import get_field_specifiers
 from sonolus.script.internal.native import native_call
-from sonolus.script.internal.value import Value
+from sonolus.script.internal.value import BackingValue, DataValue, Value
 from sonolus.script.num import Num
-from sonolus.script.pointer import _deref
+from sonolus.script.pointer import _backing_deref, _deref
 from sonolus.script.record import Record
 from sonolus.script.values import zeros
 
@@ -44,6 +43,17 @@ class _ArchetypeFieldInfo:
     storage: _StorageType
 
 
+class _ExportBackingValue(BackingValue):
+    def __init__(self, index: IRExpr):
+        self.index = index
+
+    def read(self) -> IRExpr:
+        raise NotImplementedError("Exported fields are write-only")
+
+    def write(self, value: IRExpr) -> IRStmt:
+        return IRInstr(Op.ExportValue, [self.index, value])
+
+
 class _ArchetypeField(SonolusDescriptor):
     def __init__(self, name: str, data_name: str, storage: _StorageType, offset: int, type_: type[Value]):
         self.name = name
@@ -70,7 +80,20 @@ class _ArchetypeField(SonolusDescriptor):
                     case _ArchetypeLevelData(values=values):
                         result = values[self.name]
             case _StorageType.EXPORTED:
-                raise RuntimeError("Exported fields are write-only")
+                match instance._data_:
+                    case _ArchetypeSelfData():
+
+                        def backing_source(i: IRExpr):
+                            return _ExportBackingValue(IRPureInstr(Op.Add, [i, IRConst(self.offset)]))
+
+                        result = _backing_deref(
+                            backing_source,
+                            self.type,
+                        )
+                    case _ArchetypeReferenceData():
+                        raise RuntimeError("Exported fields of other entities are not accessible")
+                    case _ArchetypeLevelData():
+                        raise RuntimeError("Exported fields are not available in level data")
             case _StorageType.MEMORY:
                 match instance._data_:
                     case _ArchetypeSelfData():
@@ -584,6 +607,7 @@ class _BaseArchetype:
         cls._spawn_signature_ = inspect.Signature(
             [inspect.Parameter(name, inspect.Parameter.POSITIONAL_OR_KEYWORD) for name in cls._memory_fields_]
         )
+        cls._post_init_fields()
 
     @property
     @abstractmethod
@@ -609,6 +633,10 @@ class _BaseArchetype:
             case _:
                 raise RuntimeError("Invalid entity data")
 
+    @classmethod
+    def _post_init_fields(cls):
+        pass
+
 
 class PlayArchetype(_BaseArchetype):
     """Base class for play mode archetypes.
@@ -651,7 +679,7 @@ class PlayArchetype(_BaseArchetype):
     def should_spawn(self) -> bool:
         """Return whether the entity should be spawned.
 
-        Runs when this entity is first in the spawn queue.
+        Runs each frame while the entity is the first entity in the spawn queue.
         """
 
     def initialize(self):
@@ -885,6 +913,11 @@ class WatchArchetype(_BaseArchetype):
             case _:
                 raise RuntimeError("Result is only accessible from the entity itself")
 
+    @classmethod
+    def _post_init_fields(cls):
+        if cls._exported_fields_:
+            raise RuntimeError("Watch archetypes cannot have exported fields")
+
 
 class PreviewArchetype(_BaseArchetype):
     """Base class for preview mode archetypes.
@@ -934,6 +967,11 @@ class PreviewArchetype(_BaseArchetype):
         """The index of this entity."""
         return self._info.index
 
+    @classmethod
+    def _post_init_fields(cls):
+        if cls._exported_fields_:
+            raise RuntimeError("Preview archetypes cannot have exported fields")
+
 
 @meta_fn
 def entity_info_at(index: Num) -> PlayEntityInfo | WatchEntityInfo | PreviewEntityInfo:
@@ -1066,7 +1104,7 @@ class EntityRef[A: _BaseArchetype](Record):
         """Check if entity at the index is precisely of the archetype."""
         return self.index >= 0 and self.archetype().is_at(self.index)
 
-    def _to_list_(self, level_refs: dict[Any, str] | None = None) -> list[float | str | BlockPlace]:
+    def _to_list_(self, level_refs: dict[Any, str] | None = None) -> list[DataValue | str]:
         ref = getattr(self, "_ref_", None)
         if ref is None:
             return [self.index]

sonolus/script/array.py

@@ -12,7 +12,7 @@ from sonolus.script.internal.context import ctx
 from sonolus.script.internal.error import InternalError
 from sonolus.script.internal.generic import GenericValue
 from sonolus.script.internal.impl import meta_fn, validate_value
-from sonolus.script.internal.value import Value
+from sonolus.script.internal.value import BackingSource, DataValue, Value
 from sonolus.script.num import Num
 
 
@@ -27,16 +27,14 @@ class Array[T, Size](GenericValue, ArrayLike[T]):
         ```
     """
 
-    _value: list[T] | BlockPlace
+    _value: list[T] | BlockPlace | BackingSource
 
     @classmethod
-    @meta_fn
     def element_type(cls) -> type[T] | type[Value]:
         """Return the type of elements in this array type."""
         return cls.type_var_value(T)
 
     @classmethod
-    @meta_fn
     def size(cls) -> int:
         """Return the size of this array type.
 
@@ -85,6 +83,10 @@ class Array[T, Size](GenericValue, ArrayLike[T]):
     def _is_value_type_(cls) -> bool:
         return False
 
+    @classmethod
+    def _from_backing_source_(cls, source: BackingSource) -> Self:
+        return cls._with_value(source)
+
     @classmethod
     def _from_place_(cls, place: BlockPlace) -> Self:
         return cls._with_value(place)
@@ -108,11 +110,11 @@ class Array[T, Size](GenericValue, ArrayLike[T]):
         return self
 
     @classmethod
-    def _from_list_(cls, values: Iterable[float | BlockPlace]) -> Self:
+    def _from_list_(cls, values: Iterable[DataValue]) -> Self:
         iterator = iter(values)
         return cls(*(cls.element_type()._from_list_(iterator) for _ in range(cls.size())))
 
-    def _to_list_(self, level_refs: dict[Any, str] | None = None) -> list[float | str | BlockPlace]:
+    def _to_list_(self, level_refs: dict[Any, str] | None = None) -> list[DataValue | str]:
         match self._value:
             case list():
                 return [entry for value in self._value for entry in value._to_list_(level_refs)]
@@ -124,6 +126,8 @@ class Array[T, Size](GenericValue, ArrayLike[T]):
                     ._from_place_(self._value.add_offset(i * self.element_type()._size_()))
                     ._to_list_()
                 ]
+            case backing_source if callable(backing_source):
+                return [backing_source(IRConst(i)) for i in range(self.size() * self.element_type()._size_())]
             case _:
                 assert_unreachable()
 
@@ -135,7 +139,7 @@ class Array[T, Size](GenericValue, ArrayLike[T]):
         return self
 
     def _set_(self, value: Self):
-        raise TypeError("Array does not support set_")
+        raise TypeError("Array does not support _set_")
 
     def _copy_from_(self, value: Self):
         if not isinstance(value, type(self)):
@@ -150,6 +154,7 @@ class Array[T, Size](GenericValue, ArrayLike[T]):
             result._copy_from_(self)
             return result
         else:
+            assert isinstance(self._value, list)
             return self._with_value([value._copy_() for value in self._value])
 
     @classmethod
@@ -186,22 +191,36 @@ class Array[T, Size](GenericValue, ArrayLike[T]):
                     return self._value[const_index]._get_()
                 else:
                     return self._value[const_index]._get_()._as_py_()
-            else:
+            elif isinstance(self._value, BlockPlace):
                 return (
                     self.element_type()
                     ._from_place_(self._value.add_offset(const_index * self.element_type()._size_()))
                     ._get_()
                 )
+            elif callable(self._value):
+                return self.element_type()._from_backing_source_(
+                    lambda offset: self._value((Num(offset) + Num(const_index * self.element_type()._size_())).ir())
+                )
+            else:
+                raise InternalError("Unexpected array value")
         else:
             if not ctx():
                 raise InternalError("Unexpected non-constant index")
-            base = ctx().rom[tuple(self._to_list_())] if isinstance(self._value, list) else self._value
-            place = BlockPlace(
-                block=base.block,
-                index=(Num(base.index) + index * self.element_type()._size_()).index(),
-                offset=base.offset,
-            )
-            return self.element_type()._from_place_(place)._get_()
+            if isinstance(self._value, list | BlockPlace):
+                base = ctx().rom[tuple(self._to_list_())] if isinstance(self._value, list) else self._value
+                place = BlockPlace(
+                    block=base.block,
+                    index=(Num(base.index) + index * self.element_type()._size_()).index(),
+                    offset=base.offset,
+                )
+                return self.element_type()._from_place_(place)._get_()
+            elif callable(self._value):
+                base_offset = index * Num(self.element_type()._size_())
+                return self.element_type()._from_backing_source_(
+                    lambda offset: self._value((Num(offset) + base_offset).ir())
+                )
+            else:
+                raise InternalError("Unexpected array value")
 
     @meta_fn
     def __setitem__(self, index: Num, value: T):
@@ -210,17 +229,25 @@ class Array[T, Size](GenericValue, ArrayLike[T]):
         if ctx():
             if isinstance(self._value, list):
                 raise ValueError("Cannot mutate a compile time constant array")
-            base = self._value
-            place = (
-                base.add_offset(int(index._as_py_()) * self.element_type()._size_())
-                if index._is_py_()
-                else BlockPlace(
-                    block=base.block,
-                    index=(Num(base.index) + index * self.element_type()._size_()).index(),
-                    offset=base.offset,
+            elif isinstance(self._value, BlockPlace):
+                base = self._value
+                place = (
+                    base.add_offset(int(index._as_py_()) * self.element_type()._size_())
+                    if index._is_py_()
+                    else BlockPlace(
+                        block=base.block,
+                        index=(Num(base.index) + index * self.element_type()._size_()).index(),
+                        offset=base.offset,
+                    )
+                )
+                dst = self.element_type()._from_place_(place)
+            elif callable(self._value):
+                base_offset = index * Num(self.element_type()._size_())
+                dst = self.element_type()._from_backing_source_(
+                    lambda offset: self._value((Num(offset) + base_offset).ir())
                 )
-            )
-            dst = self.element_type()._from_place_(place)
+            else:
+                raise InternalError("Unexpected array value")
             if self.element_type()._is_value_type_():
                 dst._set_(value)
             else:
@@ -259,11 +286,13 @@ class Array[T, Size](GenericValue, ArrayLike[T]):
         return hash(tuple(self[i] for i in range(self.size())))
 
     def __str__(self):
-        if isinstance(self._value, BlockPlace):
+        if isinstance(self._value, BlockPlace) or callable(self._value):
             return f"{type(self).__name__}({self._value}...)"
-        return f"{type(self).__name__}({", ".join(str(self[i]) for i in range(self.size()))})"
+        else:
+            return f"{type(self).__name__}({", ".join(str(self[i]) for i in range(self.size()))})"
 
     def __repr__(self):
-        if isinstance(self._value, BlockPlace):
+        if isinstance(self._value, BlockPlace) or callable(self._value):
             return f"{type(self).__name__}({self._value}...)"
-        return f"{type(self).__name__}({", ".join(repr(self[i]) for i in range(self.size()))})"
+        else:
+            return f"{type(self).__name__}({", ".join(repr(self[i]) for i in range(self.size()))})"

sonolus/script/array_like.py

@@ -158,10 +158,14 @@ class ArrayLike[T](Sequence, ABC):
         return min_index
 
     def _max_(self, key: Callable[T, Any] | None = None) -> T:
-        return self[self.index_of_max(key=key)]
+        index = self.index_of_max(key=key)
+        assert index != -1
+        return self[index]
 
     def _min_(self, key: Callable[T, Any] | None = None) -> T:
-        return self[self.index_of_min(key=key)]
+        index = self.index_of_min(key=key)
+        assert index != -1
+        return self[index]
 
     def swap(self, i: Num, j: Num, /):
         """Swap the values at the given indices.
@@ -184,6 +188,7 @@ class ArrayLike[T](Sequence, ABC):
         if len(self) < 15 or key is not None:
             if key is None:
                 key = _identity
+            # May be worth adding a block sort variant for better performance on large arrays in the future
             _insertion_sort(self, 0, len(self), key, reverse)
         else:
             # Heap sort is unstable, so if there's a key, we can't rely on it

sonolus/script/bucket.py

@@ -20,7 +20,7 @@ class JudgmentWindow(Record):
     """The window for judging the accuracy of a hit.
 
     Usage:
-        ```
+        ```python
         JudgmentWindow(perfect: Interval, great: Interval, good: Interval)
         ```
     """
@@ -62,7 +62,7 @@ class JudgmentWindow(Record):
             target: The target time of the hit.
 
         Returns:
-            The judgment of the hit.
+            The [`Judgment`][sonolus.script.bucket.Judgment] of the hit.
         """
         return _judge(
             actual,
@@ -133,7 +133,9 @@ class Bucket(Record):
     """A bucket for entity judgment results.
 
     Usage:
-        ```Bucket(id: int)```
+        ```python
+        Bucket(id: int)
+        ```
     """
 
     id: int

sonolus/script/containers.py

@@ -1,13 +1,44 @@
 from __future__ import annotations
 
+from sonolus.backend.visitor import compile_and_call
 from sonolus.script.array import Array
 from sonolus.script.array_like import ArrayLike
 from sonolus.script.debug import error
+from sonolus.script.internal.context import ctx
+from sonolus.script.internal.impl import meta_fn
 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
 
 
+class Box[T](Record):
+    """A box that contains a value.
+
+    This can be helpful for generic code that can handle both Num and non-Num types.
+
+    Usage:
+        ```python
+        Box[T](value: T)
+        ```
+
+    Examples:
+        ```python
+        box = Box(1)
+        box = Box[int](2)
+
+        x: T = ...
+        y: T = ...
+        box = Box(x)
+        box.value = y  # Works regardless of whether x is a Num or not
+        ```
+    """
+
+    value: T
+    """The value contained in the box."""
+
+
 class Pair[T, U](Record):
     """A generic pair of values.
 
@@ -129,6 +160,17 @@ class VarArray[T, Capacity](Record, ArrayLike[T]):
         self._array[self._size] = value
         self._size += 1
 
+    def append_unchecked(self, value: T):
+        """Append the given value to the end of the array without checking the capacity.
+
+        Use with caution as this may cause hard to debug issues if the array is full.
+
+        Args:
+            value: The value to append.
+        """
+        self._array[self._size] = value
+        self._size += 1
+
     def extend(self, values: ArrayLike[T]):
         """Appends copies of the values in the given array to the end of the array.
 
@@ -258,6 +300,130 @@ class VarArray[T, Capacity](Record, ArrayLike[T]):
         raise TypeError("unhashable type: 'VarArray'")
 
 
+class ArrayPointer[T](Record, ArrayLike[T]):
+    """An array defined by a size and pointer to the first element.
+
+    This is intended to be created internally and improper use may result in hard to debug issues.
+
+    Usage:
+        ```python
+        ArrayPointer[T](size: int, block: int, offset: int)
+        ```
+    """
+
+    size: int
+    block: int
+    offset: int
+
+    def __len__(self) -> int:
+        """Return the number of elements in the array."""
+        return self.size
+
+    @classmethod
+    def element_type(cls) -> type[T]:
+        """Return the type of the elements in the array."""
+        return cls.type_var_value(T)
+
+    def _check_index(self, index: int):
+        assert 0 <= index < self.size
+
+    @meta_fn
+    def _get_item(self, item: int) -> T:
+        if not ctx():
+            raise TypeError("ArrayPointer values cannot be accessed outside of a context")
+        return _deref(
+            self.block,
+            self.offset + Num._accept_(item) * Num._accept_(self.element_type()._size_()),
+            self.element_type(),
+        )
+
+    @meta_fn
+    def __getitem__(self, item: int) -> T:
+        compile_and_call(self._check_index, item)
+        return self._get_item(item)._get_()
+
+    @meta_fn
+    def __setitem__(self, key: int, value: T):
+        compile_and_call(self._check_index, key)
+        dst = self._get_item(key)
+        if self.element_type()._is_value_type_():
+            dst._set_(value)
+        else:
+            dst._copy_from__(value)
+
+
+class ArraySet[T, Capacity](Record):
+    """A set implemented as an array with a fixed maximum capacity.
+
+    Usage:
+        ```python
+        ArraySet[T, Capacity].new()  # Create a new empty set
+        ```
+
+    Examples:
+        ```python
+        s = ArraySet[int, 10].new()
+        s.add(1)
+        s.add(2)
+        assert 1 in s
+        assert 3 not in s
+        s.remove(1)
+        assert 1 not in s
+        ```
+    """
+
+    _values: VarArray[T, Capacity]
+
+    @classmethod
+    def new(cls):
+        """Create a new empty set."""
+        element_type = cls.type_var_value(T)
+        capacity = cls.type_var_value(Capacity)
+        return cls(VarArray[element_type, capacity].new())
+
+    def __len__(self):
+        """Return the number of elements in the set."""
+        return len(self._values)
+
+    def __contains__(self, value):
+        """Return whether the given value is present in the set."""
+        return value in self._values
+
+    def __iter__(self):
+        """Return an iterator over the values in the set."""
+        return self._values.__iter__()
+
+    def add(self, value: T) -> bool:
+        """Add a copy of the given value to the set.
+
+        This has no effect and returns False if the value is already present or if the set is full.
+
+        Args:
+            value: The value to add.
+
+        Returns:
+            True if the value was added, False otherwise.
+        """
+        return self._values.set_add(value)
+
+    def remove(self, value: T) -> bool:
+        """Remove the given value from the set.
+
+        This has no effect and returns False if the value is not present.
+
+        Args:
+            value: The value to remove.
+
+        Returns:
+            True if the value was removed, False otherwise.
+        """
+        return self._values.set_remove(value)
+
+    def clear(self):
+        """Clear the set, removing all elements."""
+        self._values.clear()
+
+
 class _ArrayMapEntry[K, V](Record):
     key: K
     value: V

sonolus/script/debug.py

@@ -16,8 +16,13 @@ debug_log_callback = ContextVar[Callable[[Num], None]]("debug_log_callback")
 
 
 @meta_fn
-def error(message: str | None = None) -> None:
-    message = message._as_py_() if message is not None else "Error"
+def error(message: str | None = None) -> Never:
+    """Raise an error.
+
+    This function is used to raise an error during runtime.
+    When this happens, the game will pause in debug mode. The current callback will also immediately return 0.
+    """
+    message = validate_value(message)._as_py_() if message is not None else "Error"
     if not isinstance(message, str):
         raise ValueError("Expected a string")
     if ctx():
@@ -28,6 +33,19 @@ def error(message: str | None = None) -> None:
         raise RuntimeError(message)
 
 
+@meta_fn
+def static_error(message: str | None = None) -> Never:
+    """Raise a static error.
+
+    This function is used to raise an error during compile-time if the compiler cannot guarantee that
+    this function will not be called during runtime.
+    """
+    message = validate_value(message)._as_py_() if message is not None else "Error"
+    if not isinstance(message, str):
+        raise ValueError("Expected a string")
+    raise RuntimeError(message)
+
+
 @meta_fn
 def debug_log(value: Num):
     """Log a value in debug mode."""

sonolus/script/engine.py

@@ -4,7 +4,7 @@ import json
 from collections.abc import Callable
 from os import PathLike
 from pathlib import Path
-from typing import Any
+from typing import Any, Literal
 
 from sonolus.build.collection import Asset, load_asset
 from sonolus.script.archetype import PlayArchetype, PreviewArchetype, WatchArchetype, _BaseArchetype
@@ -81,7 +81,7 @@ class Engine:
         meta: Additional metadata of the engine.
     """
 
-    version = 12
+    version: Literal[13] = 13
 
     def __init__(
         self,

sonolus/script/internal/constant.py

@@ -3,7 +3,7 @@ from typing import Any, ClassVar, Self
 
 from sonolus.backend.place import BlockPlace
 from sonolus.script.internal.impl import meta_fn
-from sonolus.script.internal.value import Value
+from sonolus.script.internal.value import DataValue, Value
 
 
 class _Missing:
@@ -18,6 +18,8 @@ _MISSING = _Missing()
 
 
 class ConstantValue(Value):
+    """Wraps a python constant value usable in Sonolus scripts."""
+
     _parameterized_: ClassVar[dict[Any, type[Self]]] = {}
     _value: ClassVar[Any] = _MISSING
     instance: ClassVar[Self | _Missing] = _MISSING
@@ -94,10 +96,10 @@ class ConstantValue(Value):
         return self.value()
 
     @classmethod
-    def _from_list_(cls, values: Iterable[float | BlockPlace]) -> Self:
+    def _from_list_(cls, values: Iterable[DataValue]) -> Self:
         return cls()
 
-    def _to_list_(self, level_refs: dict[Any, str] | None = None) -> list[float | str | BlockPlace]:
+    def _to_list_(self, level_refs: dict[Any, str] | None = None) -> list[DataValue | str]:
         return []
 
     @classmethod
@@ -139,5 +141,5 @@ class ConstantValue(Value):
         return hash(self.value())
 
 
-class MiscConstantValue(ConstantValue):
+class BasicConstantValue(ConstantValue):
     """For constants without any special behavior."""

sonolus/script/internal/generic.py

@@ -1,6 +1,8 @@
 from __future__ import annotations
 
+import typing
 from enum import Enum
+from types import UnionType
 from typing import Annotated, Any, ClassVar, Literal, Self, TypeVar, get_origin
 
 from sonolus.script.internal.impl import meta_fn, validate_value
@@ -31,6 +33,11 @@ def validate_type_spec(spec: Any) -> PartialGeneric | TypeVar | type[Value]:
         spec = validate_type_spec(spec.__mro__[1])
     if isinstance(spec, PartialGeneric | TypeVar) or (isinstance(spec, type) and issubclass(spec, Value)):
         return spec
+    if typing.get_origin(spec) is UnionType:
+        args = typing.get_args(spec)
+        validated_args = {validate_type_arg(arg) for arg in args}
+        if len(validated_args) == 1:
+            return validated_args.pop()
     raise TypeError(f"Invalid type spec: {spec}")