sonolus.py 0.1.0__py3-none-any.whl

sonolus/script/internal/context.py

@@ -0,0 +1,365 @@
+from __future__ import annotations
+
+from contextlib import contextmanager
+from contextvars import ContextVar
+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.place import Block, BlockPlace, TempBlock
+from sonolus.script.globals import GlobalInfo, GlobalPlaceholder
+from sonolus.script.internal.value import Value
+
+_compiler_internal_ = True
+
+context_var = ContextVar("context_var", default=None)
+
+
+class GlobalContextState:
+    archetypes: dict[type, int]
+    rom: ReadOnlyMemory
+    const_mappings: dict[Any, int]
+    environment_mappings: dict[GlobalInfo, int]
+    environment_offsets: dict[Block, int]
+    mode: Mode
+
+    def __init__(self, mode: Mode, archetypes: dict[type, int] | None = None, rom: ReadOnlyMemory | None = None):
+        self.archetypes = archetypes or {}
+        self.rom = ReadOnlyMemory() if rom is None else rom
+        self.const_mappings = {}
+        self.environment_mappings = {}
+        self.environment_offsets = {}
+        self.mode = mode
+
+
+class CallbackContextState:
+    callback: str
+    used_names: dict[str, int]
+
+    def __init__(self, callback: str):
+        self.callback = callback
+        self.used_names = {}
+
+
+class Context:
+    global_state: GlobalContextState
+    callback_state: CallbackContextState
+    statements: list[IRStmt]
+    test: IRStmt
+    outgoing: dict[float | None, Context]
+    scope: Scope
+    loop_variables: dict[str, Value]
+    live: bool
+
+    def __init__(
+        self,
+        global_state: GlobalContextState,
+        callback_state: CallbackContextState,
+        scope: Scope | None = None,
+        live: bool = True,
+    ):
+        self.global_state = global_state
+        self.callback_state = callback_state
+        self.statements = []
+        self.test = IRConst(0)
+        self.outgoing = {}
+        self.scope = scope if scope is not None else Scope()
+        self.loop_variables = {}
+        self.live = live
+
+    @property
+    def rom(self) -> ReadOnlyMemory:
+        return self.global_state.rom
+
+    @property
+    def blocks(self) -> type[Block]:
+        return self.global_state.mode.blocks
+
+    @property
+    def callback(self) -> str:
+        return self.callback_state.callback
+
+    @property
+    def used_names(self) -> dict[str, int]:
+        return self.callback_state.used_names
+
+    @property
+    def const_mappings(self) -> dict[Any, int]:
+        return self.global_state.const_mappings
+
+    def check_readable(self, place: BlockPlace):
+        if not self.callback:
+            return
+        if isinstance(place.block, BlockData) and self.callback not in self.blocks(place.block).readable:
+            raise RuntimeError(f"Block {place.block} is not readable in {self.callback}")
+
+    def check_writable(self, place: BlockPlace):
+        if not self.callback:
+            return
+        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_statements(self, *statements: IRStmt):
+        self.statements.extend(statements)
+
+    def alloc(self, name: str | None = None, size: int = 1) -> BlockPlace:
+        if size == 0:
+            return BlockPlace(TempBlock(name or "e", 0), 0)
+        name = name or ("v" if size == 1 else "a")
+        num = self._get_alloc_number(name)
+        return BlockPlace(TempBlock(f"{name}{num}", size), 0)
+
+    def _get_alloc_number(self, name: str) -> int:
+        if name not in self.used_names:
+            self.used_names[name] = 0
+        self.used_names[name] += 1
+        return self.used_names[name]
+
+    def copy_with_scope(self, scope: Scope) -> Context:
+        return Context(
+            global_state=self.global_state,
+            callback_state=self.callback_state,
+            scope=scope,
+            live=self.live,
+        )
+
+    def branch(self, condition: float | None):
+        assert condition not in self.outgoing
+        result = self.copy_with_scope(self.scope.copy())
+        self.outgoing[condition] = result
+        return result
+
+    def branch_with_scope(self, condition: float | None, scope: Scope):
+        assert condition not in self.outgoing
+        result = self.copy_with_scope(scope)
+        self.outgoing[condition] = result
+        return result
+
+    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:
+        # 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:
+            binding = self.scope.get_binding(name)
+            if not isinstance(binding, ValueBinding):
+                continue
+            value = binding.value
+            type_ = type(value)
+            if type_._is_value_type_():
+                target_value = type_._from_place_(header.alloc(size=type_._size_()))
+                with using_ctx(self):
+                    target_value._set_(value)
+                header.scope.set_value(name, target_value)
+                header.loop_variables[name] = target_value
+            else:
+                header.scope.set_binding(name, ConflictBinding())
+        return header
+
+    def branch_to_loop_header(self, header: Self):
+        if not self.live:
+            return
+        assert len(self.outgoing) == 0
+        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)
+
+    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:
+        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")
+        if value not in self.global_state.environment_mappings:
+            if value.offset is None:
+                offset = self.global_state.environment_offsets.get(block, 0)
+                self.global_state.environment_mappings[value] = offset
+                self.global_state.environment_offsets[block] = offset + value.size
+            else:
+                self.global_state.environment_mappings[value] = value.offset
+        return BlockPlace(block, self.global_state.environment_mappings[value])
+
+    @classmethod
+    def meet(cls, contexts: list[Context]) -> Context:
+        if not contexts:
+            raise RuntimeError("Cannot meet empty list of contexts")
+        if not any(context.live for context in contexts):
+            return contexts[0].into_dead()
+        contexts = [context for context in contexts if context.live]
+        assert not any(context.outgoing for context in contexts)
+        assert all(len(context.outgoing) == 0 for context in contexts)
+        target = contexts[0].copy_with_scope(Scope())
+        Scope.apply_merge(target, contexts)
+        for context in contexts:
+            context.outgoing[None] = target
+        return target
+
+
+def ctx() -> Context | None:
+    return context_var.get()
+
+
+def set_ctx(value: Context | None):
+    return context_var.set(value)
+
+
+@contextmanager
+def using_ctx(value: Context | None):
+    token = context_var.set(value)
+    try:
+        yield
+    finally:
+        context_var.reset(token)
+
+
+class ReadOnlyMemory:
+    values: list[float]
+    indexes: dict[tuple[float, ...], int]
+
+    def __init__(self):
+        self.values = []
+        self.indexes = {}
+
+    def __getitem__(self, item: tuple[float, ...]) -> BlockPlace:
+        if item not in self.indexes:
+            index = len(self.values)
+            self.indexes[item] = index
+            self.values.extend(item)
+        else:
+            index = self.indexes[item]
+        return BlockPlace(self.block, index)
+
+    @property
+    def block(self) -> Block:
+        if ctx():
+            return ctx().blocks.EngineRom
+        else:
+            return PlayBlock.EngineRom
+
+
+@dataclass
+class ValueBinding:
+    value: Value
+
+
+@dataclass
+class ConflictBinding:
+    pass
+
+
+@dataclass
+class EmptyBinding:
+    pass
+
+
+Binding = ValueBinding | ConflictBinding | EmptyBinding
+
+
+class Scope:
+    bindings: dict[str, Binding]
+
+    def __init__(self, bindings: dict[str, Binding] | None = None):
+        self.bindings = bindings or {}
+
+    def get_binding(self, name: str) -> Binding:
+        return self.bindings.get(name, EmptyBinding())
+
+    def set_binding(self, name: str, binding: Binding):
+        self.bindings[name] = binding
+
+    def get_value(self, name: str) -> Value:
+        binding = self.get_binding(name)
+        match binding:
+            case ValueBinding(value):
+                # we don't need to call _get_() here because _set_() is never called where it could be a problem
+                return value
+            case ConflictBinding():
+                raise RuntimeError(
+                    f"Binding '{name}' has multiple conflicting definitions or may not be guaranteed to be defined"
+                )
+            case EmptyBinding():
+                raise RuntimeError(f"Binding '{name}' is not defined")
+
+    def set_value(self, name: str, value: Value):
+        from sonolus.script.internal.impl import validate_value
+
+        self.bindings[name] = ValueBinding(validate_value(value))
+
+    def delete_binding(self, name: str):
+        del self.bindings[name]
+
+    def copy(self) -> Scope:
+        return Scope(self.bindings.copy())
+
+    @classmethod
+    def apply_merge(cls, target: Context, incoming: list[Context]):
+        if not incoming:
+            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}
+        for key in keys:
+            bindings = [source.get_binding(key) for source in sources]
+            if not all(isinstance(binding, ValueBinding) for binding in bindings):
+                target.scope.set_binding(key, ConflictBinding())
+                continue
+            values = [binding.value for binding in bindings]
+            if len({id(value) for value in values}) == 1:
+                target.scope.set_binding(key, ValueBinding(values[0]))
+                continue
+            types = {type(value) for value in values}
+            if len(types) > 1:
+                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_()))
+                for inc in incoming:
+                    with using_ctx(inc):
+                        target_value._set_(inc.scope.get_value(key))
+                target.scope.set_value(key, target_value)
+                continue
+            else:
+                target.scope.set_binding(key, ConflictBinding())
+                continue
+
+
+def iter_contexts(context: Context):
+    seen = set()
+    queue = [context]
+    while queue:
+        current = queue.pop()
+        if current in seen:
+            continue
+        seen.add(current)
+        yield current
+        queue.extend(current.outgoing.values())
+
+
+def context_to_cfg(context: Context) -> BasicBlock:
+    blocks = {context: BasicBlock(statements=context.statements, test=context.test)}
+    for current in iter_contexts(context):
+        for condition, target in current.outgoing.items():
+            if target not in blocks:
+                blocks[target] = BasicBlock(statements=target.statements, test=target.test)
+            edge = FlowEdge(src=blocks[current], dst=blocks[target], cond=condition)
+            blocks[current].outgoing.add(edge)
+            blocks[target].incoming.add(edge)
+    return blocks[context]

