sonolus.py 0.1.4__py3-none-any.whl → 0.1.6__py3-none-any.whl

sonolus/script/internal/builtin_impls.py

@@ -1,18 +1,39 @@
+from abc import ABC
 from collections.abc import Iterable
 from typing import overload
 
+from sonolus.script.array_like import ArrayLike
 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.iterator import ArrayLike, Enumerator, SonolusIterator
-from sonolus.script.math import MATH_BUILTIN_IMPLS
-from sonolus.script.num import is_num
-from sonolus.script.range import Range
+from sonolus.script.internal.math_impls import MATH_BUILTIN_IMPLS, _trunc
+from sonolus.script.internal.random import RANDOM_BUILTIN_IMPLS
+from sonolus.script.internal.range import Range
+from sonolus.script.internal.tuple_impl import TupleImpl
+from sonolus.script.internal.value import Value
+from sonolus.script.iterator import (
+    SonolusIterator,
+    _EmptyIterator,
+    _Enumerator,
+    _FilteringIterator,
+    _MappingIterator,
+    _Zipper,
+)
+from sonolus.script.num import Num, _is_num
 
 
 @meta_fn
 def _isinstance(value, type_):
     value = validate_value(value)
     type_ = validate_value(type_)._as_py_()
+    if type_ is dict:
+        return isinstance(value, DictImpl)
+    if type_ is tuple:
+        return isinstance(value, TupleImpl)
+    if type_ in {_int, _float, _bool}:
+        raise TypeError("Instance check against int, float, or bool is not supported, use Num instead")
+    if not (isinstance(type_, type) and issubclass(type_, Value | ABC)):
+        raise TypeError(f"Unsupported type: {type_} for isinstance")
     return validate_value(isinstance(value, type_))
 
 
@@ -34,12 +55,44 @@ def _enumerate(iterable, start=0):
     if not hasattr(iterable, "__iter__"):
         raise TypeError(f"'{type(iterable).__name__}' object is not iterable")
     if isinstance(iterable, ArrayLike):
-        return compile_and_call(iterable.enumerate, start)
+        return compile_and_call(iterable._enumerate_, start)
     else:
         iterator = compile_and_call(iterable.__iter__)
         if not isinstance(iterator, SonolusIterator):
             raise TypeError("Only subclasses of SonolusIterator are supported as iterators")
-        return Enumerator(0, start, iterator)
+        return _Enumerator(0, start, iterator)
+
+
+@meta_fn
+def _reversed(iterable):
+    from sonolus.backend.visitor import compile_and_call
+
+    iterable = validate_value(iterable)
+    if not isinstance(iterable, ArrayLike):
+        raise TypeError(f"Unsupported type: {type(iterable)} for reversed")
+    return compile_and_call(iterable.__reversed__)
+
+
+@meta_fn
+def _zip(*iterables):
+    from sonolus.backend.visitor import compile_and_call
+    from sonolus.script.containers import Pair
+
+    if not iterables:
+        return _EmptyIterator()
+
+    iterables = [validate_value(iterable) for iterable in iterables]
+    if any(isinstance(iterable, TupleImpl) for iterable in iterables):
+        if not all(isinstance(iterable, TupleImpl) for iterable in iterables):
+            raise TypeError("Cannot mix tuples with other types in zip")
+        return TupleImpl._accept_(tuple(zip(*(iterable._as_py_() for iterable in iterables), strict=False)))
+    iterators = [compile_and_call(iterable.__iter__) for iterable in iterables]
+    if not all(isinstance(iterator, SonolusIterator) for iterator in iterators):
+        raise TypeError("Only subclasses of SonolusIterator are supported as iterators")
+    v = iterators.pop()
+    while iterators:
+        v = Pair(iterators.pop(), v)
+    return _Zipper(v)
 
 
 @meta_fn
@@ -52,16 +105,20 @@ def _abs(value):
     return compile_and_call(value.__abs__)
 
 
+def _identity(value):
+    return value
+
+
 @overload
