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

sonolus/script/array.py

@@ -5,29 +5,43 @@ from collections.abc import Iterable
 from typing import Any, Self, final
 
 from sonolus.backend.place import BlockPlace
+from sonolus.script.array_like import ArrayLike
 from sonolus.script.debug import assert_unreachable
 from sonolus.script.internal.context import ctx
 from sonolus.script.internal.error import InternalError
 from sonolus.script.internal.generic import GenericValue
 from sonolus.script.internal.impl import meta_fn, validate_value
 from sonolus.script.internal.value import Value
-from sonolus.script.iterator import ArrayLike
 from sonolus.script.num import Num
 
 
 @final
 class Array[T, Size](GenericValue, ArrayLike[T]):
+    """A fixed size array of values.
+
+    Usage:
+        ```python
+        array_1 = Array(1, 2, 3)
+        array_2 = Array[int, 0]()
+        ```
+    """
+
     _value: list[T] | BlockPlace
 
     @classmethod
     @meta_fn
     def element_type(cls) -> type[T] | type[Value]:
-        return cls._get_type_arg_(T)
+        """Return the type of elements in this array type."""
+        return cls.type_var_value(T)
 
     @classmethod
     @meta_fn
     def size(cls) -> int:
-        return cls._get_type_arg_(Size)
+        """Return the size of this array type.
+
+        On instances, use `len(array)` instead.
+        """
+        return cls.type_var_value(Size)
 
     def __new__(cls, *args: T) -> Array[T, Any]:
         if cls._type_args_ is None:
@@ -97,10 +111,10 @@ class Array[T, Size](GenericValue, ArrayLike[T]):
         iterator = iter(values)
         return cls(*(cls.element_type()._from_list_(iterator) for _ in range(cls.size())))
 
-    def _to_list_(self) -> list[float | BlockPlace]:
+    def _to_list_(self, level_refs: dict[Any, int] | None = None) -> list[float | BlockPlace]:
         match self._value:
             case list():
