omlish 0.0.0.dev1__py3-none-any.whl

omlish/collections/treap.py

@@ -0,0 +1,228 @@
+"""
+MIT License
+
+Copyright (c) 2022 Gabriel Ochsenhofer
+
+Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated
+documentation files (the "Software"), to deal in the Software without restriction, including without limitation the
+rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit
+persons to whom the Software is furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all copies or substantial portions of the
+Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE
+WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR
+COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR
+OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
+"""
+import typing as ta
+
+
+T = ta.TypeVar('T')
+Comparer = ta.Callable[[T, T], int]
+
+
+class TreapNode(ta.Generic[T]):
+    __slots__ = ('_value', '_priority', '_left', '_right')
+
+    def __init__(
+            self,
+            *,
+            _value: T,
+            _priority: int,
+            _left: ta.Optional['TreapNode[T]'],
+            _right: ta.Optional['TreapNode[T]'],
+    ) -> None:
+        super().__init__()
+
+        self._value = _value
+        self._priority = _priority
+        self._left = _left
+        self._right = _right
+
+    @property
+    def value(self) -> T:
+        return self._value
+
+    @property
+    def priority(self) -> int:
+        return self._priority
+
+    @property
+    def left(self) -> ta.Optional['TreapNode[T]']:
+        return self._left
+
+    @property
+    def right(self) -> ta.Optional['TreapNode[T]']:
+        return self._right
+
+    def __iter__(self) -> ta.Iterator[T]:
+        if self._left is not None:
+            yield from self._left
+        yield self._value
+        if self._right is not None:
+            yield from self._right
+
+
+def find(n: ta.Optional[TreapNode[T]], v: T, c: Comparer[T]) -> ta.Optional[TreapNode[T]]:
+    while True:
+        if n is None:
+            return None
+        diff = c(n._value, v)  # noqa
+        if diff == 0:
+            return n
+        elif diff < 0:
+            n = n._right  # noqa
+        else:
+            n = n._left  # noqa
+
+
+def union(
+        n: ta.Optional[TreapNode[T]],
+        other: ta.Optional[TreapNode[T]],
+        c: Comparer[T],
+        overwrite: bool,
+) -> ta.Optional[TreapNode[T]]:
+    if n is None:
+        return other
+    if other is None:
+        return n
+    if n._priority < other._priority:
+        other, n, overwrite = n, other, not overwrite
+    left, dupe, right = split(other, n._value, c)
+    value = n._value
+    if overwrite and dupe is not None:
+        value = dupe._value
+    left = union(n._left, left, c, overwrite)
+    right = union(n._right, right, c, overwrite)
+    return TreapNode(_value=value, _priority=n._priority, _left=left, _right=right)
+
+
+def split(
+        n: ta.Optional[TreapNode[T]],
+        v: T,
+        c: Comparer[T],
+) -> tuple[
+    ta.Optional[TreapNode[T]],
+    ta.Optional[TreapNode[T]],
+    ta.Optional[TreapNode[T]],
+]:
+    tmp: TreapNode[T] = TreapNode(_value=None, _priority=0, _left=None, _right=None)  # type: ignore
+    leftp, rightp = [tmp, 'l'], [tmp, 'r']
+
+    def setp(p, o):
+        t, s = p
+        if s == 'l':
+            t._left = o
+        elif s == 'r':
+            t._right = o
+        else:
+            raise ValueError(p)
+
+    cur: ta.Optional[TreapNode[T]] = n
+    while True:
+        if cur is None:
+            setp(leftp, None)
+            setp(rightp, None)
+            return tmp._left, None, tmp._right  # noqa
+
+        d = c(cur._value, v)
+        if d < 0:
+            root = TreapNode(_value=cur._value, _priority=cur._priority, _left=cur._left, _right=None)
+            setp(leftp, root)
+            leftp[0], leftp[1] = root, 'r'
+            cur = cur._right
+        elif d > 0:
+            root = TreapNode(_value=cur._value, _priority=cur._priority, _left=None, _right=cur._right)
+            setp(rightp, root)
+            rightp[0], rightp[1] = root, 'l'
+            cur = cur._left
+        else:
+            root = TreapNode(_value=cur._value, _priority=cur._priority, _left=None, _right=None)
+            setp(leftp, cur._left)
+            setp(rightp, cur._right)
+            return tmp._left, root, tmp._right
+
+
+def intersect(
+        n: ta.Optional[TreapNode[T]],
+        other: ta.Optional[TreapNode[T]],
+        c: Comparer[T],
+) -> ta.Optional[TreapNode[T]]:
+    if n is None or other is None:
+        return None
+
+    if n._priority < other._priority:
+        n, other = other, n
+
+    left, found, right = split(other, n._value, c)
+    left = intersect(n._left, left, c)
+    right = intersect(n._right, right, c)
+
+    if found is None:
+        # TODO: use a destructive join as both left/right are copies
+        return _join(left, right)
+
+    return TreapNode(_value=n._value, _priority=n._priority, _left=left, _right=right)
+
+
+def delete(n: ta.Optional[TreapNode[T]], v: T, c: Comparer[T]) -> ta.Optional[TreapNode[T]]:
+    left, _, right = split(n, v, c)
+    return _join(left, right)
+
+
+def diff(n: ta.Optional[TreapNode[T]], other: ta.Optional[TreapNode[T]], c: Comparer[T]) -> ta.Optional[TreapNode[T]]:
+    if n is None or other is None:
+        return n
+
+    # TODO -- use count
+    if n._priority >= other._priority:
+        left, dupe, right = split(other, n._value, c)
+        left, right = diff(n._left, left, c), diff(n._right, right, c)
+        if dupe is not None:
+            return _join(left, right)
+        return TreapNode(_value=n._value, _priority=n._priority, _left=left, _right=right)
+
+    left, _, right = split(n, other._value, c)
+    left = diff(left, other._left, c)
+    right = diff(right, other._right, c)
+    return _join(left, right)
+
+
+def _join(n: ta.Optional[TreapNode[T]], other: ta.Optional[TreapNode[T]]) -> ta.Optional[TreapNode[T]]:
+    result: ta.Optional[TreapNode[T]] = None
+    resultp: list[ta.Any] = [None, None]
+
+    def setresultp(o):
+        t, s = resultp
+        if t is None:
+            nonlocal result
+            result = o
+        elif s == 'l':
+            t._left = o
+        elif s == 'r':
+            t._right = o
+        else:
+            raise ValueError(resultp)
+
+    cur: ta.Optional[TreapNode[T]] = n
+    while True:
+        if cur is None:
+            setresultp(other)
+            return result
+
+        if other is None:
+            setresultp(cur)
+            return result
+
+        if cur._priority <= other._priority:
+            root = TreapNode(_value=cur._value, _priority=cur._priority, _left=cur._left, _right=None)
+            setresultp(root)
+            resultp[0], resultp[1] = root, 'r'
+            cur = cur._right
+        else:
+            root = TreapNode(_value=other._value, _priority=other._priority, _left=None, _right=other._right)
+            setresultp(root)
+            resultp[0], resultp[1] = root, 'l'
+            other = other._left