-def _max[T](iterable: Iterable[T]) -> T: ...
+def _max[T](iterable: Iterable[T], *, key: callable = ...) -> T: ...
 
 
 @overload
-def _max[T](a: T, b: T, *args: T) -> T: ...
+def _max[T](a: T, b: T, *args: T, key: callable = ...) -> T: ...
 
 
 @meta_fn
-def _max(*args):
+def _max(*args, key: callable = _identity):
     from sonolus.backend.visitor import compile_and_call
 
     args = tuple(validate_value(arg) for arg in args)
@@ -70,38 +127,38 @@ def _max(*args):
     elif len(args) == 1:
         (iterable,) = args
         if isinstance(iterable, ArrayLike):
-            return iterable.max()
+            return compile_and_call(iterable._max_, key=key)
         else:
             raise TypeError(f"Unsupported type: {type(iterable)} for max")
     else:
-        if not all(is_num(arg) for arg in args):
+        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])
+            result = compile_and_call(_max2, args[0], args[1], key=key)
             for arg in args[2:]:
-                result = compile_and_call(_max2, result, arg)
+                result = compile_and_call(_max2, result, arg, key=key)
             return result
         else:
             return max(arg._as_py_() for arg in args)
 
 
-def _max2(a, b):
-    if a > b:
+def _max2(a, b, key=_identity):
+    if key(a) > key(b):
         return a
     else:
         return b
 
 
 @overload
-def _min[T](iterable: Iterable[T]) -> T: ...
+def _min[T](iterable: Iterable[T], *, key: callable = ...) -> T: ...
 
 
 @overload
-def _min[T](a: T, b: T, *args: T) -> T: ...
+def _min[T](a: T, b: T, *args: T, key: callable = ...) -> T: ...
 
 
 @meta_fn
-def _min(*args):
+def _min(*args, key: callable = _identity):
     from sonolus.backend.visitor import compile_and_call
 
     args = tuple(validate_value(arg) for arg in args)
@@ -110,35 +167,89 @@ def _min(*args):
     elif len(args) == 1:
         (iterable,) = args
         if isinstance(iterable, ArrayLike):
-            return iterable.min()
+            return compile_and_call(iterable._min_, key=key)
         else:
             raise TypeError(f"Unsupported type: {type(iterable)} for min")
     else:
-        if not all(is_num(arg) for arg in args):
+        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])
+            result = compile_and_call(_min2, args[0], args[1], key=key)
             for arg in args[2:]:
-                result = compile_and_call(_min2, result, arg)
+                result = compile_and_call(_min2, result, arg, key=key)
             return result
         else:
             return min(arg._as_py_() for arg in args)
 
 
-def _min2(a, b):
-    if a < b:
+def _min2(a, b, key=_identity):
+    if key(a) < key(b):
         return a
     else:
         return b
 
 
+@meta_fn
+def _callable(value):
+    return callable(value)
+
+
+def _map(fn, iterable, *iterables):
+    return _MappingIterator(lambda args: fn(*args), zip(iterable, *iterables))  # noqa: B905
+
+
+def _filter(fn, iterable):
+    if fn is None:
+        fn = _identity
+    return _FilteringIterator(fn, iterable.__iter__())  # noqa: PLC2801
+
+
+@meta_fn
+def _int(value=0):
+    value = validate_value(value)
+    if not _is_num(value):
+        raise TypeError("Only numeric arguments to int() are supported")
+    return _trunc(value)
+
+
+@meta_fn
+def _float(value=0.0):
+    value = validate_value(value)
+    if not _is_num(value):
+        raise TypeError("Only numeric arguments to float() are supported")
+    return value
+
+
+@meta_fn
+def _bool(value=False):
+    value = validate_value(value)
+    if not _is_num(value):
+        raise TypeError("Only numeric arguments to bool() are supported")
+    return value != 0
+
+
+_int._type_mapping_ = Num
+_float._type_mapping_ = Num
+_bool._type_mapping_ = Num
+
+# classmethod, property, staticmethod are supported as decorators, but not within functions
+
 BUILTIN_IMPLS = {
+    id(abs): _abs,
+    id(bool): _bool,
+    id(callable): _callable,
+    id(enumerate): _enumerate,
+    id(filter): _filter,
+    id(float): _float,
+    id(int): _int,
     id(isinstance): _isinstance,
     id(len): _len,
-    id(enumerate): _enumerate,
-    id(abs): _abs,
+    id(map): _map,
     id(max): _max,
     id(min): _min,
     id(range): Range,
-    **MATH_BUILTIN_IMPLS,
+    id(reversed): _reversed,
+    id(zip): _zip,
+    **MATH_BUILTIN_IMPLS,  # Includes round
+    **RANDOM_BUILTIN_IMPLS,
 }

