sonolus.py 0.3.4__py3-none-any.whl → 0.4.0__py3-none-any.whl

sonolus/script/internal/value.py

@@ -1,5 +1,6 @@
 from abc import abstractmethod
 from collections.abc import Callable, Iterable
+from types import NotImplementedType
 from typing import Any, Self
 
 from sonolus.backend.ir import IRConst, IRExpr, IRStmt
@@ -136,8 +137,19 @@ class Value:
         """
         raise NotImplementedError
 
+    def _get_readonly_(self) -> Self:
+        """Implements access to the value without copying if the underlying value is immutable.
+
+        The returned value should not be intentionally modified, but it is not guaranteed to be immutable.
+
+        For example, a Num might be backed internally by rom, which is immutable. If we aren't going to modify
+        (e.g. by putting it into a record where it can be modified), we can just return the original value, which
+        avoids unnecessary copying.
+        """
+        return self._get_()
+
     @abstractmethod
-    def _set_(self, value: Self):
+    def _set_(self, value: Any):
         """Implements assignment (=).
 
         This is only supported by value types.
@@ -149,7 +161,7 @@ class Value:
         raise NotImplementedError
 
     @abstractmethod
-    def _copy_from_(self, value: Self):
+    def _copy_from_(self, value: Any):
         """Implements copy assignment (@=).
 
         This is only supported by mutable reference types.
@@ -173,9 +185,35 @@ class Value:
         """Returns a zero-initialized value of this type."""
         raise NotImplementedError
 
+    @classmethod
+    def _get_merge_target_(cls, values: list[Any]) -> Any | NotImplementedType:
+        """Return the target when merging values from multiple code paths.
+
+        E.g. for code like this:
+        ```
+        if cond:
+            x = 1
+        else:
+            x = 2
+        do_something(x)
+        ```
+        This is called to create a target value for x after the if-else block,
+        and at the end of each block, that target value is assigned to the respective value (1 or 2).
+        This lets us keep the value of x as a constant within if and else branches, and only have it
+        become a runtime value after the if-else block.
+
+        This is an overrideable method to allow for some other special behavior, namely the Maybe type.
+        """
+        if cls._is_value_type_():
+            from sonolus.script.internal.context import ctx
+
+            return cls._from_place_(ctx().alloc(size=cls._size_()))
+        else:
+            return NotImplemented
+
     def __imatmul__(self, other):
         self._copy_from_(other)
         return self
 
 
-Value.__imatmul__._meta_fn_ = True
+Value.__imatmul__._meta_fn_ = True  # type: ignore

sonolus/script/interval.py

@@ -1,4 +1,4 @@
-from typing import Self
+from __future__ import annotations
 
 from sonolus.backend.ops import Op
 from sonolus.script.array_like import ArrayLike
@@ -22,7 +22,7 @@ class Interval(Record):
     end: float
 
     @classmethod
-    def zero(cls) -> Self:
+    def zero(cls) -> Interval:
         """Get an empty interval."""
         return cls(0, 0)
 
@@ -49,7 +49,7 @@ class Interval(Record):
         """The interval as a tuple."""
         return self.start, self.end
 
-    def __contains__(self, item: Self | float | int) -> bool:
+    def __contains__(self, item: Interval | float | int) -> bool:
         """Check if an item is within the interval.
 
         Args:
@@ -66,7 +66,7 @@ class Interval(Record):
             case _:
                 static_error("Invalid type for interval check")
 
-    def __add__(self, other: float | int) -> Self:
+    def __add__(self, other: float | int) -> Interval:
         """Add a value to both ends of the interval.
 
         Args:
@@ -77,7 +77,7 @@ class Interval(Record):
         """
         return Interval(self.start + other, self.end + other)
 
-    def __sub__(self, other: float | int) -> Self:
+    def __sub__(self, other: float | int) -> Interval:
         """Subtract a value from both ends of the interval.
 
         Args:
@@ -88,7 +88,7 @@ class Interval(Record):
         """
         return Interval(self.start - other, self.end - other)
 
-    def __mul__(self, other: float | int) -> Self:
+    def __mul__(self, other: float | int) -> Interval:
         """Multiply both ends of the interval by a value.
 
         Args:
@@ -99,7 +99,7 @@ class Interval(Record):
         """
         return Interval(self.start * other, self.end * other)
 
