pyfect-dev 0.1.0__py3-none-any.whl

pyfect/transaction.py

@@ -0,0 +1,340 @@
+"""Transactions over ``TxRef`` values (Effect's ``Effect.tx``).
+
+``tx(eff)`` runs ``eff`` all or nothing: the ``TxRef`` values it reads and
+writes go through a journal, and only a successful run publishes its writes.
+Transactions are optimistic. At the end of a run the journal is checked
+against the refs: if another transaction committed to a ref this one read,
+the run is discarded and the body runs again (a ``"conflict"``). ``tx_retry``
+discards the run on purpose and parks the fiber until one of the refs it read
+is written (a ``"retry"``).
+
+The check and the commit happen in one fiber step, so no other fiber of the
+run sees a half-published transaction. ``tx`` inside a running transaction
+joins it: nested calls share one journal and commit together, and a failure
+caught inside the transaction keeps the writes made before it.
+
+Yield ``Transaction()`` inside a body to read its ``TransactionMeta``.
+"""
+# pyright: reportPrivateUsage=false
+
+from __future__ import annotations
+
+from collections.abc import Callable, Mapping
+from types import MappingProxyType
+from typing import TYPE_CHECKING, Any, Literal, Never, final, overload
+
+from pyfect._core import (
+    Effect,
+    Restore,
+    _Callback,
+    _FailCause,
+    _FromExit,
+    _Locally,
+    _Succeed,
+    _Sync,
+    _WithFiber,
+    uninterruptible_mask,
+)
+from pyfect._effect_args import require_effect
+from pyfect._js import same_value
+from pyfect.cause import Cause
+from pyfect.cause import interrupt as interrupt_cause
+from pyfect.clock import Clock
+from pyfect.context import Tag
+from pyfect.exit import Exit, Success
+
+if TYPE_CHECKING:
+    from pyfect._fiber import Fiber
+    from pyfect.tx_ref import TxRef
+
+type TransactionRetryReason = Literal["retry", "conflict"]
+"""Why a run of a transaction body started again: ``"retry"`` after
+``tx_retry``, ``"conflict"`` after a ref it read changed before it could
+commit."""
+
+_NANOS = 1_000_000_000
+
+
+@final
+class JournalEntry:
+    """What a transaction knows about one ``TxRef``: the version it saw, its
+    value inside the transaction, and whether the transaction wrote it."""
+
+    __slots__ = ("_value", "_version", "_written")
+
+    def __init__(self, version: int, value: Any) -> None:
+        self._version = version
+        self._value = value
+        self._written = False
+
+    @property
+    def version(self) -> int:
+        return self._version
+
+    @property
+    def value(self) -> Any:
+        return self._value
+
+    @property
+    def written(self) -> bool:
+        return self._written
+
+    def __repr__(self) -> str:
+        return (
+            f"JournalEntry(version={self._version}, value={self._value!r}, written={self._written})"
+        )
+
+
+@final
+class TransactionMeta:
+    """The journal and attempt metadata of a running transaction.
+
+    ``attempt`` counts runs of the outermost body from 1; ``retry_reason`` is
+    ``None`` on the first. Times are wall-clock seconds read from the
+    ``Clock``, as in schedule metadata: ``start`` when the transaction began,
+    ``now`` when this attempt began."""
+
+    __slots__ = (
+        "_attempt",
+        "_completed",
+        "_journal",
+        "_now_ns",
+        "_previous_ns",
+        "_retry",
+        "_retry_reason",
+        "_start_ns",
+    )
+
+    def __init__(self, start_ns: int) -> None:
+        self._journal: dict[TxRef[Any], JournalEntry] = {}
+        self._retry = False
+        self._completed = False
+        self._attempt = 1
+        self._retry_reason: TransactionRetryReason | None = None
+        self._start_ns = start_ns
+        self._now_ns = start_ns
+        self._previous_ns = start_ns
+
+    @property
+    def journal(self) -> Mapping[TxRef[Any], JournalEntry]:
+        return MappingProxyType(self._journal)
+
+    @property
+    def attempt(self) -> int:
+        return self._attempt
+
+    @property
+    def retry_reason(self) -> TransactionRetryReason | None:
+        return self._retry_reason
+
+    @property
+    def start(self) -> float:
+        return self._start_ns / _NANOS
+
+    @property
+    def now(self) -> float:
+        return self._now_ns / _NANOS
+
+    @property
+    def elapsed(self) -> float:
+        return (self._now_ns - self._start_ns) / _NANOS
+
+    @property
+    def elapsed_since_previous(self) -> float:
+        return (self._now_ns - self._previous_ns) / _NANOS
+
+    def _entry(self, ref: TxRef[Any]) -> JournalEntry:
+        entry = self._journal.get(ref)
+        if entry is None:
+            entry = self._journal[ref] = JournalEntry(ref._version, ref._value)
+        return entry
+
+    def _consistent(self) -> bool:
+        return all(ref._version == entry._version for ref, entry in self._journal.items())
+
+    def _next_attempt(self, now_ns: int) -> None:
+        self._attempt += 1
+        self._retry_reason = "retry" if self._retry else "conflict"
+        self._previous_ns = self._now_ns
+        self._now_ns = now_ns
+        self._clear()
+
+    def _commit(self) -> None:
+        for ref, entry in self._journal.items():
+            if not entry._written:
+                continue
+            if not same_value(entry._value, ref._value):
+                ref._version += 1
+                ref._value = entry._value
+            if ref._pending:
+                wakes = list(ref._pending.values())
+                ref._pending.clear()
+                for wake in wakes:
+                    wake()
+
+    def _clear(self) -> None:
+        self._retry = False
+        self._journal.clear()
+
+
+class Transaction(Tag[TransactionMeta]):
+    """The running transaction's state. Nested ``tx`` calls see the same
+    ``TransactionMeta``."""
+
+
+def active_transaction(fiber: Fiber[Any, Any]) -> TransactionMeta | None:
+    """The transaction ``fiber`` is running in, if any. A child fiber
+    inherits its parent's; once that transaction has ended it no longer
+    counts, and the child's next ``tx`` starts its own."""
+    found, state = fiber._context.lookup(Transaction)
+    if not found or state._completed:
+        return None
+    return state
+
+
+@overload
+def tx[A, E](eff: Effect[A, E, Transaction]) -> Effect[A, E]: ...
+@overload
+def tx[A, E, R](eff: Effect[A, E, Transaction | R]) -> Effect[A, E, R]: ...
+def tx(eff: Effect[Any, Any, Any]) -> Effect[Any, Any, Any]:
+    """Run ``eff`` as one transaction, or as part of the running one."""
+    require_effect(eff, "eff")
+    return _WithFiber(
+        lambda fiber: eff if active_transaction(fiber) is not None else _boundary(fiber, eff)
+    )
+
+
+def tx_restoring[A, E](restore: Restore, eff: Effect[A, E, Transaction]) -> Effect[A, E]:
+    """``tx(eff)`` called from an uninterruptible region: the body and its
+    ``tx_retry`` wait run under ``restore``, while the commit and what the
+    caller does next stay uninterruptible, so an interrupt cannot land
+    between committing a claim (a permit, a lock) and guarding it."""
+    require_effect(eff, "eff")
+    return _tx_restoring(restore, eff)
+
+
+def _tx_restoring(restore: Restore, eff: Effect[Any, Any, Any]) -> Effect[Any, Any]:
+    return _WithFiber(
+        lambda fiber: (
+            restore(eff)
+            if active_transaction(fiber) is not None
+            else _boundary(fiber, eff, restore)
+        )
+    )
+
+
+def transactional[A](f: Callable[[TransactionMeta], A]) -> Effect[A]:
+    """``f`` applied to the running transaction, in a transaction of its
+    own when none is running. The building block of ``TxRef``."""
+
+    def body(fiber: Fiber[Any, Any]) -> Effect[A]:
+        state = active_transaction(fiber)
+        if state is not None:
+            return _Sync(lambda: f(state))
+        return _boundary(fiber, Transaction().map(f))
+
+    return _WithFiber(body)
+
+
+def _retry(state: TransactionMeta) -> Effect[Never]:
+    state._retry = True
+    return _WithFiber(lambda fiber: _FailCause(interrupt_cause(fiber.id)))
+
+
+tx_retry: Effect[Never, Never, Transaction] = Transaction().flat_map(_retry)
+"""Abandon this run of the transaction and run it again once a ``TxRef``
+it read has been written."""
+
+
+def _boundary[A, E](
+    fiber: Fiber[Any, Any], eff: Effect[A, E, Any], restore: Restore | None = None
+) -> Effect[A, E]:
+    """Run ``eff`` as the outermost transaction. With ``restore``, the
+    caller is already uninterruptible and ``restore`` makes the body
+    interruptible again; without it, the boundary masks itself."""
+    clock = fiber.get_ref(Clock)
+    state = TransactionMeta(clock.now_nanos_unsafe())
+
+    def run(restore: Restore) -> Effect[A, E]:
+        attempt = (
+            _Locally[A, E, Never](restore(eff), lambda ctx: ctx.add(Transaction, state))
+            .tap_cause(lambda cause: _after_failure(state, cause, restore))
+            .exit()
+        )
+
+        def step(ex: Exit[A, E]) -> Effect[A, E]:
+            if state._retry or not state._consistent():
+                state._next_attempt(clock.now_nanos_unsafe())
+                return attempt.flat_map(step)
+            if isinstance(ex, Success):
+                state._commit()
+            state._clear()
+            state._completed = True
+            return _FromExit(ex)
+
+        return attempt.flat_map(step)
+
+    return uninterruptible_mask(run) if restore is None else run(restore)
+
+
+_VOID: Effect[None] = _Succeed(None)
+
+
+def _after_failure(state: TransactionMeta, cause: Cause[Any], restore: Restore) -> Effect[None]:
+    if not state._retry:
+        return _VOID
+    if not cause.has_interrupts_only:
+        # ``tx_retry`` interrupted the body, but something else failed too
+        # (a finalizer): that failure is the transaction's outcome.
+        state._clear()
+        return _VOID
+    return restore(_await_change(state))
+
+
+def _await_change(state: TransactionMeta) -> Effect[None]:
+    """Park until a ref in the journal is written. The read set is checked
+    and the waiter registered in one step: a commit that landed after the
+    reads has already woken its waiters, so a stale read set reruns at
+    once. A commit wakes the waiter through the waiter's scheduler, not in
+    the committing step (Effect's ``scheduleTask``), so an interrupt sent
+    later in that step still finds it parked."""
+
+    def park(fiber: Fiber[Any, Any]) -> Effect[None]:
+        def register(resume: Callable[[Effect[None]], None]) -> Effect[None] | None:
+            if not state._consistent():
+                resume(_VOID)
+                return None
+            key = object()
+            refs = list(state._journal)
+
+            def clear() -> None:
+                for ref in refs:
+                    ref._pending.pop(key, None)
+
+            def wake() -> None:
+                clear()
+                resume(_VOID)
+
+            def wake_later() -> None:
+                fiber._scheduler.schedule(wake)
+
+            for ref in refs:
+                ref._pending[key] = wake_later
+            return _Sync(clear)
+
+        return _Callback(register)
+
+    return _WithFiber(park)
+
+
+__all__ = [
+    "JournalEntry",
+    "Transaction",
+    "TransactionMeta",
+    "TransactionRetryReason",
+    "active_transaction",
+    "transactional",
+    "tx",
+    "tx_restoring",
+    "tx_retry",
+]