sonolus/script/internal/constant.py

@@ -0,0 +1,139 @@
+from collections.abc import Iterable
+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
+
+
+class _Missing:
+    def __repr__(self) -> str:
+        return "MISSING"
+
+    def __bool__(self) -> bool:
+        return False
+
+
+_MISSING = _Missing()
+
+
+class ConstantValue(Value):
+    _parameterized_: ClassVar[dict[Any, type[Self]]] = {}
+    _value: ClassVar[Any] = _MISSING
+    instance: ClassVar[Self | _Missing] = _MISSING
+
+    def __new__(cls) -> Self:
+        if cls.value() is _MISSING:
+            raise TypeError(f"Class {cls.__name__} is not parameterized")
+        return cls.instance
+
+    @classmethod
+    def value(cls):
+        # We need this to avoid descriptors getting in the way
+        return cls._value[0] if cls._value is not _MISSING else _MISSING
+
+    @classmethod
+    def of(cls, value: Any) -> Self:
+        if value in cls._parameterized_:
+            return cls._parameterized_[value]()
+
+        parameterized = cls._get_parameterized(value)
+        cls._parameterized_[value] = parameterized
+        return parameterized()
+
+    @classmethod
+    def _get_parameterized(cls, parameter: Any) -> type[Self]:
+        class Parameterized(cls):
+            _value = (parameter,)
+
+        Parameterized.__name__ = f"{parameter}"
+        Parameterized.__qualname__ = f"{parameter}"
+        Parameterized.__module__ = cls.__module__
+        Parameterized.instance = object.__new__(Parameterized)
+        return Parameterized
+
+    @classmethod
+    def _is_concrete_(cls) -> bool:
+        return True
+
+    @classmethod
+    def _size_(cls) -> int:
+        return 0
+
+    @classmethod
+    def _is_value_type_(cls) -> bool:
+        return False
+
+    @classmethod
+    def _from_place_(cls, place: BlockPlace) -> Self:
+        if cls.value() is _MISSING:
+            raise TypeError(f"Class {cls.__name__} is not parameterized")
+        return cls()
+
+    @classmethod
+    def _accepts_(cls, value: Any) -> bool:
+        from sonolus.script.internal.impl import validate_value
+
+        # We rely on validate_value to create the correct instance
+        return isinstance(validate_value(value), cls)
+
+    @classmethod
+    def _accept_(cls, value: Any) -> Self:
+        from sonolus.script.internal.impl import validate_value
+
+        # We rely on validate_value to create the correct instance
+        value = validate_value(value)
+        if not isinstance(value, cls):
+            raise ValueError(f"Value {value} is not of type {cls}")
+        return value
+
+    def _is_py_(self) -> bool:
+        return True
+
+    def _as_py_(self) -> Any:
+        return self.value()
+
+    @classmethod
+    def _from_list_(cls, values: Iterable[float | BlockPlace]) -> Self:
+        return cls()
+
+    def _to_list_(self, level_refs: dict[Any, int] | None = None) -> list[float | BlockPlace]:
+        return []
+
+    @classmethod
+    def _flat_keys_(cls, prefix: str) -> list[str]:
+        return []
+
+    def _get_(self) -> Self:
+        return self
+
+    def _set_(self, value: Any) -> Self:
+        if value is not self:
+            raise ValueError(f"{type(self).__name__} is immutable")
+
+    def _copy_from_(self, value: Self):
+        if value is not self:
+            raise ValueError(f"{type(self).__name__} is immutable")
+
+    def _copy_(self) -> Self:
+        return self
+
+    @classmethod
+    def _alloc_(cls) -> Self:
+        return cls()
+
+    @meta_fn
+    def __eq__(self, other):
+        return self is other
+
+    @meta_fn
+    def __ne__(self, other):
+        return self is not other
+
+    @meta_fn
+    def __hash__(self):
+        return hash(self.value())
+
+
+class MiscConstantValue(ConstantValue):
+    """For constants without any special behavior."""