-    def __truediv__(self, other: float | int) -> Self:
+    def __truediv__(self, other: float | int) -> Interval:
         """Divide both ends of the interval by a value.
 
         Args:
@@ -110,7 +110,7 @@ class Interval(Record):
         """
         return Interval(self.start / other, self.end / other)
 
-    def __floordiv__(self, other: float | int) -> Self:
+    def __floordiv__(self, other: float | int) -> Interval:
         """Divide both ends of the interval by a value and floor the result.
 
         Args:
@@ -121,7 +121,7 @@ class Interval(Record):
         """
         return Interval(self.start // other, self.end // other)
 
-    def __and__(self, other: Self) -> Self:
+    def __and__(self, other: Interval) -> Interval:
         """Get the intersection of two intervals.
 
         The resulting interval will be empty and may have a negative length if the two intervals do not overlap.
@@ -134,7 +134,7 @@ class Interval(Record):
         """
         return Interval(max(self.start, other.start), min(self.end, other.end))
 
-    def shrink(self, value: float | int) -> Self:
+    def shrink(self, value: float | int) -> Interval:
         """Shrink the interval by a value on both ends.
 
         Args:
@@ -145,7 +145,7 @@ class Interval(Record):
         """
         return Interval(self.start + value, self.end - value)
 
-    def expand(self, value: float | int) -> Self:
+    def expand(self, value: float | int) -> Interval:
         """Expand the interval by a value on both ends.
 
         Args:
@@ -223,11 +223,11 @@ def _num_lerp_clamped(a, b, x, /):
 
 
 def _generic_lerp[T](a: T, b: T, x: float, /) -> T:
-    return a + (b - a) * x
+    return a + (b - a) * x  # type: ignore
 
 
 def _generic_lerp_clamped[T](a: T, b: T, x: float, /) -> T:
-    return a + (b - a) * max(0, min(1, x))
+    return a + (b - a) * max(0, min(1, x))  # type: ignore
 
 
 def lerp[T](a: T, b: T, x: float, /) -> T:

sonolus/script/iterator.py

@@ -1,9 +1,9 @@
 from __future__ import annotations
 
-from abc import abstractmethod
 from typing import Any
 
 from sonolus.script.internal.impl import meta_fn
+from sonolus.script.maybe import Maybe, Nothing, Some
 from sonolus.script.record import Record
 
 
@@ -12,57 +12,28 @@ class SonolusIterator[T]:
 
     This class is used to define custom iterators that can be used in Sonolus.py.
 
-    Inheritors must implement the `has_next`, `get`, and `advance` methods.
-    The `__next__` and `__iter__` methods are implemented by default.
+    Inheritors must implement the `next` method, which should return a `Maybe[T]`.
 
     Usage:
         ```python
         class MyIterator(Record, SonolusIterator):
-            def has_next(self) -> bool:
-                ...
-
-            def get(self) -> Any:
-                ...
-
-            def advance(self):
+            def next(self) -> Maybe[T]:
                 ...
         ```
     """
 
     _allow_instance_check_ = True
 
-    def next(self) -> T:
-        result = self.get()
-        self.advance()
-        return result
-
-    @abstractmethod
-    def has_next(self) -> bool:
-        """Return whether the iterator has more elements."""
-        raise NotImplementedError
-
-    @abstractmethod
-    def get(self) -> T:
-        """Return the next element of the iterator.
-
-        May be called multiple times before calling `advance`.
-
-        Must not be called if `has_next` returns `False`.
-        """
-        raise NotImplementedError
-
-    @abstractmethod
-    def advance(self):
-        """Advance the iterator to the next element.
-
-        Must not be called if `has_next` returns `False`.
-        """
+    @meta_fn
+    def next(self) -> Maybe[T]:
         raise NotImplementedError
 
     def __next__(self) -> T:
-        if not self.has_next():
+        result = self.next()
+        if result.is_some:
+            return result.get_unsafe()
+        else:
             raise StopIteration
-        return self.next()
 
     def __iter__(self) -> SonolusIterator[T]:
         return self
@@ -73,15 +44,13 @@ class _Enumerator[V: SonolusIterator](Record, SonolusIterator):
     offset: int
     iterator: V
 
-    def has_next(self) -> bool:
-        return self.iterator.has_next()
-
-    def get(self) -> tuple[int, Any]:
-        return self.i + self.offset, self.iterator.get()
-
-    def advance(self):
+    def next(self) -> Maybe[tuple[int, Any]]:
+        value = self.iterator.next()
+        if value.is_nothing:
+            return Nothing
+        result = (self.i + self.offset, value.get_unsafe())
         self.i += 1
-        self.iterator.advance()
+        return Some(result)
 
 
 class _Zipper[T](Record, SonolusIterator):
@@ -89,11 +58,6 @@ class _Zipper[T](Record, SonolusIterator):
     iterators: T
 
     @meta_fn
-    def has_next(self) -> bool:
-        from sonolus.backend.visitor import compile_and_call
-
-        return compile_and_call(self._has_next, self._get_iterators())
-
     def _get_iterators(self) -> tuple[SonolusIterator, ...]:
         from sonolus.script.containers import Pair
 
@@ -105,81 +69,48 @@ class _Zipper[T](Record, SonolusIterator):
         iterators.append(v)
         return tuple(iterators)
 
-    def _has_next(self, iterators: tuple[SonolusIterator, ...]) -> bool:
-        for iterator in iterators:  # noqa: SIM110
-            if not iterator.has_next():
-                return False
-        return True
-
     @meta_fn
-    def get(self) -> tuple[Any, ...]:
+    def _get_next_values(self) -> tuple[Any, ...]:
         from sonolus.backend.visitor import compile_and_call
 
-        return tuple(compile_and_call(iterator.get) for iterator in self._get_iterators())
+        return tuple(compile_and_call(iterator.next) for iterator in self._get_iterators())
 
     @meta_fn
-    def advance(self):
+    def _values_to_tuple(self, values: tuple[Any, ...]) -> tuple[Any, ...]:
         from sonolus.backend.visitor import compile_and_call
 
-        for iterator in self._get_iterators():
-            compile_and_call(iterator.advance)
-
+        return tuple(compile_and_call(value.get_unsafe) for value in values)
 
-class _EmptyIterator(Record, SonolusIterator):
-    def has_next(self) -> bool:
-        return False
+    def next(self) -> Maybe[tuple[Any, ...]]:
+        values = self._get_next_values()
+        for value in values:
+            if value.is_nothing:
+                return Nothing
+        return Some(self._values_to_tuple(values))
 
-    def get(self) -> Any:
-        return None
 
-    def advance(self):
-        pass
+class _EmptyIterator(Record, SonolusIterator):
+    def next(self) -> Maybe[Any]:
+        return Nothing
 
 
 class _MappingIterator[T, Fn](Record, SonolusIterator):
     fn: Fn
     iterator: T
 
-    def has_next(self) -> bool:
-        return self.iterator.has_next()
-
-    def get(self) -> Any:
-        return self.fn(self.iterator.get())
-
-    def advance(self):
-        self.iterator.advance()
+    def next(self) -> Maybe[Any]:
+        return self.iterator.next().map(self.fn)
 
 
 class _FilteringIterator[T, Fn](Record, SonolusIterator):
     fn: Fn
     iterator: T
 
-    def has_next(self) -> bool:
-        while self.iterator.has_next():
-            if self.fn(self.iterator.get()):
-                return True
-            self.iterator.advance()
-        return False
-
-    def get(self) -> Any:
-        return self.iterator.get()
-
-    def advance(self):
-        self.iterator.advance()
-
-
-class _ChainingIterator[T](Record, SonolusIterator):
-    iterator: T
-
-    def has_next(self) -> bool:
-        return self.iterator.has_next()
-
-    def get(self) -> Any:
-        return self.iterator.get().get()
-
-    def advance(self):
-        self.iterator.get().advance()
-        while not self.iterator.get().has_next():
-            self.iterator.advance()
-            if not self.iterator.has_next():
-                break
+    def next(self) -> Maybe[T]:
+        while True:
+            value = self.iterator.next()
+            if value.is_nothing:
+                return Nothing
+            inside = value.get_unsafe()
+            if self.fn(inside):
+                return Some(inside)

sonolus/script/maybe.py

@@ -0,0 +1,139 @@
+from __future__ import annotations
+
+from collections.abc import Callable
+from typing import Any
+
+from sonolus.script.internal.context import ctx
+from sonolus.script.internal.impl import meta_fn
+from sonolus.script.internal.transient import TransientValue
+from sonolus.script.internal.value import Value
+from sonolus.script.num import Num
+
+
+class Maybe[T](TransientValue):
+    """A type that either has a value or is empty.
+
+    Maybe has special behavior when returned from a function: it may be returned from multiple places
+    in a function, provided that all but one return statement returns the literal `Nothing`.
+
+    This type is not intended for other uses, such as being stored in a Record, Array, or Archetype.
+
+    Usage:
+        ```python
+        def fn(a, b):
+            if a:
+                return Some(b)
+            else:
+                return Nothing
+
+        result = fn(..., ...)
+        if result.is_some:
+            value = result.get()
+            ...
+        ```
+    """
+
+    _present: Num
+    _value: T
+
+    def __init__(self, *, present: bool, value: T):
+        self._present = Num._accept_(present)
+        self._value = value
+
+    @property
+    @meta_fn
+    def is_some(self) -> bool:
+        """Check if the value is present."""
+        if ctx():
+            return self._present._get_readonly_()
+        else:
+            return self._present._as_py_()
+
+    @property
+    def is_nothing(self) -> bool:
+        """Check if the value is empty."""
+        return not self.is_some
+
+    def get(self) -> T:
+        """Get the value if present, otherwise raise an error."""
+        assert self.is_some
+        return self.get_unsafe()
+
+    @meta_fn
+    def get_unsafe(self) -> T:
+        return self._value
+
+    def map[R](self, fn: Callable[[T], R], /) -> Maybe[R]:
+        if self.is_some:
+            return Some(fn(self.get_unsafe()))
+        return Nothing
+
+    @classmethod
+    def _accepts_(cls, value: Any) -> bool:
+        return isinstance(value, cls)
+
+    @classmethod
+    def _accept_(cls, value: Any) -> Maybe[T]:
+        if not cls._accepts_(value):
+            raise TypeError(f"Cannot accept value of type {type(value).__name__} as {cls.__name__}.")
+        return value
+
+    def _is_py_(self) -> bool:
+        return not self._present or not isinstance(self._value, Value) or self._value._is_py_()
+
+    def _as_py_(self) -> Any:
+        if not self._is_py_():
+            raise ValueError("Not a python value")
+        return self
+
+    def _copy_from_(self, value: Any):
+        raise TypeError("Maybe does not support mutation.")
+
+    def _copy_(self) -> Maybe[T]:
+        raise TypeError("Maybe does not support copying.")
+
+    def _set_(self, value: Any):
+        if not self._accepts_(value):
+            raise TypeError(f"Cannot set value of type {type(value).__name__} to {self.__class__.__name__}.")
+        if value is not Nothing and self._value is not value._value:
+            raise TypeError(f"Cannot set value of type {type(value._value).__name__} to {self.__class__.__name__}.")
+        self._present._set_(value._present)
+
+    @classmethod
+    def _get_merge_target_(cls, values: list[Any]) -> Any:
+        if not all(isinstance(v, cls) for v in values):
+            return NotImplemented
+        distinct = []
+        seen_ids = set()
+        for v in values:
+            if v is Nothing:
+                continue
+            if id(v._value) not in seen_ids:
+                distinct.append(v)
+                seen_ids.add(id(v._value))
+        match distinct:
+            case []:
+                return Nothing
+            case [v]:
+                return Maybe(present=Num._alloc_(), value=v._value)
+            case _:
+                return NotImplemented
+
+
+def Some[T](value: T) -> Maybe[T]:  # noqa: N802
+    """Create a `Maybe` instance with a value.
+
+    Args:
+        value: The contained value.
+
+    Returns:
+        A `Maybe` instance that contains the provided value.
+    """
+    return Maybe(present=True, value=value)
+
+
+Nothing: Maybe[Any] = Maybe(present=False, value=None)  # type: ignore
+
+# Note: has to come after the definition to hide the definition in the docs.
+Nothing: Maybe[Any]
+"""The empty `Maybe` instance."""

sonolus/script/num.py

@@ -1,9 +1,11 @@
+# type: ignore
 from __future__ import annotations
 
 import operator
 from collections.abc import Callable, Iterable
 from typing import TYPE_CHECKING, Any, Self, TypeGuard, final, runtime_checkable
 
+from sonolus.backend.blocks import BlockData
 from sonolus.backend.ir import IRConst, IRExpr, IRGet, IRPureInstr, IRSet
 from sonolus.backend.ops import Op
 from sonolus.backend.place import BlockPlace
@@ -114,15 +116,28 @@ class _Num(Value, metaclass=_NumMeta):
 
     def _get_(self) -> Self:
         if ctx():
+            if isinstance(self.data, BlockPlace):
+                ctx().check_readable(self.data)
             place = ctx().alloc(size=1)
+            ctx().add_statements(IRSet(place, self.ir()))
+            return Num(place)
+        else:
+            return Num(self.data)
+
+    def _get_readonly_(self) -> Self:
+        if ctx():
             if isinstance(self.data, BlockPlace):
                 ctx().check_readable(self.data)
+                if isinstance(self.data.block, BlockData) and not ctx().is_writable(self.data):
+                    # This block is immutable in the current callback, so no need to copy it in case it changes.
+                    return Num(self.data)
+            place = ctx().alloc(size=1)
             ctx().add_statements(IRSet(place, self.ir()))
             return Num(place)
         else:
             return Num(self.data)
 
-    def _set_(self, value: Self):
+    def _set_(self, value: Any):
         value = Num._accept_(value)
         if ctx():
             match self.data:
@@ -136,7 +151,7 @@ class _Num(Value, metaclass=_NumMeta):
         else:
             self.data = value.data
 
-    def _copy_from_(self, value: Self):
+    def _copy_from_(self, value: Any):
         raise ValueError("Cannot assign to a number")
 
     def _copy_(self) -> Self:

sonolus/script/options.py

@@ -175,7 +175,7 @@ def select_option(
     return _SelectOption(name, description, standard, advanced, scope, default, values)
 
 
-type Options = NewType("Options", Any)
+type Options = NewType("Options", Any)  # type: ignore
 type _OptionInfo = _SliderOption | _ToggleOption | _SelectOption
 
 

sonolus/script/particle.py

@@ -95,7 +95,7 @@ def particle(name: str) -> Any:
     return _ParticleInfo(name)
 
 
-type Particles = NewType("Particles", Any)
+type Particles = NewType("Particles", Any)  # type: ignore
 
 
 @dataclass_transform()

sonolus/script/project.py

@@ -1,10 +1,10 @@
 from __future__ import annotations
 
-from collections.abc import Sequence
+from collections.abc import Callable, Iterable, Sequence
 from dataclasses import dataclass
 from os import PathLike
 from pathlib import Path
-from typing import ClassVar, Self, TypedDict
+from typing import ClassVar, TypedDict
 
 from sonolus.backend.optimize import optimize
 from sonolus.backend.optimize.passes import CompilerPass
@@ -25,14 +25,23 @@ class Project:
     def __init__(
         self,
         engine: Engine,
-        levels: list[Level] | None = None,
+        levels: Iterable[Level] | Callable[[], Iterable[Level]] | None = None,
         resources: PathLike | None = None,
     ):
         self.engine = engine
-        self.levels = levels or []
+        match levels:
+            case Callable():
+                self._level_source = lazy_loader(levels)
+            case Iterable():
+                self._level_source = levels
+            case None:
+                self._level_source = []
+            case _:
+                raise TypeError(f"Invalid type for levels: {type(levels)}. Expected Iterable or Callable.")
+        self._levels = None
         self.resources = Path(resources or "resources")
 
-    def with_levels(self, levels: list[Level]) -> Self:
+    def with_levels(self, levels: list[Level]) -> Project:
         """Create a new project with the specified levels.
 
         Args:
@@ -78,6 +87,16 @@ class Project:
 
         return get_project_schema(self)
 
+    @property
+    def levels(self) -> list[Level]:
+        if self._levels is None:
+            self._levels = list(self._level_source)
+        return self._levels
+
+
+def lazy_loader(fn):
+    yield from fn()
+
 
 class ProjectSchema(TypedDict):
     archetypes: list[ArchetypeSchema]