Ctrie (concurrent hash-array mapped trie) in Python
Lock-free hash trie with cheap atomic snapshots, supporting linearizable lookups and insertions.
- File
ctrie.py, 320 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/ctrie/. The authors are listed in prov/3YJ33R.jsonl.
import threading
import time
from collections.abc import Callable, Generator, Hashable
from typing import Any, TypeVar
R = TypeVar("R")
SPIN = "spin"
NOTFOUND = "NOTFOUND"
RESTART = "RESTART"
Steps = Generator[str | None, None, R]
def run(steps: "Steps[R]") -> R:
"""Run a step generator to its end in the calling thread and return its result."""
while True:
try:
if next(steps) == SPIN:
time.sleep(0)
except StopIteration as stop:
return stop.value
class SNode:
"""A key and its value. Immutable."""
__slots__ = ("key", "value", "hash")
def __init__(self, key: Hashable, value: Any, h: int) -> None:
self.key, self.value, self.hash = key, value, h
class TNode:
"""A tomb: the single S-node left in an I-node that is to be removed. Immutable."""
__slots__ = ("sn",)
def __init__(self, sn: SNode) -> None:
self.sn = sn
class LNode:
"""The S-nodes whose hashes are equal in every bit. Immutable."""
__slots__ = ("sns",)
def __init__(self, sns: tuple[SNode, ...]) -> None:
self.sns = sns
def find(self, key: Hashable) -> Any:
for sn in self.sns:
if sn.key == key:
return sn.value
return NOTFOUND
class CNode:
"""A bitmap and the array of branches (I-nodes or S-nodes) for the set bits. Immutable."""
__slots__ = ("bmp", "array")
def __init__(self, bmp: int, array: tuple[Any, ...]) -> None:
self.bmp, self.array = bmp, array
def inserted(self, pos: int, flag: int, b: Any) -> "CNode":
return CNode(self.bmp | flag, self.array[:pos] + (b,) + self.array[pos:])
def updated(self, pos: int, b: Any) -> "CNode":
return CNode(self.bmp, self.array[:pos] + (b,) + self.array[pos + 1:])
def removed(self, pos: int, flag: int) -> "CNode":
return CNode(self.bmp ^ flag, self.array[:pos] + self.array[pos + 1:])
class INode:
"""The only mutable node: one word main, read and changed by compare-and-swap.
The compare-and-swap compares by identity and is simulated with a lock.
"""
__slots__ = ("_main", "_lock")
def __init__(self, main: Any) -> None:
self._main = main
self._lock = threading.Lock()
def read(self) -> Any:
with self._lock:
return self._main
def cas(self, expect: Any, new: Any) -> bool:
with self._lock:
if self._main is expect:
self._main = new
return True
return False
class Ctrie:
"""The Ctrie of Prokopec, Bagwell and Odersky, without snapshots.
Each C-node branches on W bits of the hash; when all hbits bits are used,
keys with equal hashes share an L-node. Every read of an I-node and every
compare-and-swap is one step of a generator.
"""
def __init__(self, w: int = 5, hbits: int = 30, hash_fn: Callable[[Hashable], int] | None = None) -> None:
self.w, self.hbits = w, hbits
self.mask = (1 << w) - 1
self.hash_fn = hash_fn or hash
self.root = INode(CNode(0, ()))
def _hash(self, key: Hashable) -> int:
return self.hash_fn(key) & ((1 << self.hbits) - 1)
def _flagpos(self, h: int, lev: int, bmp: int) -> tuple[int, int]:
flag = 1 << ((h >> lev) & self.mask)
return flag, bin(bmp & (flag - 1)).count("1")
def _dual(self, x: SNode, y: SNode, lev: int) -> Any:
"""The main node that holds two S-nodes whose hashes agree below level lev."""
if lev >= self.hbits:
return LNode((x, y))
fx, fy = 1 << ((x.hash >> lev) & self.mask), 1 << ((y.hash >> lev) & self.mask)
if fx == fy:
return CNode(fx, (INode(self._dual(x, y, lev + self.w)),))
return CNode(fx | fy, (x, y) if fx < fy else (y, x))
def _contract(self, cn: CNode, lev: int) -> Any:
if lev > 0 and len(cn.array) == 1 and isinstance(cn.array[0], SNode):
return TNode(cn.array[0])
return cn
def _compress(self, cn: CNode, lev: int) -> "Steps[Any]":
branches = []
for b in cn.array:
if isinstance(b, INode):
m = b.read()
yield None
if isinstance(m, TNode):
b = m.sn # resurrect
branches.append(b)
return self._contract(CNode(cn.bmp, tuple(branches)), lev)
def _clean_steps(self, p: INode | None, lev: int) -> "Steps[None]":
if p is None:
return None
m = p.read() # @line L26
yield None
if isinstance(m, CNode):
new = yield from self._compress(m, lev)
p.cas(m, new)
yield None
return None
def _clean_parent_steps(self, p: INode | None, i: INode, h: int, lev: int) -> "Steps[None]":
while p is not None:
m = p.read() # @line L27
yield None
if not isinstance(m, CNode):
return None
flag, pos = self._flagpos(h, lev, m.bmp)
if not m.bmp & flag or m.array[pos] is not i:
return None
im = i.read() # @line L28
yield None
if not isinstance(im, TNode):
return None
ok = p.cas(m, self._contract(m.updated(pos, im.sn), lev))
yield None
if ok:
return None
return None
# ---- lookup ----
def lookup_steps(self, key: Hashable) -> "Steps[Any]":
h = self._hash(key)
while True:
r = yield from self._ilookup(self.root, key, h, 0, None) # @line L1
if r is not RESTART:
return r
def _ilookup(self, i: INode, key: Hashable, h: int, lev: int, parent: INode | None) -> "Steps[Any]":
m = i.read() # @line L2
yield None
if isinstance(m, CNode): # @line L3
flag, pos = self._flagpos(h, lev, m.bmp)
if not m.bmp & flag:
return NOTFOUND
b = m.array[pos] # @line L4
if isinstance(b, INode):
return (yield from self._ilookup(b, key, h, lev + self.w, i))
return b.value if b.key == key else NOTFOUND # @line L5
if isinstance(m, TNode): # @line L6
yield from self._clean_steps(parent, lev - self.w)
return RESTART
return m.find(key) # @line L7
# ---- insert ----
def insert_steps(self, key: Hashable, value: Any) -> "Steps[None]":
h = self._hash(key)
while True:
r = yield from self._iinsert(self.root, SNode(key, value, h), 0, None) # @line L8
if r is not RESTART:
return None
def _iinsert(self, i: INode, sn: SNode, lev: int, parent: INode | None) -> "Steps[Any]":
m = i.read() # @line L9
yield None
if isinstance(m, CNode): # @line L10
flag, pos = self._flagpos(sn.hash, lev, m.bmp)
if not m.bmp & flag:
ok = i.cas(m, m.inserted(pos, flag, sn))
yield None
return "OK" if ok else RESTART
b = m.array[pos] # @line L11
if isinstance(b, INode):
return (yield from self._iinsert(b, sn, lev + self.w, i))
n = sn if b.key == sn.key else INode(self._dual(b, sn, lev + self.w)) # @line L12
ok = i.cas(m, m.updated(pos, n)) # @line L13
yield None
return "OK" if ok else RESTART
if isinstance(m, TNode): # @line L14
yield from self._clean_steps(parent, lev - self.w)
return RESTART
rest = tuple(x for x in m.sns if x.key != sn.key) # @line L15
ok = i.cas(m, LNode((sn,) + rest))
yield None
return "OK" if ok else RESTART
# ---- remove ----
def remove_steps(self, key: Hashable) -> "Steps[Any]":
h = self._hash(key)
while True:
r = yield from self._iremove(self.root, key, h, 0, None) # @line L16
if r is not RESTART:
return r
def _iremove(self, i: INode, key: Hashable, h: int, lev: int, parent: INode | None) -> "Steps[Any]":
m = i.read() # @line L17
yield None
if isinstance(m, CNode): # @line L18
flag, pos = self._flagpos(h, lev, m.bmp)
if not m.bmp & flag:
return NOTFOUND
b = m.array[pos]
if isinstance(b, INode): # @line L19
r = yield from self._iremove(b, key, h, lev + self.w, i)
elif b.key != key: # @line L20
r = NOTFOUND
else:
ok = i.cas(m, self._contract(m.removed(pos, flag), lev)) # @line L21
yield None
r = b.value if ok else RESTART
if r is not NOTFOUND and r is not RESTART: # @line L22
im = i.read()
yield None
if isinstance(im, TNode):
yield from self._clean_parent_steps(parent, i, h, lev - self.w)
return r # @line L23
if isinstance(m, TNode): # @line L24
yield from self._clean_steps(parent, lev - self.w)
return RESTART
v = m.find(key) # @line L25
if v is NOTFOUND:
return NOTFOUND
rest = tuple(x for x in m.sns if x.key != key)
new = TNode(rest[0]) if len(rest) == 1 else LNode(rest)
ok = i.cas(m, new)
yield None
return v if ok else RESTART
# ---- plain methods and checks ----
def lookup(self, key: Hashable) -> Any:
return run(self.lookup_steps(key))
def insert(self, key: Hashable, value: Any) -> None:
run(self.insert_steps(key, value))
def remove(self, key: Hashable) -> Any:
return run(self.remove_steps(key))
def items(self) -> dict[Hashable, Any]:
"""The abstract map: every S-node reachable from the root. Checks the shape on the way."""
out: dict[Hashable, Any] = {}
def add(sn: SNode, lev: int, prefix: int) -> None:
assert sn.hash & ((1 << lev) - 1) == prefix, "key on the wrong path"
assert sn.key not in out, "duplicate key"
out[sn.key] = sn.value
def walk(i: INode, lev: int, prefix: int) -> None:
m = i.read()
if isinstance(m, CNode):
assert len(m.array) == bin(m.bmp).count("1")
idxs = [k for k in range(self.mask + 1) if m.bmp >> k & 1]
for idx, b in zip(idxs, m.array):
p2 = prefix | (idx << lev)
if isinstance(b, INode):
walk(b, lev + self.w, p2)
else:
add(b, lev + self.w, p2)
elif isinstance(m, TNode):
assert lev > 0, "a tomb at the root"
add(m.sn, lev, prefix)
else:
assert lev >= self.hbits and len(m.sns) >= 2
for sn in m.sns:
add(sn, self.hbits, prefix & ((1 << self.hbits) - 1))
walk(self.root, 0, 0)
return out
test_ctrie.py 225 lines
# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/ctrie/. The authors are listed in prov/3YJ33R.jsonl.
import random
import threading
from ctrie import NOTFOUND, CNode, Ctrie, INode, LNode, TNode
# ---- deterministic interleaving driver and linearizability check ----
Op = tuple # (thread, name, arg, result, invoke time, response time)
def interleave(threads, schedule, limit=100000):
"""Run the operations of each thread, one atomic step at a time.
threads[t] is a list of (name, arg, make) where make() returns the step
generator of the operation. schedule is a sequence of thread numbers; a
pick of a thread with no work left is skipped. After the schedule, the
remaining steps run round robin. Returns the history as a list of Op.
"""
todo = [list(reversed(ops)) for ops in threads]
cur: list = [None] * len(threads)
history: list = []
clock = 0
def step(t):
nonlocal clock
if cur[t] is None:
if not todo[t]:
return False
name, arg, make = todo[t].pop()
cur[t] = (name, arg, make(), clock)
name, arg, gen, inv = cur[t]
clock += 1
try:
next(gen)
except StopIteration as stop:
history.append((t, name, arg, stop.value, inv, clock))
cur[t] = None
return True
for t in schedule:
step(t)
steps = 0
while any(cur) or any(todo):
for t in range(len(threads)):
step(t)
steps += 1
assert steps < limit, "no progress"
return history
def linearizable(history, init, apply):
"""Wing and Gong search: is there a legal order that respects real time?
apply(state, name, arg) returns (new state, result); states are hashable.
"""
n = len(history)
full = (1 << n) - 1
seen = set()
def dfs(done, state):
if done == full:
return True
if (done, state) in seen:
return False
seen.add((done, state))
for i in range(n):
if done >> i & 1:
continue
inv_i = history[i][4]
if any(not done >> j & 1 and history[j][5] < inv_i for j in range(n)):
continue
s2, r = apply(state, history[i][1], history[i][2])
if r == history[i][3] and dfs(done | 1 << i, s2):
return True
return False
return dfs(0, init)
def map_apply(state, name, arg):
"""The sequential map on a frozenset of (key, value) pairs."""
d = dict(state)
if name == "lookup":
return state, d.get(arg, NOTFOUND)
if name == "insert":
k, v = arg
d[k] = v
return frozenset(d.items()), None
r = d.pop(arg, NOTFOUND)
return frozenset(d.items()), r
def finish(gen):
try:
while True:
next(gen)
except StopIteration as stop:
return stop.value
def small(**kw):
"""A trie with 4-way C-nodes and 4-bit hashes, so that paths and L-nodes are short."""
return Ctrie(w=2, hbits=4, hash_fn=lambda k: k, **kw)
# ---- tests ----
def test_example():
t = small()
for k in (0, 4, 1):
t.insert(k, str(k))
root = t.root.read()
assert root.bmp == 0b11 and isinstance(root.array[0], INode) and root.array[1].key == 1
inner = root.array[0].read() # 0 and 4 agree in the low 2 bits, and split on the next 2
assert [sn.key for sn in inner.array] == [0, 4]
assert t.lookup(4) == "4" and t.lookup(8) == NOTFOUND
rem = t.remove_steps(4)
while not isinstance(inner_main := root.array[0].read(), TNode):
next(rem) # the removal contracts the inner C-node to a tomb holding 0
assert inner_main.sn.key == 0
look = t.lookup_steps(0)
next(look) # the lookup has read the old root C-node and will meet the tomb
assert finish(look) == "0" # it cleans the root, which resurrects 0, and restarts
assert finish(rem) == "4"
root = t.root.read()
assert [b.key for b in root.array] == [0, 1] # the tomb is gone
assert t.items() == {0: "0", 1: "1"}
def test_full_collisions_use_an_lnode():
t = small()
for k in (3, 19, 35): # all equal to 3 in the low 4 bits
t.insert(k, k)
path = t.root.read().array[0].read().array[0].read() # two C-nodes deep, the hash bits run out
assert isinstance(path, LNode) and len(path.sns) == 3
assert t.lookup(19) == 19 and t.lookup(51) == NOTFOUND
assert t.remove(19) == 19 and t.remove(35) == 35
assert t.lookup(3) == 3 and t.items() == {3: 3}
def test_edges():
t = small()
assert t.lookup(1) == NOTFOUND and t.remove(1) == NOTFOUND
t.insert(1, "a")
t.insert(1, "b") # an update replaces the value
assert t.lookup(1) == "b" and t.items() == {1: "b"}
assert t.remove(1) == "b" and t.items() == {}
assert isinstance(t.root.read(), CNode) # the root is never a tomb
d = Ctrie()
for k in ("x", "y", 7):
d.insert(k, k)
assert d.lookup("y") == "y" and d.remove("x") == "x" and d.items() == {"y": "y", 7: 7}
def test_random_sequential_against_dict():
for seed in range(20):
rng = random.Random(seed)
t, model = small(), {}
for _ in range(400):
k = rng.randrange(64)
op = rng.random()
if op < 0.4:
t.insert(k, seed)
model[k] = seed
elif op < 0.75:
assert t.remove(k) == model.pop(k, NOTFOUND)
else:
assert t.lookup(k) == model.get(k, NOTFOUND)
assert t.items() == model
def test_random_interleavings_are_linearizable():
for seed in range(400):
rng = random.Random(seed)
t = small()
init = {}
for _ in range(rng.randrange(6)):
k = rng.choice((0, 4, 8, 1, 16, 32))
t.insert(k, -1)
init[k] = -1
keys = (0, 4, 8, 1, 16, 32)
work = []
for th in range(3):
ops = []
for _ in range(rng.randrange(1, 4)):
k, r = rng.choice(keys), rng.random()
if r < 0.4:
ops.append(("insert", (k, th), lambda k=k, th=th: t.insert_steps(k, th)))
elif r < 0.75:
ops.append(("remove", k, lambda k=k: t.remove_steps(k)))
else:
ops.append(("lookup", k, lambda k=k: t.lookup_steps(k)))
work.append(ops)
history = interleave(work, [rng.randrange(3) for _ in range(rng.randrange(60))])
final = t.items() # also checks the shape
history.append((9, "items", None, None, max(h[5] for h in history) + 1, 10 ** 9))
def apply(state, name, arg):
if name == "items":
return state, None if dict(state) == final else "different"
return map_apply(state, name, arg)
assert linearizable(history, frozenset(init.items()), apply), seed
def test_threads():
t = Ctrie(w=3, hbits=9, hash_fn=lambda k: k * 2654435761)
def worker(w):
for k in range(w, 2000, 4):
t.insert(k, w)
for k in range(w, 2000, 8):
assert t.remove(k) == w
ts = [threading.Thread(target=worker, args=(w,)) for w in range(4)]
for th in ts:
th.start()
for th in ts:
th.join()
expect = {k: k % 4 for k in range(2000) if k % 8 >= 4}
assert t.items() == expect
assert all(t.lookup(k) == v for k, v in expect.items())
What the next level needs
- the tests did not run in this build
Read the full page about Ctrie (concurrent hash-array mapped trie)