sonolus/script/internal/context.py

@@ -6,11 +6,11 @@ from dataclasses import dataclass
 from typing import Any, Self
 
 from sonolus.backend.blocks import BlockData, PlayBlock
-from sonolus.backend.flow import BasicBlock, FlowEdge
 from sonolus.backend.ir import IRConst, IRStmt
 from sonolus.backend.mode import Mode
+from sonolus.backend.optimize.flow import BasicBlock, FlowEdge
 from sonolus.backend.place import Block, BlockPlace, TempBlock
-from sonolus.script.globals import GlobalInfo, GlobalPlaceholder
+from sonolus.script.globals import _GlobalInfo, _GlobalPlaceholder
 from sonolus.script.internal.value import Value
 
 _compiler_internal_ = True
@@ -22,7 +22,7 @@ class GlobalContextState:
     archetypes: dict[type, int]
     rom: ReadOnlyMemory
     const_mappings: dict[Any, int]
-    environment_mappings: dict[GlobalInfo, int]
+    environment_mappings: dict[_GlobalInfo, int]
     environment_offsets: dict[Block, int]
     mode: Mode
 
@@ -60,6 +60,7 @@ class Context:
         callback_state: CallbackContextState,
         scope: Scope | None = None,
         live: bool = True,
+        foldable_constants: dict[TempBlock, list[float | None]] | None = None,
     ):
         self.global_state = global_state
         self.callback_state = callback_state
@@ -69,6 +70,7 @@ class Context:
         self.scope = scope if scope is not None else Scope()
         self.loop_variables = {}
         self.live = live
+        self.foldable_constants = foldable_constants if foldable_constants is not None else {}
 
     @property
     def rom(self) -> ReadOnlyMemory:
@@ -102,8 +104,14 @@ class Context:
         if isinstance(place.block, BlockData) and self.callback not in self.blocks(place.block).writable:
             raise RuntimeError(f"Block {place.block} is not writable in {self.callback}")
 
+    def add_statement(self, statement: IRStmt):
+        if not self.live:
+            return
+        self.statements.append(statement)
+
     def add_statements(self, *statements: IRStmt):
-        self.statements.extend(statements)
+        for statement in statements:
+            self.add_statement(statement)
 
     def alloc(self, name: str | None = None, size: int = 1) -> BlockPlace:
         if size == 0:
@@ -124,6 +132,7 @@ class Context:
             callback_state=self.callback_state,
             scope=scope,
             live=self.live,
+            foldable_constants={**self.foldable_constants},  # We want this to persist even if the scope changes
         )
 
     def branch(self, condition: float | None):
@@ -191,7 +200,7 @@ class Context:
             self.const_mappings[value] = len(self.const_mappings)
         return self.const_mappings[value]
 
-    def get_global_base(self, value: GlobalInfo | GlobalPlaceholder) -> BlockPlace:
+    def get_global_base(self, value: _GlobalInfo | _GlobalPlaceholder) -> BlockPlace:
         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")

sonolus/script/internal/dict_impl.py