omlish/collections/treapmap.py

@@ -0,0 +1,144 @@
+"""
+MIT License
+
+Copyright (c) 2022 Gabriel Ochsenhofer
+
+Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated
+documentation files (the "Software"), to deal in the Software without restriction, including without limitation the
+rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit
+persons to whom the Software is furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all copies or substantial portions of the
+Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE
+WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR
+COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR
+OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
+"""
+import random
+import typing as ta
+
+from . import treap
+from .persistent import PersistentMap
+
+
+K = ta.TypeVar('K')
+V = ta.TypeVar('V')
+
+
+class TreapMap(PersistentMap[K, V]):
+    __slots__ = ('_n', '_c')
+
+    def __init__(
+            self,
+            *,
+            _n: ta.Optional[treap.TreapNode[tuple[K, V]]],
+            _c: treap.Comparer[tuple[K, V]],
+    ) -> None:
+        super().__init__()
+
+        self._n = _n
+        self._c = _c
+
+    def __len__(self) -> int:
+        # TODO: memo
+        # TODO: itertools lol
+        if self._n is None:
+            return 0
+        result = 0
+        for _ in self._n:
+            result += 1
+        return result
+
+    def __contains__(self, item: K) -> bool:
+        try:
+            self[item]  # noqa
+        except KeyError:
+            return False
+        else:
+            return True
+
+    def __getitem__(self, item: K) -> V:
+        n = treap.find(self._n, (item, None), self._c)  # type: ignore
+        if n is None:
+            raise KeyError(item)
+        return n.value
+
+    def __iter__(self) -> ta.Iterator[tuple[K, V]]:
+        i = self.iterate()
+        while i.has_next():
+            yield i.next()
+
+    def iterate(self) -> 'TreapMapIterator[K, V]':
+        i = TreapMapIterator(
+            _st=[],
+            _n=self._n,
+            _b=False,
+        )
+        while (n := i._n) is not None and n.left is not None:  # noqa
+            i._st.append(n)  # noqa
+            i._n = n.left  # noqa
+        return i
+
+    def with_(self, k: K, v: V) -> 'TreapMap[K, V]':
+        node = treap.TreapNode(
+            _value=(k, v),
+            _priority=int(random.random() * 0xFFFFFFFF),
+            _left=None,
+            _right=None,
+        )
+        n = treap.union(self._n, node, self._c, True)
+        return TreapMap(_n=n, _c=self._c)
+
+    def without(self, k: K) -> 'TreapMap[K, V]':
+        n = treap.delete(self._n, (k, None), self._c)  # type: ignore
+        return TreapMap(_n=n, _c=self._c)
+
+    def default(self, k: K, v: V) -> 'TreapMap[K, V]':
+        try:
+            self[k]  # noqa
+        except KeyError:
+            return self.with_(k, v)
+        else:
+            return self
+
+
+def new_treap_map(cmp: ta.Callable[[tuple[K, V], tuple[K, V]], int]) -> PersistentMap[K, V]:
+    return TreapMap(_n=None, _c=cmp)
+
+
+class TreapMapIterator(ta.Generic[K, V]):
+    __slots__ = ('_st', '_n', '_b')
+
+    def __init__(
+            self,
+            *,
+            _st: list[treap.TreapNode[tuple[K, V]]],
+            _n: ta.Optional[treap.TreapNode[tuple[K, V]]],
+            _b: bool,
+    ) -> None:
+        super().__init__()
+
+        self._st = _st
+        self._n = _n
+        self._b = _b
+
+    def has_next(self) -> bool:
+        return self._n is not None
+
+    def next(self) -> tuple[K, V]:
+        n = self._n
+        if n is None:
+            raise StopIteration
+        if n.right is not None:
+            self._n = n.right
+            while self._n.left is not None:
+                self._st.append(self._n)
+                self._n = self._n.left
+        elif len(self._st) > 0:
+            self._n = self._st[-1]
+            self._st.pop()
+        else:
+            self._n = None
+        return n.value

