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

sonolus/script/internal/context.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+from collections.abc import Iterable
 from contextlib import contextmanager
 from contextvars import ContextVar
 from dataclasses import dataclass
@@ -7,7 +8,7 @@ from threading import Lock
 from typing import Any, Self
 
 from sonolus.backend.blocks import BlockData, PlayBlock
-from sonolus.backend.ir import IRConst, IRStmt
+from sonolus.backend.ir import IRConst, IRExpr, IRStmt
 from sonolus.backend.mode import Mode
 from sonolus.backend.optimize.flow import BasicBlock, FlowEdge
 from sonolus.backend.place import Block, BlockPlace, TempBlock
@@ -16,7 +17,7 @@ from sonolus.script.internal.value import Value
 
 _compiler_internal_ = True
 
-context_var = ContextVar("context_var", default=None)
+context_var: ContextVar[Context | None] = ContextVar("context_var", default=None)  # type: ignore
 
 
 @dataclass(frozen=True)
@@ -70,7 +71,7 @@ class Context:
     global_state: GlobalContextState
     callback_state: CallbackContextState
     statements: list[IRStmt]
-    test: IRStmt
+    test: IRExpr
     outgoing: dict[float | None, Context]
     scope: Scope
     loop_variables: dict[str, Value]
@@ -115,7 +116,7 @@ class Context:
             return
         if not self.callback:
             return
-        if isinstance(place.block, BlockData) and self.callback not in self.blocks(place.block).readable:
+        if isinstance(place.block, BlockData) and not self.is_readable(place):
             raise RuntimeError(f"Block {place.block} is not readable in {self.callback}")
 
     def check_writable(self, place: BlockPlace):
@@ -123,9 +124,19 @@ class Context:
             return
         if not self.callback:
             return
-        if isinstance(place.block, BlockData) and self.callback not in self.blocks(place.block).writable:
+        if isinstance(place.block, BlockData) and not self.is_writable(place):
             raise RuntimeError(f"Block {place.block} is not writable in {self.callback}")
 
+    def is_readable(self, place: BlockPlace) -> bool:
+        if debug_config().unchecked_reads:
+            return True
+        return self.callback and self.callback in self.blocks(place.block).readable
+
+    def is_writable(self, place: BlockPlace) -> bool:
+        if debug_config().unchecked_writes:
+            return True
+        return self.callback and self.callback in self.blocks(place.block).writable
+
     def add_statement(self, statement: IRStmt):
         if not self.live:
             return
@@ -169,20 +180,26 @@ class Context:
         self.outgoing[condition] = result
         return result
 
+    def new_disconnected(self):
+        return self.copy_with_scope(self.scope.copy())
+
+    def new_empty_disconnected(self):
+        return self.copy_with_scope(Scope())
+
     def into_dead(self):
         """Create a new context for code that is unreachable, like after a return statement."""
         result = self.copy_with_scope(self.scope.copy())
         result.live = False
         return result
 
-    def prepare_loop_header(self, to_merge: set[str]) -> Self:
+    def prepare_loop_header(self, to_merge: set[str]) -> Context:
         # to_merge is the set of bindings set anywhere in the loop
         # we need to invalidate them in the header if they're reference types
         # or merge them if they're value types
         # structure is self -> intermediate -> header -> body (continue -> header) | exit
         assert len(self.outgoing) == 0
         header = self.branch(None)
-        for name in to_merge:
+        for name in sorted(to_merge):
             binding = self.scope.get_binding(name)
             if not isinstance(binding, ValueBinding):
                 continue
@@ -205,12 +222,15 @@ class Context:
         assert len(self.outgoing) == 0
         self.outgoing[None] = header
         values = {}
-        # First do a pass through and copy every value
+        # First do a pass through and get every value
         for name, target_value in header.loop_variables.items():
             with using_ctx(self):
                 if type(target_value)._is_value_type_():
                     value = self.scope.get_value(name)
