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

sonolus/script/internal/builtin_impls.py

@@ -1,12 +1,12 @@
-from collections.abc import Iterable
-from typing import overload
-
+from sonolus.backend.ops import Op
 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.dict_impl import DictImpl
 from sonolus.script.internal.impl import meta_fn, validate_value
 from sonolus.script.internal.math_impls import MATH_BUILTIN_IMPLS, _trunc
+from sonolus.script.internal.native import native_function
 from sonolus.script.internal.random import RANDOM_BUILTIN_IMPLS
 from sonolus.script.internal.range import Range
 from sonolus.script.internal.tuple_impl import TupleImpl
@@ -21,6 +21,8 @@ from sonolus.script.iterator import (
 )
 from sonolus.script.num import Num, _is_num
 
+_empty = object()
+
 
 @meta_fn
 def _isinstance(value, type_):
@@ -44,7 +46,7 @@ def _len(value):
     value = validate_value(value)
     if not hasattr(value, "__len__"):
         raise TypeError(f"object of type '{type(value).__name__}' has no len()")
-    return compile_and_call(value.__len__)
+    return compile_and_call(value.__len__)  # type: ignore
 
 
 @meta_fn
@@ -59,7 +61,7 @@ def _enumerate(iterable, start=0):
     elif isinstance(iterable, ArrayLike):
         return compile_and_call(iterable._enumerate_, start)
     else:
-        iterator = compile_and_call(iterable.__iter__)
+        iterator = compile_and_call(iterable.__iter__)  # type: ignore
         if not isinstance(iterator, SonolusIterator):
             raise TypeError("Only subclasses of SonolusIterator are supported as iterators")
         return _Enumerator(0, start, iterator)
@@ -104,25 +106,20 @@ def _abs(value):
     value = validate_value(value)
     if not hasattr(value, "__abs__"):
         raise TypeError(f"bad operand type for abs(): '{type(value).__name__}'")
-    return compile_and_call(value.__abs__)
+    return compile_and_call(value.__abs__)  # type: ignore
 
 
 def _identity(value):
     return value
 
 
-@overload
-def _max[T](iterable: Iterable[T], *, key: callable = ...) -> T: ...
-
-
-@overload
-def _max[T](a: T, b: T, *args: T, key: callable = ...) -> T: ...
-
-
 @meta_fn
-def _max(*args, key: callable = _identity):
+def _max(*args, default=_empty, key=None):
     from sonolus.backend.visitor import compile_and_call
 
+    if key is None:
+        key = _identity
+
     args = tuple(validate_value(arg) for arg in args)
     if len(args) == 0:
         raise ValueError("Expected at least one argument to max")
@@ -131,40 +128,91 @@ def _max(*args, key: callable = _identity):
         if isinstance(iterable, ArrayLike):
             return compile_and_call(iterable._max_, key=key)
         elif isinstance(iterable, TupleImpl) and all(_is_num(v) for v in iterable.value):
+            if len(iterable.value) == 0:
+                if default is not _empty:
+                    return default
+                raise ValueError("max() arg is an empty sequence")
             return compile_and_call(Array(*iterable.value)._max_, key=key)
+        elif isinstance(iterable, SonolusIterator):
+            if not (default is _empty or Num._accepts_(default)):
+                raise TypeError("default argument must be a number")
+            return compile_and_call(
+                _max_num_iterator,
+                iterable,
+                Num._accept_(default) if default is not _empty else None,
+                key=key if key is not _identity else None,
+            )
         else:
             raise TypeError(f"Unsupported type: {type(iterable)} for max")
     else:
+        if default is not _empty:
+            raise TypeError("default argument is not supported for max with multiple arguments")
         if not all(_is_num(arg) for arg in args):
             raise TypeError("Arguments to max must be numbers")
         if ctx():
-            result = compile_and_call(_max2, args[0], args[1], key=key)
+            result = _max2(args[0], args[1], key=key)
             for arg in args[2:]:
-                result = compile_and_call(_max2, result, arg, key=key)
+                result = _max2(result, arg, key=key)
             return result
         else:
             return max(arg._as_py_() for arg in args)
 
 
 def _max2(a, b, key=_identity):