-                return [entry for value in self._value for entry in value._to_list_()]
+                return [entry for value in self._value for entry in value._to_list_(level_refs)]
             case BlockPlace():
                 return [
                     entry
@@ -130,7 +144,7 @@ class Array[T, Size](GenericValue, ArrayLike[T]):
 
     def _copy_(self) -> Self:
         if ctx():
-            place = ctx().alloc(size=self.size())
+            place = ctx().alloc(size=self._size_())
             result: Self = self._from_place_(place)
             result._copy_from_(self)
             return result
@@ -145,16 +159,17 @@ class Array[T, Size](GenericValue, ArrayLike[T]):
         else:
             return cls._with_value([cls.element_type()._alloc_() for _ in range(cls.size())])
 
+    def __len__(self):
+        return self.size()
+
     @meta_fn
     def __getitem__(self, index: Num) -> T:
         index: Num = Num._accept_(index)
-        if index._is_py_():
+        if index._is_py_() and 0 <= index._as_py_() < self.size():
             const_index = index._as_py_()
             if isinstance(const_index, float) and not const_index.is_integer():
                 raise ValueError("Array index must be an integer")
             const_index = int(const_index)
-            if not 0 <= const_index < self.size():
-                raise IndexError("Array index out of range")
             if isinstance(self._value, list):
                 if ctx():
                     return self._value[const_index]._get_()
@@ -215,7 +230,9 @@ class Array[T, Size](GenericValue, ArrayLike[T]):
                 dst._copy_from_(value)
 
     def __eq__(self, other):
-        if self.size() != other.size():
+        if not isinstance(other, ArrayLike):
+            return False
+        if len(self) != len(other):
             return False
         i = 0
         while i < self.size():

sonolus/script/array_like.py

@@ -0,0 +1,297 @@
+from __future__ import annotations
+
+import random
+from abc import ABC, abstractmethod
+from collections.abc import Callable, Sequence
+from typing import Any
+
+from sonolus.script.iterator import SonolusIterator
+from sonolus.script.num import Num
+from sonolus.script.record import Record
+from sonolus.script.values import copy
+
+# Note: we don't use Range in this file because Range itself inherits from ArrayLike
+
+
+class ArrayLike[T](Sequence, ABC):
+    """Mixin for array-like objects.
+
+    Inheritors must implement `__len__`, `__getitem__`, and `__setitem__`.
+
+    Usage:
+        ```python
+        class MyArrayLike[T](Record, ArrayLike[T]):
+            def __len__(self) -> int:
+                ...
+
+            def __getitem__(self, index: Num) -> T:
+                ...
+
+            def __setitem__(self, index: Num, value: T):
+                ...
+        ```
+    """
+
+    @abstractmethod
+    def __len__(self) -> int:
+        """Return the length of the array."""
+
+    @abstractmethod
+    def __getitem__(self, index: Num) -> T:
+        """Return the item at the given index.
+
+        Args:
+            index: The index of the item. Must be an integer between 0 and `len(self) - 1`.
+        """
+
+    @abstractmethod
+    def __setitem__(self, index: Num, value: T):
+        """Set the value of the item at the given index.
+
+        Args:
+            index: The index of the item. Must be an integer between 0 and `len(self) - 1`.
+            value: The value to set.
+        """
+
+    def __iter__(self) -> SonolusIterator[T]:
+        """Return an iterator over the array."""
+        return _ArrayIterator(0, self)
+
+    def __contains__(self, value: Any) -> bool:
+        """Return whether any element in the array is equal to the given value.
+
+        Args:
+            value: The value to check for.
+        """
+        i = 0
+        while i < len(self):
+            if self[i] == value:
+                return True
+            i += 1
+        return False
+
+    def __reversed__(self):
+        """Return a reversed view of the array."""
+        return _ArrayReverser(self)
+
+    def _enumerate_(self, start: Num = 0) -> SonolusIterator[T]:
+        return _ArrayEnumerator(0, start, self)
+
+    def index(self, value: T, start: Num = 0, stop: Num | None = None) -> Num:
+        """Return the index of the value in the array equal to the given value.
+
+        Args:
+            value: The value to search for.
+            start: The index to start searching from.
+            stop: The index to stop searching at. If `None`, search to the end of the array.
+        """
+        if stop is None:
+            stop = len(self)
+        i = start
+        while i < stop:
+            if self[i] == value:
+                return i
+            i += 1
+        return -1
+
+    def count(self, value: T) -> Num:
+        """Return the number of elements in the array equal to the given value.
+
+        Args:
+            value: The value to count.
+        """
+        count = 0
+        i = 0
+        while i < len(self):
+            if self[i] == value:
+                count += 1
+            i += 1
+        return count
+
+    def last_index(self, value: T) -> Num:
+        """Return the last index of the value in the array equal to the given value.
+
+        Args:
+            value: The value to search for.
+        """
+        i = len(self) - 1
+        while i >= 0:
+            if self[i] == value:
+                return i
+            i -= 1
+        return -1
+
+    def index_of_max(self, *, key: Callable[T, Any] | None = None) -> Num:
+        """Return the index of the maximum value in the array.
+
+        Args:
+            key: A one-argument ordering function to use for comparison like the one used in `max()`.
+        """
+        if len(self) == 0:
+            return -1
+        if key is None:
+            key = _identity
+        max_index = 0
+        i = 1
+        while i < len(self):
+            if key(self[i]) > key(self[max_index]):
+                max_index = i
+            i += 1
+        return max_index
+
+    def index_of_min(self, *, key: Callable[T, Any] | None = None) -> Num:
+        """Return the index of the minimum value in the array.
+
+        Args:
+            key: A one-argument ordering function to use for comparison like the one used in `min()`.
+        """
+        if len(self) == 0:
+            return -1
+        if key is None:
+            key = _identity
+        min_index = 0
+        i = 1
+        while i < len(self):
+            if key(self[i]) < key(self[min_index]):
+                min_index = i
+            i += 1
+        return min_index
+
+    def _max_(self, key: Callable[T, Any] | None = None) -> T:
+        return self[self.index_of_max(key=key)]
+
+    def _min_(self, key: Callable[T, Any] | None = None) -> T:
+        return self[self.index_of_min(key=key)]
+
+    def swap(self, i: Num, j: Num, /):
+        """Swap the values at the given indices.
+
+        Args:
+            i: The first index.
+            j: The second index.
+        """
+        temp = copy(self[i])
+        self[i] = self[j]
+        self[j] = temp
+
+    def sort(self, *, key: Callable[T, Any] | None = None, reverse: bool = False):
+        """Sort the values in the array in place.
+
+        Args:
+            key: A one-argument ordering function to use for comparison.
+            reverse: If `True`, sort in descending order, otherwise sort in ascending order.
+        """
+        if len(self) < 15 or key is not None:
+            if key is None:
+                key = _identity
+            _insertion_sort(self, 0, len(self), key, reverse)
+        else:
+            # Heap sort is unstable, so if there's a key, we can't rely on it
+            _heap_sort(self, 0, len(self), reverse)
+
+    def shuffle(self):
+        """Shuffle the values in the array in place."""
+        random.shuffle(self)  # type: ignore
+
+    def reverse(self):
+        """Reverse the values in the array in place."""
+        i = 0
+        j = len(self) - 1
+        while i < j:
+            self.swap(i, j)
+            i += 1
+            j -= 1
+
+
+def _identity[T](value: T) -> T:
+    return value
+
+
+def _insertion_sort[T](array: ArrayLike[T], start: Num, end: Num, key: Callable[T, Any], reverse: bool):
+    i = start + 1
+    if reverse:
+        while i < end:
+            j = i
+            while j > start and key(array[j - 1]) < key(array[j]):
+                array.swap(j - 1, j)
+                j -= 1
+            i += 1
+    else:
+        while i < end:
+            j = i
+            while j > start and key(array[j - 1]) > key(array[j]):
+                array.swap(j - 1, j)
+                j -= 1
+            i += 1
+
+
+def _heapify[T](array: ArrayLike[T], end: Num, index: Num, reverse: bool):
+    while True:
+        left = index * 2 + 1
+        right = left + 1
+        largest = index
+        if left < end and (array[left] > array[largest]) != reverse:
+            largest = left
+        if right < end and (array[right] > array[largest]) != reverse:
+            largest = right
+        if largest == index:
+            break
+        array.swap(index, largest)
+        index = largest
+
+
+def _heap_sort[T](array: ArrayLike[T], start: Num, end: Num, reverse: bool):
+    i = end // 2 - 1
+    while i >= start:
+        _heapify(array, end, i, reverse)
+        i -= 1
+    i = end - 1
+    while i > start:
+        array.swap(start, i)
+        _heapify(array, i, start, reverse)
+        i -= 1
+
+
+class _ArrayIterator[V: ArrayLike](Record, SonolusIterator):
+    i: int
+    array: V
+
+    def has_next(self) -> bool:
+        return self.i < len(self.array)
+
+    def get(self) -> V:
+        return self.array[self.i]
+
+    def advance(self):
+        self.i += 1
+
+
+class _ArrayReverser[V: ArrayLike](Record, ArrayLike):
+    array: V
+
+    def __len__(self) -> int:
+        return len(self.array)
+
+    def __getitem__(self, index: Num) -> V:
+        return self.array[len(self) - 1 - index]
+
+    def __setitem__(self, index: Num, value: V):
+        self.array[len(self) - 1 - index] = value
+
+    def reversed(self) -> ArrayLike[V]:
+        return self.array
+
+
+class _ArrayEnumerator[V: ArrayLike](Record, SonolusIterator):
+    i: int
+    offset: int
+    array: V
+
+    def has_next(self) -> bool:
+        return self.i < len(self.array)
+
+    def get(self) -> tuple[int, Any]:
+        return self.i + self.offset, self.array[self.i]
+
+    def advance(self):
+        self.i += 1