omlish/collections/unmodifiable.py

@@ -0,0 +1,174 @@
+import typing as ta
+
+from .. import lang
+
+
+T = ta.TypeVar('T')
+K = ta.TypeVar('K')
+V = ta.TypeVar('V')
+
+
+class Unmodifiable(lang.Abstract):
+    pass
+
+
+class UnmodifiableSequence(ta.Sequence[T], Unmodifiable, lang.Final):
+
+    def __init__(self, target: ta.Sequence[T]) -> None:
+        super().__init__()
+
+        if target is None:
+            raise TypeError(target)
+        self._target = target
+
+    def __repr__(self) -> str:
+        return '%s(%r)' % (type(self).__name__, self._target)
+
+    def __contains__(self, x: object) -> bool:
+        return x in self._target
+
+    def __eq__(self, o: object) -> bool:
+        return self._target == o
+
+    def __ge__(self, other: object) -> bool:
+        return self._target >= other  # type: ignore
+
+    def __getitem__(self, i: ta.Union[int, slice]) -> T:  # type: ignore
+        return self._target[i]  # type: ignore
+
+    def __gt__(self, other: object) -> bool:
+        return self._target > other  # type: ignore
+
+    def __iter__(self) -> ta.Iterator[T]:
+        return iter(self._target)
+
+    def __le__(self, other: object) -> bool:
+        return self._target <= other  # type: ignore
+
+    def __len__(self) -> int:
+        return len(self._target)
+
+    def __lt__(self, other: object) -> bool:
+        return self._target < other  # type: ignore
+
+    def __ne__(self, o: object) -> bool:
+        return self._target != o
+
+    def __reversed__(self) -> ta.Iterator[T]:
+        return reversed(self._target)
+
+    def count(self, x: ta.Any) -> int:
+        return self._target.count(x)
+
+    def index(self, x: ta.Any, *args, **kwargs) -> int:  # type: ignore
+        return self._target.index(x, *args, **kwargs)
+
+
+class UnmodifiableSet(ta.AbstractSet[T], Unmodifiable, lang.Final):
+
+    def __init__(self, target: ta.AbstractSet[T]) -> None:
+        super().__init__()
+
+        if target is None:
+            raise TypeError(target)
+        self._target = target
+
+    def __repr__(self) -> str:
+        return '%s(%r)' % (type(self).__name__, self._target)
+
+    def __and__(self, s: ta.AbstractSet[ta.Any]) -> ta.AbstractSet[T]:
+        return self._target & s
+
+    def __contains__(self, x: object) -> bool:
+        return x in self._target
+
+    def __eq__(self, o: object) -> bool:
+        return self._target == o
+
+    def __ge__(self, s: ta.AbstractSet[ta.Any]) -> bool:
+        return self._target >= s
+
+    def __gt__(self, s: ta.AbstractSet[ta.Any]) -> bool:
+        return self._target > s
+
+    def __iter__(self) -> ta.Iterator[T]:
+        return iter(self._target)
+
+    def __le__(self, s: ta.AbstractSet[ta.Any]) -> bool:
+        return self._target <= s
+
+    def __len__(self) -> int:
+        return len(self._target)
+
+    def __lt__(self, s: ta.AbstractSet[ta.Any]) -> bool:
+        return self._target > s
+
+    def __ne__(self, o: object) -> bool:
+        return self._target != o
+
+    def __or__(self, s: ta.AbstractSet[T]) -> ta.AbstractSet[T]:  # type: ignore
+        return self._target | s
+
+    def __sub__(self, s: ta.AbstractSet[ta.Any]) -> ta.AbstractSet[T]:
+        return self._target - s
+
+    def __xor__(self, s: ta.AbstractSet[T]) -> ta.AbstractSet[T]:  # type: ignore
+        return self._target ^ s
+
+    def isdisjoint(self, s: ta.Iterable[ta.Any]) -> bool:
+        return self._target.isdisjoint(s)
+
+
+class UnmodifiableMapping(ta.Mapping[K, V], Unmodifiable, lang.Final):
+
+    def __init__(self, target: ta.Mapping[K, V]) -> None:
+        super().__init__()
+
+        if target is None:
+            raise TypeError(target)
+        self._target = target
+
+    def __repr__(self) -> str:
+        return '%s(%r)' % (type(self).__name__, self._target)
+
+    def __contains__(self, o: object) -> bool:
+        return o in self._target
+
+    def __eq__(self, o: object) -> bool:
+        return self._target == o
+
+    def __ge__(self, other: object) -> bool:
+        return self._target >= other  # type: ignore
+
+    def __getitem__(self, k: K) -> V:
+        return self._target[k]
+
+    def __gt__(self, other: object) -> bool:
+        return self._target > other  # type: ignore
+
+    def __iter__(self) -> ta.Iterator[T]:
+        return iter(self._target)  # type: ignore
+
+    def __le__(self, other: object) -> bool:
+        return self._target <= other  # type: ignore
+
+    def __len__(self) -> int:
+        return len(self._target)
+
+    def __lt__(self, other: object) -> bool:
+        return self._target < other  # type: ignore
+
+    def __ne__(self, o: object) -> bool:
+        return self._target != o
+
+    def get(self, k: K, default=None) -> ta.Optional[V]:  # type: ignore
+        return self._target.get(k)
+
+    def items(self) -> ta.ItemsView[K, V]:
+        return self._target.items()
+
+    def keys(self) -> ta.KeysView[K]:
+        return self._target.keys()
+
+    def values(self) -> ta.ValuesView[V]:
+        return self._target.values()

