Persistent stack in Python
An immutable linked list in which push shares the existing tail and pop returns the tail.
- File
persistent_stack.py, 72 lines- Level
- I0 not checked yet
- Toolchain
- Python 3.14, pytest 9
- Test framework
- pytest
- Run the tests
python3 -m pytest
# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/persistent-stack/. The authors are listed in prov/WF0RT5.jsonl.
from __future__ import annotations
from typing import Generic, Iterator, TypeVar
T = TypeVar("T")
class _Node(Generic[T]):
"""One cell of the list. Its fields never change after it is made."""
__slots__ = ("item", "next", "size")
def __init__(self, item: T, nxt: _Node[T] | None) -> None:
self.item = item
self.next = nxt
self.size = 1 + (nxt.size if nxt is not None else 0)
class PersistentStack(Generic[T]):
"""An immutable stack. Each update returns a new stack and keeps the old one.
push, pop, top, is_empty and size take O(1) worst-case time. A push adds one
node; the new stack shares all nodes of the old one.
"""
__slots__ = ("_head",)
def __init__(self, _head: _Node[T] | None = None) -> None:
self._head = _head
def push(self, x: T) -> PersistentStack[T]:
"""Return the stack with x on top of self. O(1)."""
return PersistentStack(_Node(x, self._head)) # @line L1
def pop(self) -> tuple[T, PersistentStack[T]]:
"""Return (top item, the stack below it). Raise IndexError on the empty stack. O(1)."""
s = self._head
if s is None: # @line L2
raise IndexError("pop from an empty stack")
return s.item, PersistentStack(s.next) # @line L3
def top(self) -> T:
"""Return the top item. Raise IndexError on the empty stack. O(1)."""
s = self._head
if s is None: # @line L4
raise IndexError("top of an empty stack")
return s.item # @line L5
def is_empty(self) -> bool:
"""True when the stack has no item. O(1)."""
return self._head is None
def size(self) -> int:
"""The number of items. O(1), because each node stores the size of its list."""
return self._head.size if self._head is not None else 0
def __iter__(self) -> Iterator[T]:
"""The items from the top down, with a loop (no recursion)."""
s = self._head
while s is not None:
yield s.item
s = s.next
def check_invariant(self) -> None:
"""Check that each stored size is 1 plus the size of the next node."""
s = self._head
while s is not None:
assert s.size == 1 + (s.next.size if s.next is not None else 0)
s = s.next
test_persistent_stack.py 70 lines
# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/persistent-stack/. The authors are listed in prov/WF0RT5.jsonl.
import random
import pytest
from persistent_stack import PersistentStack
def test_example() -> None:
v0: PersistentStack[int] = PersistentStack()
v1 = v0.push(1)
v2 = v1.push(2)
v3 = v1.push(3)
x, v4 = v2.pop()
v5 = v4.push(4)
assert x == 2
assert [list(v) for v in (v0, v1, v2, v3, v4, v5)] == [[], [1], [2, 1], [3, 1], [1], [4, 1]]
# v2, v3 and v5 share the node of 1, and v4 is that node.
assert v2._head is not None and v3._head is not None
assert v2._head.next is v3._head.next is v1._head is v4._head
def test_edges() -> None:
e: PersistentStack[str] = PersistentStack()
assert e.is_empty() and e.size() == 0 and list(e) == []
with pytest.raises(IndexError):
e.pop()
with pytest.raises(IndexError):
e.top()
one = e.push("a")
assert one.top() == "a" and one.size() == 1 and e.is_empty()
x, back = one.pop()
assert x == "a" and back.is_empty()
def test_deep_stack_has_no_recursion() -> None:
s: PersistentStack[int] = PersistentStack()
for i in range(100000):
s = s.push(i)
assert s.size() == 100000 and s.top() == 99999
assert sum(1 for _ in s) == 100000
def test_random_version_tree() -> None:
rng = random.Random(4127)
for _ in range(500):
versions: list[PersistentStack[int]] = [PersistentStack()]
models: list[list[int]] = [[]]
for _ in range(rng.randint(1, 40)):
j = rng.randrange(len(versions))
s, m = versions[j], models[j]
if m and rng.random() < 0.4:
x, t = s.pop()
assert x == m[-1] and s.top() == m[-1]
versions.append(t)
models.append(m[:-1])
elif not m and rng.random() < 0.2:
with pytest.raises(IndexError):
s.pop()
else:
x = rng.randint(-(2**63), 2**63 - 1)
versions.append(s.push(x))
models.append(m + [x])
versions[-1].check_invariant()
for s, m in zip(versions, models):
assert list(s) == m[::-1]
assert s.size() == len(m)
assert s.is_empty() == (not m)
What the next level needs
- the tests did not run in this build