@@ -0,0 +1,65 @@
+from typing import Any, Self
+
+from sonolus.script.internal.impl import meta_fn, validate_value
+from sonolus.script.internal.transient import TransientValue
+
+
+class DictImpl(TransientValue):
+    def __init__(self, value: dict):
+        self.value = value
+
+    @meta_fn
+    def __getitem__(self, item):
+        item = validate_value(item)
+        if not item._is_py_():
+            raise TypeError("Key must be a compile-time constant")
+        item = item._as_py_()
+        if item not in self.value:
+            raise KeyError(item)
+        return self.value[item]
+
+    @meta_fn
+    def __contains__(self, item):
+        item = validate_value(item)
+        if not item._is_py_():
+            raise TypeError("Key must be a compile-time constant")
+        item = item._as_py_()
+        return item in self.value
+
+    @meta_fn
+    def __len__(self):
+        return len(self.value)
+
+    def __eq__(self, other):
+        raise TypeError("Comparing dicts is not supported")
+
+    __hash__ = None
+
+    @meta_fn
+    def __or__(self, other):
+        if not isinstance(other, DictImpl):
+            raise TypeError("Only dicts can be merged")
+        return DictImpl({**self.value, **other.value})
+
+    @classmethod
+    def _accepts_(cls, value: Any) -> bool:
+        return isinstance(value, cls | dict)
+
+    @classmethod
+    def _accept_(cls, value: Any) -> Self:
+        if not cls._accepts_(value):
+            raise TypeError(f"Cannot accept {value} as {cls.__name__}")
+        if isinstance(value, cls):
+            return value
+        else:
+            return cls({validate_value(k)._as_py_(): validate_value(v) for k, v in value.items()})
+
+    def _is_py_(self) -> bool:
+        return all(v._is_py_() for v in self.value.values())
+
+    def _as_py_(self) -> Any:
+        return {k: v._as_py_() for k, v in self.value.items()}
+
+
+DictImpl.__name__ = "dict"
+DictImpl.__qualname__ = "dict"

sonolus/script/internal/generic.py

@@ -3,19 +3,19 @@ from __future__ import annotations
 from enum import Enum
 from typing import Annotated, Any, ClassVar, Literal, Self, TypeVar, get_origin
 
-from sonolus.script.internal.impl import meta_fn
+from sonolus.script.internal.impl import meta_fn, validate_value
 from sonolus.script.internal.value import Value
 
 type AnyType = type[Value] | PartialGeneric | TypeVar
 
 
 def validate_type_arg(arg: Any) -> Any:
-    from sonolus.script.internal.impl import validate_value
-
     arg = validate_value(arg)
     if not arg._is_py_():
         raise TypeError(f"Expected a compile-time constant type argument, got {arg}")
     result = arg._as_py_()
+    if hasattr(result, "_type_mapping_"):
+        return result._type_mapping_
     if get_origin(result) is Annotated:
         return result.__args__[0]
     if get_origin(result) is Literal:
@@ -69,7 +69,7 @@ class GenericValue(Value):
             raise TypeError(f"Missing type arguments for {self.__class__.__name__}")
 
     @classmethod
