sonolus.py 0.1.3__py3-none-any.whl → 0.1.5__py3-none-any.whl

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):
@@ -164,7 +173,8 @@ class Context:
                 header.scope.set_value(name, target_value)
                 header.loop_variables[name] = target_value
             else:
-                header.scope.set_binding(name, ConflictBinding())
+                header.scope.set_value(name, value)
+                header.loop_variables[name] = value
         return header
 
     def branch_to_loop_header(self, header: Self):
@@ -174,16 +184,23 @@ class Context:
         self.outgoing[None] = header
         for name, target_value in header.loop_variables.items():
             with using_ctx(self):
-                value = self.scope.get_value(name)
-                value = type(target_value)._accept_(value)
-                target_value._set_(value)
+                if type(target_value)._is_value_type_():
+                    value = self.scope.get_value(name)
+                    value = type(target_value)._accept_(value)
+                    target_value._set_(value)
+                else:
+                    value = self.scope.get_value(name)
+                    if target_value is not value:
+                        raise RuntimeError(
+                            f"Variable '{name}' may have conflicting definitions between loop iterations"
+                        )
 
     def map_constant(self, value: Any) -> int:
         if value not in self.const_mappings:
             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/descriptor.py

@@ -1,17 +1,17 @@
-from abc import abstractmethod
-
-
-class SonolusDescriptor:
-    """Base class for Sonolus descriptors.
-
-    The compiler checks if a descriptor is an instance of a subclass of this class,
-    so it knows that it's a supported descriptor.
-    """
-
-    @abstractmethod
-    def __get__(self, instance, owner):
-        pass
-
-    @abstractmethod
-    def __set__(self, instance, value):
-        pass
+from abc import abstractmethod
+
+
+class SonolusDescriptor:
+    """Base class for Sonolus descriptors.
+
+    The compiler checks if a descriptor is an instance of a subclass of this class,
+    so it knows that it's a supported descriptor.
+    """
+
+    @abstractmethod
+    def __get__(self, instance, owner):
+        pass
+
+    @abstractmethod
+    def __set__(self, instance, value):
+        pass

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

sonolus/script/internal/introspection.py

@@ -1,14 +1,17 @@
-import inspect
-from typing import Annotated
-
-_missing = object()
-
-
-def get_field_specifiers(cls, *, globals=None, locals=None, eval_str=True):  # noqa: A002
-    """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():
-        class_value = getattr(cls, key, _missing)
-        if class_value is not _missing:
-            results[key] = Annotated[value, class_value]
-    return results
+import inspect
+from typing import Annotated
+
+_missing = object()
+
+
+def get_field_specifiers(cls, *, skip: set[str] = frozenset(), globals=None, locals=None, eval_str=True):  # noqa: A002
+    """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():
+        class_value = getattr(cls, key, _missing)
+        if class_value is not _missing and key not in skip:
+            results[key] = Annotated[value, class_value]
+    for key, value in cls.__dict__.items():
+        if key not in results and key not in skip and not key.startswith("__") and not callable(value):
+            raise ValueError(f"Missing annotation for {cls.__name__}.{key}")
+    return results

sonolus/script/internal/math_impls.py

@@ -0,0 +1,121 @@
+import math
+
+from sonolus.backend.ops import Op
+from sonolus.script.internal.native import native_function
+
+
+@native_function(Op.Sin)
+def _sin(x: float) -> float:
+    return math.sin(x)
+
+
+@native_function(Op.Cos)
+def _cos(x: float) -> float:
+    return math.cos(x)
+
+
+@native_function(Op.Tan)
+def _tan(x: float) -> float:
+    return math.tan(x)
+
+
+@native_function(Op.Arcsin)
+def _asin(x: float) -> float:
+    return math.asin(x)
+
+
+@native_function(Op.Arccos)
+def _acos(x: float) -> float:
+    return math.acos(x)
+
+
+@native_function(Op.Arctan)
+def _atan(x: float) -> float:
+    return math.atan(x)
+
+
+@native_function(Op.Arctan2)
+def _atan2(y: float, x: float) -> float:
+    return math.atan2(y, x)
+
+
+@native_function(Op.Sinh)
+def _sinh(x: float) -> float:
+    return math.sinh(x)
+
+
+@native_function(Op.Cosh)
+def _cosh(x: float) -> float:
+    return math.cosh(x)
+
+
+@native_function(Op.Tanh)
+def _tanh(x: float) -> float:
+    return math.tanh(x)
+
+
+@native_function(Op.Floor)
+def _floor(x: float) -> float:
+    return math.floor(x)
+
+
+@native_function(Op.Ceil)
+def _ceil(x: float) -> float:
+    return math.ceil(x)
+
+
+@native_function(Op.Trunc)
+def _trunc(x: float) -> float:
+    return math.trunc(x)
+
+
+@native_function(Op.Round)
+def __round(x: float) -> float:
+    return round(x)
+
+
+def _round(x: float, n: int = 0) -> float:
+    if n == 0:
+        return __round(x)
+    return __round(x * 10**n) / 10**n
+
+
+@native_function(Op.Frac)
+def frac(x: float) -> float:
+    return x % 1
+
+
+@native_function(Op.Log)
+def _ln(x: float) -> float:
+    return math.log(x)
+
+
+def _log(x: float, base: float | None = None) -> float:
+    if base is None:
+        return _ln(x)
+    return _ln(x) / _ln(base)
+
+
+@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
+    return math.copysign(abs(x) % abs(y), x)
+
+
+MATH_BUILTIN_IMPLS = {
+    id(math.sin): _sin,
+    id(math.cos): _cos,
+    id(math.tan): _tan,
+    id(math.asin): _asin,
+    id(math.acos): _acos,
+    id(math.atan): _atan,
+    id(math.atan2): _atan2,
+    id(math.sinh): _sinh,
+    id(math.cosh): _cosh,
+    id(math.tanh): _tanh,
+    id(math.floor): _floor,
+    id(math.ceil): _ceil,
+    id(math.trunc): _trunc,
+    id(round): _round,
+    id(math.log): _log,
+}