-                    values[name] = value._get_()  # _get_() will make a copy on value types
+                    # We make this call to _get_readonly_() to ensure that we're reading the value at this
+                    # point in time specifically, since _get_readonly_ will make a copy if the value is
+                    # e.g. a Num backed by a TempBlock which could be mutated.
+                    values[name] = value._get_readonly_()
         # Then actually set them
         for name, target_value in header.loop_variables.items():
             with using_ctx(self):
@@ -236,7 +256,7 @@ class Context:
         with self.global_state.lock:
             block = value.blocks.get(self.global_state.mode)
             if block is None:
-                raise RuntimeError(f"Global {value.name} is not available in '{self.global_state.mode.name}' mode")
+                raise RuntimeError(f"Global {value} is not available in '{self.global_state.mode.name}' mode")
             if value not in self.global_state.environment_mappings:
                 if value.offset is None:
                     offset = self.global_state.environment_offsets.get(block, 0)
@@ -268,7 +288,7 @@ class Context:
             return self.global_state.archetypes[type_]
 
 
-def ctx() -> Context | None:
+def ctx() -> Context | Any:  # Using Any to silence type checker warnings if it's None
     return context_var.get()
 
 
@@ -368,7 +388,7 @@ class Scope:
     def set_binding(self, name: str, binding: Binding):
         self.bindings[name] = binding
 
-    def get_value(self, name: str) -> Value:
+    def get_value(self, name: str) -> Value | Any:
         binding = self.get_binding(name)
         match binding:
             case ValueBinding(value):
@@ -398,7 +418,7 @@ class Scope:
             return
         assert all(len(inc.outgoing) == 0 for inc in incoming)
         sources = [context.scope for context in incoming]
-        keys = {key for source in sources for key in source.bindings}
+        keys = unique(key for source in sources for key in source.bindings)
         for key in keys:
             bindings = [source.get_binding(key) for source in sources]
             if not all(isinstance(binding, ValueBinding) for binding in bindings):
@@ -413,8 +433,9 @@ class Scope:
                 target.scope.set_binding(key, ConflictBinding())
                 continue
             common_type: type[Value] = types.pop()
-            if common_type._is_value_type_():
-                target_value = common_type._from_place_(target.alloc(size=common_type._size_()))
+            with using_ctx(target):
+                target_value = common_type._get_merge_target_(values)
+            if target_value is not NotImplemented:
                 for inc in incoming:
                     with using_ctx(inc):
                         target_value._set_(inc.scope.get_value(key))
@@ -447,3 +468,13 @@ def context_to_cfg(context: Context) -> BasicBlock:
             blocks[current].outgoing.add(edge)
             blocks[target].incoming.add(edge)
     return blocks[context]
+
+
+def unique[T](iterable: Iterable[T]) -> list[T]:
+    result = []
+    seen = set()
+    for item in iterable:
+        if item not in seen:
+            seen.add(item)
+            result.append(item)
+    return result

sonolus/script/internal/impl.py

@@ -3,7 +3,7 @@ from __future__ import annotations
 from collections.abc import Callable
 from enum import Enum
 from types import EllipsisType, FunctionType, MethodType, ModuleType, NoneType, NotImplementedType, UnionType
-from typing import TYPE_CHECKING, Annotated, Any, Literal, TypeVar, get_origin, overload
+from typing import TYPE_CHECKING, Annotated, Any, Literal, TypeVar, Union, get_origin, overload
 
 if TYPE_CHECKING:
     from sonolus.script.internal.value import Value
@@ -33,7 +33,7 @@ def meta_fn(fn=None):
     return decorator(fn)
 
 
-def validate_value(value: Any) -> Value:
+def validate_value[T](value: T) -> Value | T:
     result = try_validate_value(value)
     if result is None:
         raise TypeError(f"Unsupported value: {value!r}")
@@ -42,7 +42,7 @@ def validate_value(value: Any) -> Value:
 
 def try_validate_value(value: Any) -> Value | None:
     from sonolus.script.globals import _GlobalPlaceholder