sonolus/script/internal/descriptor.py

@@ -0,0 +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

sonolus/script/internal/error.py

@@ -0,0 +1,15 @@
+class InternalError(RuntimeError):
+    """Represents an error occurring due to a violation of an internal invariant.
+
+    This indicates there is a bug in sonolus.py, or that internal details have been used incorrectly.
+    """
+
+    def __init__(self, message: str):
+        super().__init__(message)
+
+
+class CompilationError(RuntimeError):
+    """Represents an error occurring during the compilation of a script."""
+
+    def __init__(self, message: str):
+        super().__init__(message)

sonolus/script/internal/generic.py

@@ -0,0 +1,197 @@
+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.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 get_origin(result) is Annotated:
+        return result.__args__[0]
+    if get_origin(result) is Literal:
+        return result.__args__[0]
+    return result
+
+
+def validate_type_spec(spec: Any) -> PartialGeneric | TypeVar | type[Value]:
+    spec = validate_type_arg(spec)
+    if isinstance(spec, type) and issubclass(spec, Enum):
+        # For values like IntEnum subclasses, this will call validate_type_spec(IntEnum),
+        # which in turn will call it on int, so this works.
+        spec = validate_type_spec(spec.__mro__[1])
+    if isinstance(spec, PartialGeneric | TypeVar) or (isinstance(spec, type) and issubclass(spec, Value)):
+        return spec
+    raise TypeError(f"Invalid type spec: {spec}")
+
+
+def validate_concrete_type(spec: Any) -> type[Value]:
+    spec = validate_type_spec(spec)
+    if isinstance(spec, type) and issubclass(spec, Value) and spec._is_concrete_():
+        return spec
+    raise TypeError(f"Expected a concrete type, got {spec}")
+
+
+def validate_type_args(args) -> tuple[Any, ...]:
+    if not isinstance(args, tuple):
+        args = (args,)
+    return tuple(validate_type_arg(arg) for arg in args)
+
+
+def contains_incomplete_type(args) -> bool:
+    if not isinstance(args, tuple):
+        args = (args,)
+    return any(isinstance(arg, TypeVar | PartialGeneric) for arg in args)
+
+
+def format_type_arg(arg: Any) -> str:
+    if isinstance(arg, type):
+        return arg.__name__
+    return f"{arg}"
+
+
+class GenericValue(Value):
+    _parameterized_: ClassVar[dict[tuple[Any, ...], type[Self]]] = {}
+    _type_args_: ClassVar[tuple[Any, ...] | None] = None
+    _type_vars_to_args_: ClassVar[dict[TypeVar, Any] | None] = None
+
+    def __init__(self):
+        if self._type_args_ is None:
+            raise TypeError(f"Missing type arguments for {self.__class__.__name__}")
+
+    @classmethod
+    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.
+        """
+        if len(args) != len(cls.__type_params__):
+            raise TypeError(f"Expected {len(cls.__type_params__)} type arguments, got {len(args)}")
+        return args
+
+    @classmethod
+    def _is_concrete_(cls) -> bool:
+        return cls._type_args_ is not None
+
+    @classmethod
+    @meta_fn
+    def _get_type_arg_(cls, var: TypeVar) -> Any:
+        if isinstance(var, Value):
+            var = var._as_py_()
+        if cls._type_args_ is None:
+            raise TypeError(f"Type {cls.__name__} is not parameterized")
+        if var in cls._type_vars_to_args_:
+            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)
+        if contains_incomplete_type(args):
+            return PartialGeneric(cls, args)
+        if args not in cls._parameterized_:
+            cls._parameterized_[args] = cls._get_parameterized(args)
+        return cls._parameterized_[args]
+
+    @classmethod
+    def _get_parameterized(cls, args: tuple[Any, ...]) -> type[Self]:
+        class Parameterized(cls):
+            _type_args_ = args
+            _type_vars_to_args_ = dict(zip(cls.__type_params__, args, strict=True))  # noqa: RUF012
+
+        if args:
+            Parameterized.__name__ = f"{cls.__name__}[{', '.join(format_type_arg(arg) for arg in args)}]"
+            Parameterized.__qualname__ = f"{cls.__qualname__}[{', '.join(format_type_arg(arg) for arg in args)}]"
+        else:
+            Parameterized.__name__ = cls.__name__
+            Parameterized.__qualname__ = cls.__qualname__
+        Parameterized.__module__ = cls.__module__
+        return Parameterized
+
+
+class PartialGeneric[T: GenericValue]:
+    def __init__(self, base: type[T], args: tuple[Any, ...]):
+        self.base = base
+        self.args = args
+
+    def __repr__(self):
+        return f"PartialGeneric({self.base!r}, {self.args!r})"
+
+    def __str__(self):
+        params = ", ".join(format_type_arg(arg) for arg in self.args)
+        return f"{self.base.__name__}?[{params}]"
+
+    def __call__(self, *args, **kwargs):
+        instance = self.base(*args, **kwargs)
+        # Throw an error if it fails
+        accept_and_infer_types(self, instance, {})
+        return instance
+
+
+def infer_and_validate_types(dst: Any, src: Any, results: dict[TypeVar, Any] | None = None) -> dict[TypeVar, Any]:
+    results = results if results is not None else {}
+    match dst:
+        case TypeVar():
+            if dst in results:
+                if results[dst] != src:
+                    raise TypeError(
+                        f"Conflicting types for {dst}: {format_type_arg(results[dst])} and {format_type_arg(src)}"
+                    )
+            else:
+                results[dst] = validate_type_arg(src)
+        case PartialGeneric():
+            if not isinstance(src, type) or not issubclass(src, dst.base):
+                raise TypeError(f"Expected type {dst.base.__name__}, got {format_type_arg(src)}")
+            assert issubclass(src, GenericValue)
+            for d, s in zip(dst.args, src._type_args_, strict=True):
+                infer_and_validate_types(d, s, results)
+        case _:
+            if (
+                src != dst  # noqa: PLR1714
+                and dst != Any
+                and not (isinstance(dst, type) and isinstance(src, type) and issubclass(src, dst))
+            ):
+                raise TypeError(f"Expected {format_type_arg(dst)}, got {format_type_arg(src)}")
+    return results
+
+
+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():
+            infer_and_validate_types(dst, type(val), results)
+            return val
+        case PartialGeneric():
+            val = dst.base._accept_(val)
+            infer_and_validate_types(dst, type(val), results)
+            return val
+        case type():
+            return dst._accept_(val)
+        case _:
+            raise TypeError(f"Expected a type, got {format_type_arg(dst)}")
+
+
+def validate_and_resolve_type(dst: Any, results: dict[TypeVar, Any]) -> Any:
+    match dst:
+        case TypeVar():
+            if dst in results:
+                return results[dst]
+            raise TypeError(f"Missing type for {dst}")
+        case PartialGeneric():
+            return dst.base[tuple(validate_and_resolve_type(arg, results) for arg in dst.args)]
+        case _:
+            return dst

sonolus/script/internal/impl.py

@@ -0,0 +1,69 @@
+from __future__ import annotations
+
+from collections.abc import Callable
+from enum import Enum
+from types import FunctionType, MethodType, NoneType, NotImplementedType
+from typing import TYPE_CHECKING, Annotated, Any, Literal, TypeVar, get_origin, overload
+
+if TYPE_CHECKING:
+    from sonolus.script.internal.value import Value
+
+
+@overload
+def meta_fn[T: Callable](fn: T) -> T: ...
+
+
+@overload
+def meta_fn[T: Callable]() -> Callable[[T], T]: ...
+
+
+def meta_fn(fn=None):
+    # noinspection PyShadowingNames
+    def decorator(fn):
+        fn._meta_fn_ = True
+        return fn
+
+    if fn is None:
+        return decorator
+    return decorator(fn)
+
+
+def validate_value(value: Any) -> Value:
+    result = try_validate_value(value)
+    if result is None:
+        raise TypeError(f"Unsupported value: {value!r}")
+    return result
+
+
+def try_validate_value(value: Any) -> Value | None:
+    from sonolus.script.comptime import Comptime
+    from sonolus.script.globals import GlobalPlaceholder
+    from sonolus.script.internal.generic import PartialGeneric
+    from sonolus.script.internal.value import Value
+    from sonolus.script.num import Num
+
+    match value:
+        case Enum():
+            return validate_value(value.value)
+        case Value():
+            return value
+        case type():
+            if value in {int, float, bool}:
+                return Comptime.accept_unchecked(Num)
+            return Comptime.accept_unchecked(value)
+        case int() | float() | bool():
+            return Num._accept_(value)
+        case tuple():
+            return Comptime.accept_unchecked(tuple(validate_value(v) for v in 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 value.get()
+        case comptime_value if getattr(comptime_value, "_is_comptime_value_", False):
+            return Comptime.accept_unchecked(comptime_value)
+        case _:
+            return None

sonolus/script/internal/introspection.py

@@ -0,0 +1,14 @@
+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