-    if key(a) > key(b):
+    from sonolus.backend.visitor import compile_and_call
+
+    a = validate_value(a)
+    b = validate_value(b)
+    if _is_num(a) and _is_num(b) and key == _identity:
+        return compile_and_call(_max2_num, a, b)
+    return compile_and_call(_max2_generic, a, b, key=key)
+
+
+@native_function(Op.Max)
+def _max2_num(a, b):
+    if a > b:
         return a
     else:
         return b
 
 
-@overload
-def _min[T](iterable: Iterable[T], *, key: callable = ...) -> T: ...
+def _max2_generic(a, b, key=_identity):
+    if key(a) > key(b):
+        return a
+    else:
+        return b
 
 
-@overload
-def _min[T](a: T, b: T, *args: T, key: callable = ...) -> T: ...
+def _max_num_iterator(iterable, default, key):
+    iterator = iterable.__iter__()  # noqa: PLC2801
+    initial = iterator.next()
+    if initial.is_nothing:
+        assert default is not None
+        return default
+    if key is not None:
+        result = initial.get_unsafe()
+        best_key = key(result)
+        for value in iterator:
+            new_key = key(value)
+            if new_key > best_key:
+                result = value
+                best_key = new_key
+        return result
+    else:
+        result = initial.get_unsafe()
+        for value in iterator:
+            if value > result:  # noqa: PLR1730
+                result = value
+        return result
 
 
 @meta_fn
-def _min(*args, key: callable = _identity):
+def _min(*args, default=_empty, key=None):
     from sonolus.backend.visitor import compile_and_call
 
+    if key is None:
+        key = _identity
+
     args = tuple(validate_value(arg) for arg in args)
     if len(args) == 0:
         raise ValueError("Expected at least one argument to min")
@@ -173,34 +221,92 @@ def _min(*args, key: callable = _identity):
         if isinstance(iterable, ArrayLike):
             return compile_and_call(iterable._min_, key=key)
         elif isinstance(iterable, TupleImpl) and all(_is_num(v) for v in iterable.value):
+            if len(iterable.value) == 0:
+                if default is not _empty:
+                    return default
+                raise ValueError("min() arg is an empty sequence")
             return compile_and_call(Array(*iterable.value)._min_, key=key)
+        elif isinstance(iterable, SonolusIterator):
+            if not (default is _empty or Num._accepts_(default)):
+                raise TypeError("default argument must be a number")
+            return compile_and_call(
+                _min_num_iterator,
+                iterable,
+                Num._accept_(default) if default is not _empty else None,
+                key=key if key is not _identity else None,
+            )
         else:
             raise TypeError(f"Unsupported type: {type(iterable)} for min")
     else:
+        if default is not _empty:
+            raise TypeError("default argument is not supported for min with multiple arguments")
         if not all(_is_num(arg) for arg in args):
             raise TypeError("Arguments to min must be numbers")
         if ctx():
-            result = compile_and_call(_min2, args[0], args[1], key=key)
+            result = _min2(args[0], args[1], key=key)
             for arg in args[2:]:
-                result = compile_and_call(_min2, result, arg, key=key)
+                result = _min2(result, arg, key=key)
             return result
         else:
             return min(arg._as_py_() for arg in args)
 
 
 def _min2(a, b, key=_identity):
+    from sonolus.backend.visitor import compile_and_call
+
+    a = validate_value(a)
+    b = validate_value(b)
+    if _is_num(a) and _is_num(b) and key == _identity:
+        return compile_and_call(_min2_num, a, b)
+    return compile_and_call(_min2_generic, a, b, key=key)
+
+
+@native_function(Op.Min)
+def _min2_num(a, b):
+    if a < b:
+        return a
+    else:
+        return b
+
+
+def _min2_generic(a, b, key=_identity):
     if key(a) < key(b):
         return a
     else:
         return b
 
 