-    from sonolus.script.internal.constant import BasicConstantValue
+    from sonolus.script.internal.constant import BasicConstantValue, TypingSpecialFormConstant
     from sonolus.script.internal.dict_impl import DictImpl
     from sonolus.script.internal.generic import PartialGeneric
     from sonolus.script.internal.tuple_impl import TupleImpl
@@ -85,6 +85,12 @@ def try_validate_value(value: Any) -> Value | None:
             | EllipsisType()
         ):
             return BasicConstantValue.of(value)
+        case special_form if value in {
+            Literal,
+            Annotated,
+            Union,
+        }:
+            return TypingSpecialFormConstant.of(special_form)
         case other_type if get_origin(value) in {Literal, Annotated, UnionType, tuple}:
             return BasicConstantValue.of(other_type)
         case _GlobalPlaceholder():

sonolus/script/internal/introspection.py

@@ -4,7 +4,14 @@ from typing import Annotated
 _missing = object()
 
 
-def get_field_specifiers(cls, *, skip: set[str] = frozenset(), globals=None, locals=None, eval_str=True):  # noqa: A002
+def get_field_specifiers(
+    cls,
+    *,
+    skip: frozenset[str] | set[str] = frozenset(),
+    globals=None,  # noqa: A002
+    locals=None,  # noqa: A002
+    eval_str=True,
+):
     """Like inspect.get_annotations, but also turns class attributes into Annotated."""
     results = inspect.get_annotations(cls, globals=globals, locals=locals, eval_str=eval_str)
     for key, value in results.items():

sonolus/script/internal/math_impls.py

@@ -97,6 +97,22 @@ def _log(x: float, base: float | None = None) -> float:
         return _ln(x) / _ln(base)
 
 
+def _sqrt(x: float) -> float:
+    return x**0.5
+
+
+@native_function(Op.Degree)
+def _degrees(x: float) -> float:
+    """Convert radians to degrees."""
+    return math.degrees(x)
+
+
+@native_function(Op.Radian)
+def _radians(x: float) -> float:
+    """Convert degrees to radians."""
+    return math.radians(x)
+
+
 @native_function(Op.Rem)
 def _remainder(x: float, y: float) -> float:
     # This is different from math.remainder in Python's math package, which could be confusing
@@ -119,4 +135,5 @@ MATH_BUILTIN_IMPLS = {
     id(math.trunc): _trunc,
     id(round): _round,
     id(math.log): _log,
+    id(math.sqrt): _sqrt,
 }

sonolus/script/internal/native.py

@@ -9,7 +9,7 @@ from sonolus.script.internal.impl import meta_fn, validate_value
 from sonolus.script.num import Num, _is_num
 
 
-def native_call(op: Op, *args: Num) -> Num:
+def native_call(op: Op, *args: int | float | bool) -> Num:
     if not ctx():
         raise RuntimeError("Unexpected native call")
     args = tuple(validate_value(arg) for arg in args)
@@ -21,20 +21,20 @@ def native_call(op: Op, *args: Num) -> Num:
 
 
 def native_function[**P, R](op: Op) -> Callable[[Callable[P, R]], Callable[P, R]]:
-    def decorator(fn: Callable[P, Num]) -> Callable[P, Num]:
+    def decorator(fn: Callable[P, int | float | bool]) -> Callable[P, Num]:
         signature = inspect.signature(fn)
 
         @functools.wraps(fn)
         @meta_fn
-        def wrapper(*args: Num) -> Num:
+        def wrapper(*args: int | float | bool) -> Num:
             if len(args) < sum(1 for p in signature.parameters.values() if p.default == inspect.Parameter.empty):
                 raise TypeError(f"Expected {len(signature.parameters)} arguments, got {len(args)}")
             if ctx():
                 bound_args = signature.bind(*args)
                 bound_args.apply_defaults()
                 return native_call(op, *(Num._accept_(arg) for arg in bound_args.args))
-            return fn(*args)
+            return fn(*args)  # type: ignore
 
         return wrapper
 
-    return decorator
+    return decorator  # type: ignore

sonolus/script/internal/range.py