pyfect/trie.py

@@ -0,0 +1,317 @@
+"""``Trie[V]``: an immutable map from strings to values, iterated in key
+order, with prefix queries (Effect's ``Trie``)::
+
+    t = Trie.make(("she", 0), ("shells", 1), ("sea", 2))
+    list(t.keys_with_prefix("she"))  # ["she", "shells"]
+    t.longest_prefix_of("shellsort")  # Some(("shells", 1))
+
+Keys are ordered by their UTF-16 code units, as JavaScript compares strings
+(so a character past U+FFFF sorts below U+E000, where Python's ``<`` puts it
+above). The empty string is never a key: inserting it returns the same
+trie. Every update returns a new trie and leaves this one as it was; like
+Effect's, the trie is persistent, so an update copies only the nodes along
+its key.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Callable, Iterable, Iterator
+from typing import Any, TypeIs, overload
+
+from pyfect._equal import equal_key, equals
+from pyfect._js import js_sort_key
+from pyfect._unsubscripted import Unsubscripted
+from pyfect.option import Option, Some
+from pyfect.result import Ok, Result
+
+
+class _Node:
+    """A node of the character trie; never changed once built. ``count``
+    is the number of keys at or below it."""
+
+    __slots__ = ("_order", "children", "count", "has", "value")
+
+    def __init__(self, has: bool, value: object, children: dict[str, _Node]) -> None:
+        self.has = has
+        self.value = value
+        self.children = children
+        self.count = int(has) + sum(child.count for child in children.values())
+        self._order: tuple[tuple[str, _Node], ...] | None = None
+
+    def ordered(self) -> tuple[tuple[str, _Node], ...]:
+        """The children in UTF-16 order, sorted once."""
+        if self._order is None:
+            self._order = tuple(sorted(self.children.items(), key=lambda kv: js_sort_key(kv[0])))
+        return self._order
+
+
+_EMPTY = _Node(False, None, {})
+
+
+def _descend(node: _Node, key: str) -> _Node | None:
+    for c in key:
+        child = node.children.get(c)
+        if child is None:
+            return None
+        node = child
+    return node
+
+
+def _walk(node: _Node, prefix: str) -> Iterator[tuple[str, object]]:
+    """The entries at and below ``node`` in key order, without recursion."""
+    stack: list[tuple[str, _Node]] = [(prefix, node)]
+    while stack:
+        key, n = stack.pop()
+        if n.has:
+            yield key, n.value
+        stack.extend((key + c, child) for c, child in reversed(n.ordered()))
+
+
+def _replace(root: _Node, key: str, leaf: _Node | None) -> _Node:
+    """``root`` with the node at ``key`` (which exists up to its parent)
+    replaced by ``leaf``; ``None`` drops it, and parents left empty."""
+    path = [root]
+    for c in key[:-1]:
+        path.append(path[-1].children[c])
+    node = leaf
+    for c, parent in zip(reversed(key), reversed(path), strict=True):
+        children = dict(parent.children)
+        if node is None:
+            del children[c]
+        else:
+            children[c] = node
+        node = (
+            None
+            if not children and not parent.has and parent is not root
+            else (_Node(parent.has, parent.value, children))
+        )
+    return _EMPTY if node is None else node
+
+
+class Trie[V = Unsubscripted]:
+    """Build one with ``Trie.empty()``, ``Trie.make(*entries)`` or
+    ``Trie.from_iterable(entries)``."""
+
+    __slots__ = ("_root",)
+
+    _root: _Node
+
+    def __init__(self, root: _Node = _EMPTY) -> None:
+        self._root = root
+
+    @classmethod
+    def empty(cls) -> Trie[V]:
+        """Spelled ``Trie[int].empty()``."""
+        return Trie()
+
+    @overload
+    @classmethod
+    def make[V2](cls: type[Trie[Unsubscripted]], *entries: tuple[str, V2]) -> Trie[V2]: ...
+    @overload
+    @classmethod
+    def make(cls, *entries: tuple[str, V]) -> Trie[V]: ...
+    @classmethod
+    def make[V2](cls, *entries: tuple[str, V2]) -> Trie[V2]:
+        return Trie.from_iterable(entries)
+
+    @overload
+    @classmethod
+    def from_iterable[V2](
+        cls: type[Trie[Unsubscripted]], entries: Iterable[tuple[str, V2]]
+    ) -> Trie[V2]: ...
+    @overload
+    @classmethod
+    def from_iterable(cls, entries: Iterable[tuple[str, V]]) -> Trie[V]: ...
+    @classmethod
+    def from_iterable[V2](cls, entries: Iterable[tuple[str, V2]]) -> Trie[V2]:
+        """The trie of ``entries``; a later value for a key wins."""
+        return Trie[V2]().insert_many(entries)
+
+    # --- reading ------------------------------------------------------------------
+
+    def _node(self, key: str) -> _Node | None:
+        node = _descend(self._root, key)
+        return node if node is not None and node.has else None
+
+    def __len__(self) -> int:
+        return self._root.count
+
+    def __bool__(self) -> bool:
+        return self._root.count > 0
+
+    def __iter__(self) -> Iterator[tuple[str, V]]:
+        return self.entries()
+
+    def __contains__(self, key: object) -> bool:
+        return isinstance(key, str) and self.has(key)
+
+    def is_empty(self) -> bool:
+        return self._root.count == 0
+
+    def get(self, key: str) -> Option[V]:
+        node = self._node(key)
+        return None if node is None else Some(self._value(node))
+
+    def has(self, key: str) -> bool:
+        return key != "" and self._node(key) is not None
+
+    def get_unsafe(self, key: str) -> V:
+        """The value for ``key``; raises ``KeyError`` without one."""
+        node = self._node(key)
+        if node is None:
+            raise KeyError("Expected trie to contain key")
+        return self._value(node)
+
+    def _value(self, node: _Node) -> V:
+        return node.value  # pyright: ignore[reportReturnType]
+
+    def keys(self) -> Iterator[str]:
+        """The keys in order."""
+        return (k for k, _ in self.entries())
+
+    def values(self) -> Iterator[V]:
+        """The values in key order."""
+        return (v for _, v in self.entries())
+
+    def entries(self) -> Iterator[tuple[str, V]]:
+        """The ``(key, value)`` pairs in key order."""
+        return self.entries_with_prefix("")
+
+    def to_entries(self) -> list[tuple[str, V]]:
+        return list(self.entries())
+
+    def keys_with_prefix(self, prefix: str) -> Iterator[str]:
+        """The keys starting with ``prefix``, in order."""
+        return (k for k, _ in self.entries_with_prefix(prefix))
+
+    def values_with_prefix(self, prefix: str) -> Iterator[V]:
+        return (v for _, v in self.entries_with_prefix(prefix))
+
+    def entries_with_prefix(self, prefix: str) -> Iterator[tuple[str, V]]:
+        node = _descend(self._root, prefix)
+        if node is None:
+            return iter(())
+        return _walk(node, prefix)  # pyright: ignore[reportReturnType]
+
+    def to_entries_with_prefix(self, prefix: str) -> list[tuple[str, V]]:
+        return list(self.entries_with_prefix(prefix))
+
+    def longest_prefix_of(self, key: str) -> Option[tuple[str, V]]:
+        """The entry whose key is the longest one that ``key`` starts with."""
+        node, found = self._root, None
+        for end, c in enumerate(key, 1):
+            child = node.children.get(c)
+            if child is None:
+                break
+            node = child
+            if node.has:
+                found = (end, node)
+        if found is None:
+            return None
+        return Some((key[: found[0]], self._value(found[1])))
+
+    def reduce[Z](self, zero: Z, f: Callable[[Z, V, str], Z]) -> Z:
+        """Fold the entries in key order."""
+        out = zero
+        for k, v in self.entries():
+            out = f(out, v, k)
+        return out
+
+    def for_each(self, f: Callable[[V, str], object]) -> None:
+        for k, v in self.entries():
+            f(v, k)
+
+    # --- updates --------------------------------------------------------------------
+
+    def insert(self, key: str, value: V) -> Trie[V]:
+        """The trie with ``key`` set to ``value``; the same trie for ``""``."""
+        if not key:
+            return self
+        root, depth = self._root, 0
+        node: _Node | None = root
+        while depth < len(key) and node is not None:
+            node = node.children.get(key[depth])
+            depth += 1
+        if node is not None:
+            return Trie(_replace(root, key, _Node(True, value, node.children)))
+        # ``key[:depth - 1]`` exists: grow the rest as a chain below it.
+        leaf = _Node(True, value, {})
+        for c in reversed(key[depth:]):
+            leaf = _Node(False, None, {c: leaf})
+        return Trie(_replace(root, key[:depth], leaf))
+
+    def insert_many(self, entries: Iterable[tuple[str, V]]) -> Trie[V]:
+        out = self
+        for key, value in entries:
+            out = out.insert(key, value)
+        return out
+
+    def remove(self, key: str) -> Trie[V]:
+        """The trie without ``key``; the same trie when it has no ``key``."""
+        node = self._node(key) if key else None
+        if node is None:
+            return self
+        leaf = _Node(False, None, node.children) if node.children else None
+        return Trie(_replace(self._root, key, leaf))
+
+    def remove_many(self, keys: Iterable[str]) -> Trie[V]:
+        out = self
+        for key in keys:
+            out = out.remove(key)
+        return out
+
+    def modify(self, key: str, f: Callable[[V], V]) -> Trie[V]:
+        """The trie with ``key``'s value replaced by ``f(value)``; the same
+        trie without ``key``."""
+        node = self._node(key) if key else None
+        if node is None:
+            return self
+        return self.insert(key, f(self._value(node)))
+
+    def _rebuilt[A](self, entries: Iterable[tuple[str, A]]) -> Trie[A]:
+        return Trie[A]().insert_many(list(entries))
+
+    def map[A](self, f: Callable[[V, str], A]) -> Trie[A]:
+        return self._rebuilt((k, f(v, k)) for k, v in self.entries())
+
+    @overload
+    def filter[B](self, predicate: Callable[[V, str], TypeIs[B]]) -> Trie[B]: ...
+    @overload
+    def filter(self, predicate: Callable[[V, str], bool]) -> Trie[V]: ...
+    def filter(self, predicate: Callable[[V, str], bool]) -> Trie[Any]:
+        return self._rebuilt((k, v) for k, v in self.entries() if predicate(v, k))
+
+    def filter_map[A](self, f: Callable[[V, str], Result[A, object]]) -> Trie[A]:
+        """The entries ``f(value, key)`` keeps (``Ok``), with its values."""
+
+        def kept() -> Iterator[tuple[str, A]]:
+            for k, v in self.entries():
+                result = f(v, k)
+                if isinstance(result, Ok):
+                    yield k, result.value
+
+        return self._rebuilt(kept())
+
+    def compact[A](self: Trie[Option[A]]) -> Trie[A]:
+        """The ``Some`` values, unwrapped."""
+        return self._rebuilt((k, v.value) for k, v in self.entries() if isinstance(v, Some))
+
+    # --- equality -------------------------------------------------------------------
+
+    def __eq__(self, other: object) -> bool:
+        if not isinstance(other, Trie):
+            return NotImplemented
+        theirs = Trie[object].to_entries(other)  # pyright: ignore[reportUnknownArgumentType]
+        mine = self.to_entries()
+        return len(mine) == len(theirs) and all(
+            a[0] == b[0] and equals(a[1], b[1]) for a, b in zip(mine, theirs, strict=True)
+        )
+
+    def __hash__(self) -> int:
+        return hash(tuple((k, equal_key(v)) for k, v in self.entries()))
+
+    def __repr__(self) -> str:
+        return f"Trie({self.to_entries()!r})"
+
+
+__all__ = ["Trie"]