+def _min_num_iterator(iterable, default, key):
+    iterator = iterable.__iter__()  # noqa: PLC2801
+    initial = iterator.next()
+    if initial.is_nothing:
+        assert default is not None
+        return default
+    if key is not None:
+        result = initial.get_unsafe()
+        best_key = key(result)
+        for value in iterator:
+            new_key = key(value)
+            if new_key < best_key:
+                result = value
+                best_key = new_key
+        return result
+    else:
+        result = initial.get_unsafe()
+        for value in iterator:
+            if value < result:  # noqa: PLR1730
+                result = value
+        return result
+
+
 @meta_fn
 def _callable(value):
     return callable(value)
 
 
 def _map(fn, iterable, *iterables):
+    if len(iterables) == 0:
+        return _MappingIterator(fn, iterable.__iter__())  # noqa: PLC2801
     return _MappingIterator(lambda args: fn(*args), zip(iterable, *iterables))  # noqa: B905
 
 
@@ -234,9 +340,9 @@ def _bool(value=False):
         return False
 
 
-_int._type_mapping_ = Num
-_float._type_mapping_ = Num
-_bool._type_mapping_ = Num
+_int._type_mapping_ = Num  # type: ignore
+_float._type_mapping_ = Num  # type: ignore
+_bool._type_mapping_ = Num  # type: ignore
 
 
 def _any(iterable):
@@ -253,12 +359,31 @@ def _all(iterable):
     return True
 
 
+def _sum(iterable, /, start=0):
+    for value in iterable:
+        start += value
+    return start
+
+
+def _next(iterator):
+    assert isinstance(iterator, SonolusIterator)
+    value = iterator.next()
+    if value.is_some:
+        return value.get_unsafe()
+    error("Iterator has been exhausted")
+
+
+def _iter(iterable):
+    return iterable.__iter__()  # type: ignore # noqa: PLC2801
+
+
 # classmethod, property, staticmethod are supported as decorators, but not within functions
 
 BUILTIN_IMPLS = {
     id(abs): _abs,
     id(all): _all,
     id(any): _any,
+    id(sum): _sum,
     id(bool): _bool,
     id(callable): _callable,
     id(enumerate): _enumerate,
@@ -266,10 +391,12 @@ BUILTIN_IMPLS = {
     id(float): _float,
     id(int): _int,
     id(isinstance): _isinstance,
+    id(iter): _iter,
     id(len): _len,
     id(map): _map,
     id(max): _max,
     id(min): _min,
+    id(next): _next,
     id(range): Range,
     id(reversed): _reversed,
     id(zip): _zip,

sonolus/script/internal/constant.py

@@ -48,8 +48,8 @@ class ConstantValue(Value):
         class Parameterized(cls):
             _value = (parameter,)
 
-        Parameterized.__name__ = f"{parameter}"
-        Parameterized.__qualname__ = f"{parameter}"
+        Parameterized.__name__ = f"Const[{object.__repr__(parameter)}]"  # noqa: PLC2801
+        Parameterized.__qualname__ = Parameterized.__name__
         Parameterized.__module__ = cls.__module__
         Parameterized.instance = object.__new__(Parameterized)
         return Parameterized
@@ -113,7 +113,7 @@ class ConstantValue(Value):
         if value is not self:
             raise ValueError(f"{type(self).__name__} is immutable")
 
-    def _copy_from_(self, value: Self):
+    def _copy_from_(self, value: Any):
         if value is not self:
             raise ValueError(f"{type(self).__name__} is immutable")
 
@@ -143,3 +143,13 @@ class ConstantValue(Value):
 
 class BasicConstantValue(ConstantValue):
     """For constants without any special behavior."""
+
+
+class TypingSpecialFormConstant(ConstantValue):
+    """For constants that are typing special forms that have a [] operator."""
+
+    @meta_fn
+    def __getitem__(self, item: Any) -> Self:
+        if not item._is_py_():
+            raise TypeError(f"Invalid value for type parameter: {item}")
+        return self.value()[item._as_py_()]

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/native.py

@@ -33,8 +33,8 @@ def native_function[**P, R](op: Op) -> Callable[[Callable[P, R]], Callable[P, R]
                 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,6 +2,7 @@ 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
 
@@ -70,17 +71,13 @@ 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

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: