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

sonolus/script/interval.py

@@ -3,118 +3,336 @@ from typing import Self
 from sonolus.backend.ops import Op
 from sonolus.script.debug import error
 from sonolus.script.internal.native import native_function
+from sonolus.script.num import Num
 from sonolus.script.record import Record
 
 
 class Interval(Record):
-    """A closed interval."""
+    """A closed interval.
+
+    Usage:
+        ```python
+        Interval(start: float, end: float)
+        ```
+    """
 
     start: float
     end: float
 
     @property
     def length(self) -> float:
+        """The length of the interval.
+
+        May be negative if the end is less than the start.
+        """
         return self.end - self.start
 
     @property
     def is_empty(self) -> bool:
+        """Whether the interval has length of zero or less."""
         return self.start > self.end
 
     @property
     def mid(self) -> float:
+        """The midpoint of the interval."""
         return (self.start + self.end) / 2
 
     @property
     def tuple(self):
+        """The interval as a tuple."""
         return self.start, self.end
 
     def __contains__(self, item: Self | float | int) -> bool:
+        """Check if an item is within the interval.
+
+        Args:
+            item: The item to check. If it is an interval, it must be fully contained within this interval.
+
+        Returns:
+            True if the item is within the interval, False otherwise.
+        """
         match item:
             case Interval(start, end):
                 return self.start <= start and end <= self.end
-            case float() | int() as value:
+            case Num(value):
                 return self.start <= value <= self.end
             case _:
                 error("Invalid type for interval check")
 
     def __add__(self, other: float | int) -> Self:
+        """Add a value to both ends of the interval.
+
+        Args:
+            other: The value to add.
+
+        Returns:
+            A new interval with the value added to both ends.
+        """
         return Interval(self.start + other, self.end + other)
 
     def __sub__(self, other: float | int) -> Self:
+        """Subtract a value from both ends of the interval.
+
+        Args:
+            other: The value to subtract.
+
+        Returns:
+            A new interval with the value subtracted from both ends.
+        """
         return Interval(self.start - other, self.end - other)
 
     def __mul__(self, other: float | int) -> Self:
+        """Multiply both ends of the interval by a value.
+
+        Args:
+            other: The value to multiply by.
+
+        Returns:
+            A new interval with both ends multiplied by the value.
+        """
         return Interval(self.start * other, self.end * other)
 
     def __truediv__(self, other: float | int) -> Self:
+        """Divide both ends of the interval by a value.
+
+        Args:
+            other: The value to divide by.
+
+        Returns:
+            A new interval with both ends divided by the value.
+        """
         return Interval(self.start / other, self.end / other)
 
     def __floordiv__(self, other: float | int) -> Self:
+        """Divide both ends of the interval by a value and floor the result.
+
+        Args:
+            other: The value to divide by.
+
+        Returns:
+            A new interval with both ends divided by the value and floored.
+        """
         return Interval(self.start // other, self.end // other)
 
     def __and__(self, other: Self) -> Self:
+        """Get the intersection of two intervals.
+
+        The resulting interval will be empty and may have a negative length if the two intervals do not overlap.
+
+        Args:
+            other: The other interval.
+
+        Returns:
+            A new interval representing the intersection of the two intervals.
+        """
         return Interval(max(self.start, other.start), min(self.end, other.end))
 
     def shrink(self, value: float | int) -> Self:
+        """Shrink the interval by a value on both ends.
+
+        Args:
+            value: The value to shrink by.
+
+        Returns:
+            A new interval with the value subtracted from the start and added to the end.
+        """
         return Interval(self.start + value, self.end - value)
 
     def expand(self, value: float | int) -> Self:
+        """Expand the interval by a value on both ends.
+
+        Args:
+            value: The value to expand by.
+
+        Returns:
+            A new interval with the value subtracted from the start and added to the end.
+        """
         return Interval(self.start - value, self.end + value)
 
     def lerp(self, x: float, /) -> float:
+        """Linearly interpolate a value within the interval.
+
+        Args:
+            x: The interpolation factor.
+
+        Returns:
+            The interpolated value.
+        """
         return lerp(self.start, self.end, x)
 
     def lerp_clamped(self, x: float, /) -> float:
+        """Linearly interpolate a value within the interval, clamped to the interval.
+
+        Args:
+            x: The interpolation factor.
+
+        Returns:
+            The interpolated value.
+        """
         return lerp_clamped(self.start, self.end, x)
 
     def unlerp(self, x: float, /) -> float:
+        """Inverse linear interpolation of a value within the interval.
+
+        Args:
+            x: The value to unlerp.
+
+        Returns:
+            The unlerped value.
+        """
         return unlerp(self.start, self.end, x)
 
     def unlerp_clamped(self, x: float, /) -> float:
+        """Inverse linear interpolation of a value within the interval, clamped to the interval.
+
+        Args:
+            x: The value to unlerp.
+
+        Returns:
+            The unlerped value.
+        """
         return unlerp_clamped(self.start, self.end, x)
 
     def clamp(self, x: float, /) -> float:
+        """Clamp a value to the interval.
+
+        Args:
+            x: The value to clamp.
+
+        Returns:
+            The clamped value.
+        """
         return clamp(x, self.start, self.end)
 
 
 @native_function(Op.Lerp)
-def lerp(a, b, x, /):
+def _num_lerp(a, b, x, /):
     return a + (b - a) * x
 
 
 @native_function(Op.LerpClamped)
-def lerp_clamped(a, b, x, /):
+def _num_lerp_clamped(a, b, x, /):
+    return a + (b - a) * max(0, min(1, x))
+
+
+def _generic_lerp[T](a: T, b: T, x: float, /) -> T:
+    return a + (b - a) * x
+
+
+def _generic_lerp_clamped[T](a: T, b: T, x: float, /) -> T:
     return a + (b - a) * max(0, min(1, x))
 
 
+def lerp[T](a: T, b: T, x: float, /) -> T:
+    """Linearly interpolate between two values.
+
+    Args:
+        a: The start value.
+        b: The end value.
+        x: The interpolation factor.
+
+    Returns:
+        The interpolated value.
+    """
+    match a, b:
+        case (Num(a), Num(b)):
+            return _num_lerp(a, b, x)
+        case _:
+            return _generic_lerp(a, b, x)
+
+
+def lerp_clamped[T](a: T, b: T, x: float, /) -> T:
+    """Linearly interpolate between two values, clamped to the interval.
+
+    Args:
+        a: The start value.
+        b: The end value.
+        x: The interpolation factor.
+
+    Returns:
+        The interpolated value.
+    """
+    match a, b:
+        case (Num(a), Num(b)):
+            return _num_lerp_clamped(a, b, x)
+        case _:
+            return _generic_lerp_clamped(a, b, x)
+
+
 @native_function(Op.Unlerp)
-def unlerp(a, b, x, /):
+def unlerp(a: float, b: float, x: float, /) -> float:
+    """Inverse linear interpolation.
+
+    Args:
+        a: The start value.
+        b: The end value.
+        x: The value to unlerp.
+
+    Returns:
+        The unlerped value.
+    """
     return (x - a) / (b - a)
 
 
 @native_function(Op.UnlerpClamped)
-def unlerp_clamped(a, b, x, /):
+def unlerp_clamped(a: float, b: float, x: float, /) -> float:
+    """Inverse linear interpolation, clamped to the interval.
+
+    Args:
+        a: The start value.
+        b: The end value.
+        x: The value to unlerp.
+
+    Returns:
+        The unlerped value.
+    """
     return max(0, min(1, (x - a) / (b - a)))
 
 
 @native_function(Op.Remap)
-def remap(a, b, c, d, x, /):
+def remap(a: float, b: float, c: float, d: float, x: float, /) -> float:
+    """Remap a value from one interval to another.
+
+    Args:
+        a: The start of the input interval.
+        b: The end of the input interval.
+        c: The start of the output interval.
+        d: The end of the output interval.
+        x: The value to remap.
+
+    Returns:
+        The remapped value.
+    """
     return c + (d - c) * (x - a) / (b - a)
 
 
 @native_function(Op.RemapClamped)
-def remap_clamped(a, b, c, d, x, /):
+def remap_clamped(a: float, b: float, c: float, d: float, x: float, /) -> float:
+    """Remap a value from one interval to another, clamped to the output interval.
+
+    Args:
+        a: The start of the input interval.
+        b: The end of the input interval.
+        c: The start of the output interval.
+        d: The end of the output interval.
+        x: The value to remap.
+
+    Returns:
+        The remapped value.
+    """
     return c + (d - c) * max(0, min(1, (x - a) / (b - a)))
 
 
 @native_function(Op.Clamp)
-def clamp(x, a, b, /):
-    return max(a, min(b, x))
-
-
-def generic_lerp[T](a: T, b: T, x: float, /) -> T:
-    return a + (b - a) * x
+def clamp(x: float, a: float, b: float, /) -> float:
+    """Clamp a value to an interval.
 
+    Args:
+        x: The value to clamp.
+        a: The start of the interval.
+        b: The end of the interval.
 
-def generic_lerp_clamped[T](a: T, b: T, x: float, /) -> T:
-    return a + (b - a) * max(0, min(1, x))
+    Returns:
+        The clamped value.
+    """
+    return max(a, min(b, x))

sonolus/script/iterator.py

@@ -1,20 +1,61 @@
 from __future__ import annotations
 
 from abc import abstractmethod
-from collections.abc import Collection, Iterator
+from collections.abc import Iterator
+from typing import Any
 
-from sonolus.script.num import Num
+from sonolus.script.internal.impl import meta_fn
 from sonolus.script.record import Record
-from sonolus.script.values import copy
 
 
 class SonolusIterator[T](Iterator[T]):
+    """Base class for Sonolus iterators.
+
+    This class is used to define custom iterators that can be used in Sonolus.py.
+
+    Inheritors must implement the `has_next`, `get`, and `advance` methods.
+    The `__next__` and `__iter__` methods are implemented by default.
+
+    Usage:
+        ```python
+        class MyIterator(Record, SonolusIterator):
+            def has_next(self) -> bool:
+                ...
+
+            def get(self) -> Any:
+                ...
+
+            def advance(self):
+                ...
+        ```
+    """
+
+    def next(self) -> T:
+        result = self.get()
+        self.advance()
+        return result
+
     @abstractmethod
     def has_next(self) -> bool:
+        """Return whether the iterator has more elements."""
         raise NotImplementedError
 
     @abstractmethod
-    def next(self) -> T:
+    def get(self) -> T:
+        """Return the next element of the iterator.
+
+        May be called multiple times before calling `advance`.
+
+        Must not be called if `has_next` returns `False`.
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    def advance(self):
+        """Advance the iterator to the next element.
+
+        Must not be called if `has_next` returns `False`.
+        """
         raise NotImplementedError
 
     def __next__(self) -> T:
@@ -22,193 +63,105 @@ class SonolusIterator[T](Iterator[T]):
             raise StopIteration
         return self.next()
 
+    def __iter__(self) -> SonolusIterator[T]:
+        return self
 
-class ArrayLike[T](Collection):
-    @abstractmethod
-    def size(self) -> int:
-        pass
 
-    @abstractmethod
-    def __getitem__(self, index: Num) -> T:
-        pass
+class _Enumerator[V: SonolusIterator](Record, SonolusIterator):
+    i: int
+    offset: int
+    iterator: V
 
-    @abstractmethod
-    def __setitem__(self, index: Num, value: T):
-        pass
+    def has_next(self) -> bool:
+        return self.iterator.has_next()
 
-    def __len__(self) -> int:
-        return self.size()
+    def get(self) -> tuple[int, Any]:
+        return self.i + self.offset, self.iterator.get()
 
-    def __iter__(self) -> SonolusIterator[T]:
-        return ArrayIterator(0, self)
+    def advance(self):
+        self.i += 1
+        self.iterator.advance()
 
-    def __contains__(self, value: T) -> bool:
-        i = 0
-        while i < self.size():
-            if self[i] == value:
-                return True
-            i += 1
-        return False
 
-    def reversed(self) -> ArrayLike[T]:
-        return ArrayReverser(self)
-
-    def iter(self) -> SonolusIterator[T]:
-        return self.__iter__()  # noqa: PLC2801
-
-    def enumerate(self, start: Num = 0) -> SonolusIterator[T]:
-        return ArrayEnumerator(0, start, self)
-
-    def index_of(self, value: T, start: Num = 0) -> Num:
-        i = start
-        while i < self.size():
-            if self[i] == value:
-                return i
-            i += 1
-        return -1
-
-    def last_index_of(self, value: T) -> Num:
-        i = self.size() - 1
-        while i >= 0:
-            if self[i] == value:
-                return i
-            i -= 1
-        return -1
-
-    def index_of_max(self) -> Num:
-        if self.size() == 0:
-            return -1
-        max_index = 0
-        i = 1
-        while i < self.size():
-            if self[i] > self[max_index]:
-                max_index = i
-            i += 1
-        return max_index
-
-    def index_of_min(self) -> Num:
-        if self.size() == 0:
-            return -1
-        min_index = 0
-        i = 1
-        while i < self.size():
-            if self[i] < self[min_index]:
-                min_index = i
-            i += 1
-        return min_index
-
-    def max(self) -> T:
-        return self[self.index_of_max()]
-
-    def min(self) -> T:
-        return self[self.index_of_min()]
-
-    def swap(self, i: Num, j: Num):
-        temp = copy(self[i])
-        self[i] = self[j]
-        self[j] = temp
-
-    def sort(self, *, reverse: bool = False):
-        if self.size() < 15:
-            _insertion_sort(self, 0, self.size(), reverse)
-        else:
-            _heap_sort(self, 0, self.size(), reverse)
-
-
-def _insertion_sort[T](array: ArrayLike[T], start: Num, end: Num, reverse: bool):
-    i = start + 1
-    while i < end:
-        value = copy(array[i])
-        j = i - 1
-        while j >= start and (array[j] > value) != reverse:
-            array[j + 1] = array[j]
-            j -= 1
-        array[j + 1] = value
-        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
-
-
-# Heap sort is simple to implement iteratively without dynamic memory allocation
-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
+class _Zipper[T](Record, SonolusIterator):
+    # Can be a, Pair[a, b], Pair[a, Pair[b, c]], etc.
+    iterators: T
 
+    @meta_fn
     def has_next(self) -> bool:
-        return self.i < self.array.size()
+        from sonolus.backend.visitor import compile_and_call
 
-    def next(self) -> V:
-        value = self.array[self.i]
-        self.i += 1
-        return value
+        return compile_and_call(self._has_next, self._get_iterators())
 
+    def _get_iterators(self) -> tuple[SonolusIterator, ...]:
+        from sonolus.script.containers import Pair
 
-class ArrayReverser[V: ArrayLike](Record, ArrayLike):
-    array: V
+        iterators = []
+        v = self.iterators
+        while isinstance(v, Pair):
+            iterators.append(v.first)
+            v = v.second
+        iterators.append(v)
+        return tuple(iterators)
 
-    def size(self) -> int:
-        return self.array.size()
+    def _has_next(self, iterators: tuple[SonolusIterator, ...]) -> bool:
+        for iterator in iterators:  # noqa: SIM110
+            if not iterator.has_next():
+                return False
+        return True
 
-    def __getitem__(self, index: Num) -> V:
-        return self.array[self.size() - 1 - index]
+    @meta_fn
+    def get(self) -> tuple[Any, ...]:
+        from sonolus.backend.visitor import compile_and_call
 
-    def __setitem__(self, index: Num, value: V):
-        self.array[self.size() - 1 - index] = value
+        return tuple(compile_and_call(iterator.get) for iterator in self._get_iterators())
 
-    def reversed(self) -> ArrayLike[V]:
-        return self.array
+    @meta_fn
+    def advance(self):
+        from sonolus.backend.visitor import compile_and_call
 
+        for iterator in self._get_iterators():
+            compile_and_call(iterator.advance)
 
-class Enumerator[V: SonolusIterator](Record, SonolusIterator):
-    i: int
-    offset: int
-    iterator: V
+
+class _EmptyIterator(Record, SonolusIterator):
+    def has_next(self) -> bool:
+        return False
+
+    def get(self) -> Any:
+        return None
+
+    def advance(self):
+        pass
+
+
+class _MappingIterator[T, Fn](Record, SonolusIterator):
+    fn: Fn
+    iterator: T
 
     def has_next(self) -> bool:
         return self.iterator.has_next()
 
-    def next(self):
-        value = self.iterator.next()
-        index = self.i + self.offset
-        self.i += 1
-        return index, value
+    def get(self) -> Any:
+        return self.fn(self.iterator.get())
 
+    def advance(self):
+        self.iterator.advance()
 
-class ArrayEnumerator[V: ArrayLike](Record, SonolusIterator):
-    i: int
-    offset: int
-    array: V
+
+class _FilteringIterator[T, Fn](Record, SonolusIterator):
+    fn: Fn
+    iterator: T
 
     def has_next(self) -> bool:
-        return self.i < self.array.size()
+        while self.iterator.has_next():
+            if self.fn(self.iterator.get()):
+                return True
+            self.iterator.advance()
+        return False
 
-    def next(self):
-        value = self.array[self.i]
-        index = self.i + self.offset
-        self.i += 1
-        return index, value
+    def get(self) -> Any:
+        return self.iterator.get()
+
+    def advance(self):
+        self.iterator.advance()