-    def _validate__type_args_(cls, args: tuple[Any, ...]) -> tuple[Any, ...]:
+    def _validate_type_args_(cls, args: tuple[Any, ...]) -> tuple[Any, ...]:
         """Validate the type arguments and return them as a tuple.
 
         This may be called with PartialGeneric or TypeVar instances inside args.
@@ -84,7 +84,7 @@ class GenericValue(Value):
 
     @classmethod
     @meta_fn
-    def _get_type_arg_(cls, var: TypeVar) -> Any:
+    def type_var_value(cls, var: TypeVar) -> Any:
         if isinstance(var, Value):
             var = var._as_py_()
         if cls._type_args_ is None:
@@ -93,12 +93,11 @@ class GenericValue(Value):
             return cls._type_vars_to_args_[var]
         raise TypeError(f"Missing type argument for {var}")
 
-    @classmethod
     def __class_getitem__(cls, args: Any) -> type[Self]:
         if cls._type_args_ is not None:
             raise TypeError(f"Type {cls.__name__} is already parameterized")
         args = validate_type_args(args)
-        args = cls._validate__type_args_(args)
+        args = cls._validate_type_args_(args)
         if contains_incomplete_type(args):
             return PartialGeneric(cls, args)
         if args not in cls._parameterized_:
@@ -168,8 +167,6 @@ def infer_and_validate_types(dst: Any, src: Any, results: dict[TypeVar, Any] | N
 
 
 def accept_and_infer_types(dst: Any, val: Any, results: dict[TypeVar, Any]) -> Value:
-    from sonolus.script.internal.impl import validate_value
-
     val = validate_value(val)
     match dst:
         case TypeVar():

sonolus/script/internal/impl.py

@@ -2,7 +2,7 @@ from __future__ import annotations
 
 from collections.abc import Callable
 from enum import Enum
-from types import FunctionType, MethodType, NoneType, NotImplementedType
+from types import FunctionType, MethodType, ModuleType, NoneType, NotImplementedType, UnionType
 from typing import TYPE_CHECKING, Annotated, Any, Literal, TypeVar, get_origin, overload
 
 if TYPE_CHECKING:
@@ -18,6 +18,11 @@ def meta_fn[T: Callable]() -> Callable[[T], T]: ...
 
 
 def meta_fn(fn=None):
+    """Marks a function as a meta function to be called directly without the AST visitor.
+
+    This can also improve performance in some cases by avoiding the overhead of the AST visitor.
+    """
+
     # noinspection PyShadowingNames
     def decorator(fn):
         fn._meta_fn_ = True
@@ -36,12 +41,23 @@ def validate_value(value: Any) -> Value:
 
 
 def try_validate_value(value: Any) -> Value | None:
-    from sonolus.script.comptime import Comptime
-    from sonolus.script.globals import GlobalPlaceholder
+    from sonolus.script.globals import _GlobalPlaceholder
+    from sonolus.script.internal.constant import MiscConstantValue
+    from sonolus.script.internal.dict_impl import DictImpl
     from sonolus.script.internal.generic import PartialGeneric
+    from sonolus.script.internal.tuple_impl import TupleImpl
     from sonolus.script.internal.value import Value
     from sonolus.script.num import Num
 
+    try:
+        # Unfortunately this is called during import, so this may fail
+        from sonolus.script.internal.builtin_impls import BUILTIN_IMPLS
+
+        if id(value) in BUILTIN_IMPLS:
+            return validate_value(BUILTIN_IMPLS[id(value)])
+    except ImportError:
+        pass
+
     match value:
         case Enum():
             return validate_value(value.value)
@@ -49,21 +65,30 @@ def try_validate_value(value: Any) -> Value | None:
             return value
         case type():
             if value in {int, float, bool}:
-                return Comptime.accept_unchecked(Num)
-            return Comptime.accept_unchecked(value)
+                return MiscConstantValue.of(Num)
+            return MiscConstantValue.of(value)
         case int() | float() | bool():
             return Num._accept_(value)
         case tuple():
-            return Comptime.accept_unchecked(tuple(validate_value(v) for v in value))
+            return TupleImpl._accept_(value)
         case dict():
-            return Comptime.accept_unchecked({validate_value(k)._as_py_(): validate_value(v) for k, v in value.items()})
-        case PartialGeneric() | TypeVar() | FunctionType() | MethodType() | NotImplementedType() | str() | NoneType():
-            return Comptime.accept_unchecked(value)
-        case other_type if get_origin(value) in {Literal, Annotated}:
-            return Comptime.accept_unchecked(other_type)
-        case GlobalPlaceholder():
+            return DictImpl._accept_(value)
+        case (
+            PartialGeneric()
+            | TypeVar()
+            | FunctionType()
+            | MethodType()
+            | NotImplementedType()
+            | str()
+            | NoneType()
+            | ModuleType()
+        ):
+            return MiscConstantValue.of(value)
+        case other_type if get_origin(value) in {Literal, Annotated, UnionType, tuple}:
+            return MiscConstantValue.of(other_type)
+        case _GlobalPlaceholder():
             return value.get()
         case comptime_value if getattr(comptime_value, "_is_comptime_value_", False):
-            return Comptime.accept_unchecked(comptime_value)
+            return MiscConstantValue.of(comptime_value)
         case _:
             return None