omlish/collections/utils.py

@@ -0,0 +1,110 @@
+import functools
+import itertools
+import typing as ta
+
+from .. import check
+from .identity import IdentityKeyDict
+from .identity import IdentitySet
+
+
+T = ta.TypeVar('T')
+K = ta.TypeVar('K')
+V = ta.TypeVar('V')
+
+
+def mut_toposort(data: dict[T, set[T]]) -> ta.Iterator[set[T]]:
+    for k, v in data.items():
+        v.discard(k)
+    extra_items_in_deps = functools.reduce(set.union, data.values()) - set(data.keys())
+    data.update({item: set() for item in extra_items_in_deps})
+    while True:
+        ordered = set(item for item, dep in data.items() if not dep)
+        if not ordered:
+            break
+        yield ordered
+        data = {item: (dep - ordered) for item, dep in data.items() if item not in ordered}
+    if data:
+        raise ValueError('Cyclic dependencies exist among these items: ' + ' '.join(repr(x) for x in data.items()))
+
+
+def toposort(data: ta.Mapping[T, ta.AbstractSet[T]]) -> ta.Iterator[set[T]]:
+    return mut_toposort({k: set(v) for k, v in data.items()})
+
+
+def partition(items: ta.Iterable[T], pred: ta.Callable[[T], bool]) -> tuple[list[T], list[T]]:
+    t: list[T] = []
+    f: list[T] = []
+    for e in items:
+        if pred(e):
+            t.append(e)
+        else:
+            f.append(e)
+    return t, f
+
+
+def unique(it: ta.Iterable[T], *, identity: bool = False) -> list[T]:
+    if isinstance(it, str):
+        raise TypeError(it)
+    ret: list[T] = []
+    seen: ta.MutableSet[T] = IdentitySet() if identity else set()
+    for e in it:
+        if e not in seen:
+            seen.add(e)
+            ret.append(e)
+    return ret
+
+
+def unique_dict(items: ta.Iterable[tuple[K, V]], *, identity: bool = False) -> ta.MutableMapping[K, V]:
+    dct: ta.MutableMapping[K, V] = IdentityKeyDict() if identity else {}
+    for k, v in items:
+        if k in dct:
+            raise KeyError(k)
+        dct[k] = v
+    return dct
+
+
+def all_equal(it: ta.Iterable[T]) -> bool:
+    i = iter(it)
+    try:
+        l = next(i)
+    except StopIteration:
+        return True
+    return all(r == l for r in i)
+
+
+def all_not_equal(it: ta.Iterable[T]) -> bool:
+    s = set()
+    for v in it:
+        if v in s:
+            return False
+        s.add(v)
+    return True
+
+
+def key_cmp(fn: ta.Callable[[K, K], int]) -> ta.Callable[[tuple[K, V], tuple[K, V]], int]:
+    return lambda t0, t1: fn(t0[0], t1[0])
+
+
+def indexes(it: ta.Iterable[T]) -> dict[T, int]:
+    return {e: i for i, e in enumerate(it)}
+
+
+def mut_unify_sets(sets: ta.Iterable[set[T]]) -> list[set[T]]:
+    rem: list[set[T]] = list(sets)
+    ret: list[set[T]] = []
+    while rem:
+        cur = rem.pop()
+        while True:
+            moved = False
+            for i in range(len(rem) - 1, -1, -1):
+                if any(e in cur for e in rem[i]):
+                    cur.update(rem.pop(i))
+                    moved = True
+            if not moved:
+                break
+        ret.append(cur)
+    if ret:
+        all_ = set(itertools.chain.from_iterable(ret))
+        num = sum(map(len, ret))
+        check.equal(len(all_), num)
+    return ret