@@ -2,16 +2,17 @@ from sonolus.script.array_like import ArrayLike, get_positive_index
 from sonolus.script.internal.context import ctx
 from sonolus.script.internal.impl import meta_fn, validate_value
 from sonolus.script.iterator import SonolusIterator
+from sonolus.script.maybe import Maybe, Nothing, Some
 from sonolus.script.num import Num
 from sonolus.script.record import Record
 
 
-class Range(Record, ArrayLike[Num]):
+class Range(Record, ArrayLike[int]):
     start: int
     stop: int
     step: int
 
-    def __new__(cls, start: Num, stop: Num | None = None, step: Num = 1):
+    def __new__(cls, start: int, stop: int | None = None, step: int = 1):
         if stop is None:
             start, stop = 0, start
         return super().__new__(cls, start, stop, step)
@@ -37,14 +38,14 @@ class Range(Record, ArrayLike[Num]):
                 return 0
             return (diff - self.step - 1) // -self.step
 
-    def __getitem__(self, index: Num) -> Num:
+    def __getitem__(self, index: int) -> int:
         return self.start + get_positive_index(index, len(self)) * self.step
 
-    def __setitem__(self, index: Num, value: Num):
+    def __setitem__(self, index: int, value: int):
         raise TypeError("Range does not support item assignment")
 
     @property
-    def last(self) -> Num:
+    def last(self) -> int:
         return self[len(self) - 1]
 
     def __eq__(self, other):
@@ -70,21 +71,17 @@ class RangeIterator(Record, SonolusIterator):
     stop: int
     step: int
 
-    def has_next(self) -> bool:
-        if self.step > 0:
-            return self.value < self.stop
-        else:
-            return self.value > self.stop
-
-    def get(self) -> int:
-        return self.value
-
-    def advance(self):
-        self.value += self.step
+    def next(self) -> Maybe[int]:
+        has_next = self.value < self.stop if self.step > 0 else self.value > self.stop
+        if has_next:
+            current = self.value
+            self.value += self.step
+            return Some(current)
+        return Nothing
 
 
 @meta_fn
-def range_or_tuple(start: Num, stop: Num | None = None, step: Num = 1) -> Range | tuple[Num, ...]:
+def range_or_tuple(start: int, stop: int | None = None, step: int = 1) -> Range | tuple[int, ...]:
     if stop is None:
         start, stop = 0, start
     if not ctx():

sonolus/script/internal/simulation_context.py

@@ -126,6 +126,6 @@ class SimulationContext:
             SimulationContext._active_context = None
 
 
-def sim_ctx() -> SimulationContext | None:
+def sim_ctx() -> SimulationContext | Any:
     """Get the current simulation context, or None if not active."""
     return SimulationContext._active_context

sonolus/script/internal/transient.py

@@ -36,11 +36,11 @@ class TransientValue(Value):
     def _get_(self) -> Self:
         return self
 
-    def _set_(self, value: Self) -> None:
+    def _set_(self, value: Any) -> None:
         if value is not self:
             raise TypeError(f"{type(self).__name__} is immutable")
 
-    def _copy_from_(self, value: Self):
+    def _copy_from_(self, value: Any):
         raise TypeError(f"{type(self).__name__} is immutable")
 
     def _copy_(self) -> Self:

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
@@ -122,7 +123,7 @@ class Value:
         For instance:
         ```
         class X(Record):
-            v: Num
+            v: int
 
         a = 1
         b = X(a) # (1) _get_() is called on a
@@ -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)
 
@@ -36,7 +36,7 @@ class Interval(Record):
 
     @property
     def is_empty(self) -> bool:
-        """Whether the has a start greater than its end."""
+        """Whether the interval has a start greater than its end."""
         return self.start > self.end
 
     @property
@@ -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:
@@ -377,7 +377,7 @@ def interp_clamped(
     xp: ArrayLike[float] | tuple[float, ...],
     fp: ArrayLike[float] | tuple[float, ...],
     x: float,
-):
+) -> float:
     """Linearly interpolate a value within a sequence of points.
 
     The sequence must have at least 2 elements and be sorted in increasing order of x-coordinates.