btree.py 30 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850
  1. # Copyright (C) Dnspython Contributors, see LICENSE for text of ISC license
  2. """
  3. A BTree in the style of Cormen, Leiserson, and Rivest's "Algorithms" book, with
  4. copy-on-write node updates, cursors, and optional space optimization for mostly-in-order
  5. insertion.
  6. """
  7. from collections.abc import MutableMapping, MutableSet
  8. from typing import Any, Callable, Generic, Optional, Tuple, TypeVar, cast
  9. DEFAULT_T = 127
  10. KT = TypeVar("KT") # the type of a key in Element
  11. class Element(Generic[KT]):
  12. """All items stored in the BTree are Elements."""
  13. def key(self) -> KT:
  14. """The key for this element; the returned type must implement comparison."""
  15. raise NotImplementedError # pragma: no cover
  16. ET = TypeVar("ET", bound=Element) # the type of a value in a _KV
  17. def _MIN(t: int) -> int:
  18. """The minimum number of keys in a non-root node for a BTree with the specified
  19. ``t``
  20. """
  21. return t - 1
  22. def _MAX(t: int) -> int:
  23. """The maximum number of keys in node for a BTree with the specified ``t``"""
  24. return 2 * t - 1
  25. class _Creator:
  26. """A _Creator class instance is used as a unique id for the BTree which created
  27. a node.
  28. We use a dedicated creator rather than just a BTree reference to avoid circularity
  29. that would complicate GC.
  30. """
  31. def __str__(self): # pragma: no cover
  32. return f"{id(self):x}"
  33. class _Node(Generic[KT, ET]):
  34. """A Node in the BTree.
  35. A Node (leaf or internal) of the BTree.
  36. """
  37. __slots__ = ["t", "creator", "is_leaf", "elts", "children"]
  38. def __init__(self, t: int, creator: _Creator, is_leaf: bool):
  39. assert t >= 3
  40. self.t = t
  41. self.creator = creator
  42. self.is_leaf = is_leaf
  43. self.elts: list[ET] = []
  44. self.children: list[_Node[KT, ET]] = []
  45. def is_maximal(self) -> bool:
  46. """Does this node have the maximal number of keys?"""
  47. assert len(self.elts) <= _MAX(self.t)
  48. return len(self.elts) == _MAX(self.t)
  49. def is_minimal(self) -> bool:
  50. """Does this node have the minimal number of keys?"""
  51. assert len(self.elts) >= _MIN(self.t)
  52. return len(self.elts) == _MIN(self.t)
  53. def search_in_node(self, key: KT) -> tuple[int, bool]:
  54. """Get the index of the ``Element`` matching ``key`` or the index of its
  55. least successor.
  56. Returns a tuple of the index and an ``equal`` boolean that is ``True`` iff.
  57. the key was found.
  58. """
  59. l = len(self.elts)
  60. if l > 0 and key > self.elts[l - 1].key():
  61. # This is optimizing near in-order insertion.
  62. return l, False
  63. l = 0
  64. i = len(self.elts)
  65. r = i - 1
  66. equal = False
  67. while l <= r:
  68. m = (l + r) // 2
  69. k = self.elts[m].key()
  70. if key == k:
  71. i = m
  72. equal = True
  73. break
  74. elif key < k:
  75. i = m
  76. r = m - 1
  77. else:
  78. l = m + 1
  79. return i, equal
  80. def maybe_cow_child(self, index: int) -> "_Node[KT, ET]":
  81. assert not self.is_leaf
  82. child = self.children[index]
  83. cloned = child.maybe_cow(self.creator)
  84. if cloned:
  85. self.children[index] = cloned
  86. return cloned
  87. else:
  88. return child
  89. def _get_node(self, key: KT) -> Tuple[Optional["_Node[KT, ET]"], int]:
  90. """Get the node associated with key and its index, doing
  91. copy-on-write if we have to descend.
  92. Returns a tuple of the node and the index, or the tuple ``(None, 0)``
  93. if the key was not found.
  94. """
  95. i, equal = self.search_in_node(key)
  96. if equal:
  97. return (self, i)
  98. elif self.is_leaf:
  99. return (None, 0)
  100. else:
  101. child = self.maybe_cow_child(i)
  102. return child._get_node(key)
  103. def get(self, key: KT) -> ET | None:
  104. """Get the element associated with *key* or return ``None``"""
  105. i, equal = self.search_in_node(key)
  106. if equal:
  107. return self.elts[i]
  108. elif self.is_leaf:
  109. return None
  110. else:
  111. return self.children[i].get(key)
  112. def optimize_in_order_insertion(self, index: int) -> None:
  113. """Try to minimize the number of Nodes in a BTree where the insertion
  114. is done in-order or close to it, by stealing as much as we can from our
  115. right sibling.
  116. If we don't do this, then an in-order insertion will produce a BTree
  117. where most of the nodes are minimal.
  118. """
  119. if index == 0:
  120. return
  121. left = self.children[index - 1]
  122. if len(left.elts) == _MAX(self.t):
  123. return
  124. left = self.maybe_cow_child(index - 1)
  125. while len(left.elts) < _MAX(self.t):
  126. if not left.try_right_steal(self, index - 1):
  127. break
  128. def insert_nonfull(self, element: ET, in_order: bool) -> ET | None:
  129. assert not self.is_maximal()
  130. while True:
  131. key = element.key()
  132. i, equal = self.search_in_node(key)
  133. if equal:
  134. # replace
  135. old = self.elts[i]
  136. self.elts[i] = element
  137. return old
  138. elif self.is_leaf:
  139. self.elts.insert(i, element)
  140. return None
  141. else:
  142. child = self.maybe_cow_child(i)
  143. if child.is_maximal():
  144. self.adopt(*child.split())
  145. # Splitting might result in our target moving to us, so
  146. # search again.
  147. continue
  148. oelt = child.insert_nonfull(element, in_order)
  149. if in_order:
  150. self.optimize_in_order_insertion(i)
  151. return oelt
  152. def split(self) -> tuple["_Node[KT, ET]", ET, "_Node[KT, ET]"]:
  153. """Split a maximal node into two minimal ones and a central element."""
  154. assert self.is_maximal()
  155. right = self.__class__(self.t, self.creator, self.is_leaf)
  156. right.elts = list(self.elts[_MIN(self.t) + 1 :])
  157. middle = self.elts[_MIN(self.t)]
  158. self.elts = list(self.elts[: _MIN(self.t)])
  159. if not self.is_leaf:
  160. right.children = list(self.children[_MIN(self.t) + 1 :])
  161. self.children = list(self.children[: _MIN(self.t) + 1])
  162. return self, middle, right
  163. def try_left_steal(self, parent: "_Node[KT, ET]", index: int) -> bool:
  164. """Try to steal from this Node's left sibling for balancing purposes.
  165. Returns ``True`` if the theft was successful, or ``False`` if not.
  166. """
  167. if index != 0:
  168. left = parent.children[index - 1]
  169. if not left.is_minimal():
  170. left = parent.maybe_cow_child(index - 1)
  171. elt = parent.elts[index - 1]
  172. parent.elts[index - 1] = left.elts.pop()
  173. self.elts.insert(0, elt)
  174. if not left.is_leaf:
  175. assert not self.is_leaf
  176. child = left.children.pop()
  177. self.children.insert(0, child)
  178. return True
  179. return False
  180. def try_right_steal(self, parent: "_Node[KT, ET]", index: int) -> bool:
  181. """Try to steal from this Node's right sibling for balancing purposes.
  182. Returns ``True`` if the theft was successful, or ``False`` if not.
  183. """
  184. if index + 1 < len(parent.children):
  185. right = parent.children[index + 1]
  186. if not right.is_minimal():
  187. right = parent.maybe_cow_child(index + 1)
  188. elt = parent.elts[index]
  189. parent.elts[index] = right.elts.pop(0)
  190. self.elts.append(elt)
  191. if not right.is_leaf:
  192. assert not self.is_leaf
  193. child = right.children.pop(0)
  194. self.children.append(child)
  195. return True
  196. return False
  197. def adopt(self, left: "_Node[KT, ET]", middle: ET, right: "_Node[KT, ET]") -> None:
  198. """Adopt left, middle, and right into our Node (which must not be maximal,
  199. and which must not be a leaf). In the case were we are not the new root,
  200. then the left child must already be in the Node."""
  201. assert not self.is_maximal()
  202. assert not self.is_leaf
  203. key = middle.key()
  204. i, equal = self.search_in_node(key)
  205. assert not equal
  206. self.elts.insert(i, middle)
  207. if len(self.children) == 0:
  208. # We are the new root
  209. self.children = [left, right]
  210. else:
  211. assert self.children[i] == left
  212. self.children.insert(i + 1, right)
  213. def merge(self, parent: "_Node[KT, ET]", index: int) -> None:
  214. """Merge this node's parent and its right sibling into this node."""
  215. right = parent.children.pop(index + 1)
  216. self.elts.append(parent.elts.pop(index))
  217. self.elts.extend(right.elts)
  218. if not self.is_leaf:
  219. self.children.extend(right.children)
  220. def minimum(self) -> ET:
  221. """The least element in this subtree."""
  222. if self.is_leaf:
  223. return self.elts[0]
  224. else:
  225. return self.children[0].minimum()
  226. def maximum(self) -> ET:
  227. """The greatest element in this subtree."""
  228. if self.is_leaf:
  229. return self.elts[-1]
  230. else:
  231. return self.children[-1].maximum()
  232. def balance(self, parent: "_Node[KT, ET]", index: int) -> None:
  233. """This Node is minimal, and we want to make it non-minimal so we can delete.
  234. We try to steal from our siblings, and if that doesn't work we will merge
  235. with one of them."""
  236. assert not parent.is_leaf
  237. if self.try_left_steal(parent, index):
  238. return
  239. if self.try_right_steal(parent, index):
  240. return
  241. # Stealing didn't work, so both siblings must be minimal.
  242. if index == 0:
  243. # We are the left-most node so merge with our right sibling.
  244. self.merge(parent, index)
  245. else:
  246. # Have our left sibling merge with us. This lets us only have "merge right"
  247. # code.
  248. left = parent.maybe_cow_child(index - 1)
  249. left.merge(parent, index - 1)
  250. def delete(
  251. self, key: KT, parent: Optional["_Node[KT, ET]"], exact: ET | None
  252. ) -> ET | None:
  253. """Delete an element matching *key* if it exists. If *exact* is not ``None``
  254. then it must be an exact match with that element. The Node must not be
  255. minimal unless it is the root."""
  256. assert parent is None or not self.is_minimal()
  257. i, equal = self.search_in_node(key)
  258. original_key = None
  259. if equal:
  260. # Note we use "is" here as we meant "exactly this object".
  261. if exact is not None and self.elts[i] is not exact:
  262. raise ValueError("exact delete did not match existing elt")
  263. if self.is_leaf:
  264. return self.elts.pop(i)
  265. # Note we need to ensure exact is None going forward as we've
  266. # already checked exactness and are about to change our target key
  267. # to the least successor.
  268. exact = None
  269. original_key = key
  270. least_successor = self.children[i + 1].minimum()
  271. key = least_successor.key()
  272. i = i + 1
  273. if self.is_leaf:
  274. # No match
  275. if exact is not None:
  276. raise ValueError("exact delete had no match")
  277. return None
  278. # recursively delete in the appropriate child
  279. child = self.maybe_cow_child(i)
  280. if child.is_minimal():
  281. child.balance(self, i)
  282. # Things may have moved.
  283. i, equal = self.search_in_node(key)
  284. assert not equal
  285. child = self.children[i]
  286. assert not child.is_minimal()
  287. elt = child.delete(key, self, exact)
  288. if original_key is not None:
  289. node, i = self._get_node(original_key)
  290. assert node is not None
  291. assert elt is not None
  292. oelt = node.elts[i]
  293. node.elts[i] = elt
  294. elt = oelt
  295. return elt
  296. def visit_in_order(self, visit: Callable[[ET], None]) -> None:
  297. """Call *visit* on all of the elements in order."""
  298. for i, elt in enumerate(self.elts):
  299. if not self.is_leaf:
  300. self.children[i].visit_in_order(visit)
  301. visit(elt)
  302. if not self.is_leaf:
  303. self.children[-1].visit_in_order(visit)
  304. def _visit_preorder_by_node(self, visit: Callable[["_Node[KT, ET]"], None]) -> None:
  305. """Visit nodes in preorder. This method is only used for testing."""
  306. visit(self)
  307. if not self.is_leaf:
  308. for child in self.children:
  309. child._visit_preorder_by_node(visit)
  310. def maybe_cow(self, creator: _Creator) -> Optional["_Node[KT, ET]"]:
  311. """Return a clone of this Node if it was not created by *creator*, or ``None``
  312. otherwise (i.e. copy for copy-on-write if we haven't already copied it)."""
  313. if self.creator is not creator:
  314. return self.clone(creator)
  315. else:
  316. return None
  317. def clone(self, creator: _Creator) -> "_Node[KT, ET]":
  318. """Make a shallow-copy duplicate of this node."""
  319. cloned = self.__class__(self.t, creator, self.is_leaf)
  320. cloned.elts.extend(self.elts)
  321. if not self.is_leaf:
  322. cloned.children.extend(self.children)
  323. return cloned
  324. def __str__(self): # pragma: no cover
  325. if not self.is_leaf:
  326. children = " " + " ".join([f"{id(c):x}" for c in self.children])
  327. else:
  328. children = ""
  329. return f"{id(self):x} {self.creator} {self.elts}{children}"
  330. class Cursor(Generic[KT, ET]):
  331. """A seekable cursor for a BTree.
  332. If you are going to use a cursor on a mutable BTree, you should use it
  333. in a ``with`` block so that any mutations of the BTree automatically park
  334. the cursor.
  335. """
  336. def __init__(self, btree: "BTree[KT, ET]"):
  337. self.btree = btree
  338. self.current_node: _Node | None = None
  339. # The current index is the element index within the current node, or
  340. # if there is no current node then it is 0 on the left boundary and 1
  341. # on the right boundary.
  342. self.current_index: int = 0
  343. self.recurse = False
  344. self.increasing = True
  345. self.parents: list[tuple[_Node, int]] = []
  346. self.parked = False
  347. self.parking_key: KT | None = None
  348. self.parking_key_read = False
  349. def _seek_least(self) -> None:
  350. # seek to the least value in the subtree beneath the current index of the
  351. # current node
  352. assert self.current_node is not None
  353. while not self.current_node.is_leaf:
  354. self.parents.append((self.current_node, self.current_index))
  355. self.current_node = self.current_node.children[self.current_index]
  356. assert self.current_node is not None
  357. self.current_index = 0
  358. def _seek_greatest(self) -> None:
  359. # seek to the greatest value in the subtree beneath the current index of the
  360. # current node
  361. assert self.current_node is not None
  362. while not self.current_node.is_leaf:
  363. self.parents.append((self.current_node, self.current_index))
  364. self.current_node = self.current_node.children[self.current_index]
  365. assert self.current_node is not None
  366. self.current_index = len(self.current_node.elts)
  367. def park(self):
  368. """Park the cursor.
  369. A cursor must be "parked" before mutating the BTree to avoid undefined behavior.
  370. Cursors created in a ``with`` block register with their BTree and will park
  371. automatically. Note that a parked cursor may not observe some changes made when
  372. it is parked; for example a cursor being iterated with next() will not see items
  373. inserted before its current position.
  374. """
  375. if not self.parked:
  376. self.parked = True
  377. def _maybe_unpark(self):
  378. if self.parked:
  379. if self.parking_key is not None:
  380. # remember our increasing hint, as seeking might change it
  381. increasing = self.increasing
  382. if self.parking_key_read:
  383. # We've already returned the parking key, so we want to be before it
  384. # if decreasing and after it if increasing.
  385. before = not self.increasing
  386. else:
  387. # We haven't returned the parking key, so we've parked right
  388. # after seeking or are on a boundary. Either way, the before
  389. # hint we want is the value of self.increasing.
  390. before = self.increasing
  391. self.seek(self.parking_key, before)
  392. self.increasing = increasing # might have been altered by seek()
  393. self.parked = False
  394. self.parking_key = None
  395. def prev(self) -> ET | None:
  396. """Get the previous element, or return None if on the left boundary."""
  397. self._maybe_unpark()
  398. self.parking_key = None
  399. if self.current_node is None:
  400. # on a boundary
  401. if self.current_index == 0:
  402. # left boundary, there is no prev
  403. return None
  404. else:
  405. assert self.current_index == 1
  406. # right boundary; seek to the actual boundary
  407. # so we can do a prev()
  408. self.current_node = self.btree.root
  409. self.current_index = len(self.btree.root.elts)
  410. self._seek_greatest()
  411. while True:
  412. if self.recurse:
  413. if not self.increasing:
  414. # We only want to recurse if we are continuing in the decreasing
  415. # direction.
  416. self._seek_greatest()
  417. self.recurse = False
  418. self.increasing = False
  419. self.current_index -= 1
  420. if self.current_index >= 0:
  421. elt = self.current_node.elts[self.current_index]
  422. if not self.current_node.is_leaf:
  423. self.recurse = True
  424. self.parking_key = elt.key()
  425. self.parking_key_read = True
  426. return elt
  427. else:
  428. if len(self.parents) > 0:
  429. self.current_node, self.current_index = self.parents.pop()
  430. else:
  431. self.current_node = None
  432. self.current_index = 0
  433. return None
  434. def next(self) -> ET | None:
  435. """Get the next element, or return None if on the right boundary."""
  436. self._maybe_unpark()
  437. self.parking_key = None
  438. if self.current_node is None:
  439. # on a boundary
  440. if self.current_index == 1:
  441. # right boundary, there is no next
  442. return None
  443. else:
  444. assert self.current_index == 0
  445. # left boundary; seek to the actual boundary
  446. # so we can do a next()
  447. self.current_node = self.btree.root
  448. self.current_index = 0
  449. self._seek_least()
  450. while True:
  451. if self.recurse:
  452. if self.increasing:
  453. # We only want to recurse if we are continuing in the increasing
  454. # direction.
  455. self._seek_least()
  456. self.recurse = False
  457. self.increasing = True
  458. if self.current_index < len(self.current_node.elts):
  459. elt = self.current_node.elts[self.current_index]
  460. self.current_index += 1
  461. if not self.current_node.is_leaf:
  462. self.recurse = True
  463. self.parking_key = elt.key()
  464. self.parking_key_read = True
  465. return elt
  466. else:
  467. if len(self.parents) > 0:
  468. self.current_node, self.current_index = self.parents.pop()
  469. else:
  470. self.current_node = None
  471. self.current_index = 1
  472. return None
  473. def _adjust_for_before(self, before: bool, i: int) -> None:
  474. if before:
  475. self.current_index = i
  476. else:
  477. self.current_index = i + 1
  478. def seek(self, key: KT, before: bool = True) -> None:
  479. """Seek to the specified key.
  480. If *before* is ``True`` (the default) then the cursor is positioned just
  481. before *key* if it exists, or before its least successor if it doesn't. A
  482. subsequent next() will retrieve this value. If *before* is ``False``, then
  483. the cursor is positioned just after *key* if it exists, or its greatest
  484. precessessor if it doesn't. A subsequent prev() will return this value.
  485. """
  486. self.current_node = self.btree.root
  487. assert self.current_node is not None
  488. self.recurse = False
  489. self.parents = []
  490. self.increasing = before
  491. self.parked = False
  492. self.parking_key = key
  493. self.parking_key_read = False
  494. while not self.current_node.is_leaf:
  495. i, equal = self.current_node.search_in_node(key)
  496. if equal:
  497. self._adjust_for_before(before, i)
  498. if before:
  499. self._seek_greatest()
  500. else:
  501. self._seek_least()
  502. return
  503. self.parents.append((self.current_node, i))
  504. self.current_node = self.current_node.children[i]
  505. assert self.current_node is not None
  506. i, equal = self.current_node.search_in_node(key)
  507. if equal:
  508. self._adjust_for_before(before, i)
  509. else:
  510. self.current_index = i
  511. def seek_first(self) -> None:
  512. """Seek to the left boundary (i.e. just before the least element).
  513. A subsequent next() will return the least element if the BTree isn't empty."""
  514. self.current_node = None
  515. self.current_index = 0
  516. self.recurse = False
  517. self.increasing = True
  518. self.parents = []
  519. self.parked = False
  520. self.parking_key = None
  521. def seek_last(self) -> None:
  522. """Seek to the right boundary (i.e. just after the greatest element).
  523. A subsequent prev() will return the greatest element if the BTree isn't empty.
  524. """
  525. self.current_node = None
  526. self.current_index = 1
  527. self.recurse = False
  528. self.increasing = False
  529. self.parents = []
  530. self.parked = False
  531. self.parking_key = None
  532. def __enter__(self):
  533. self.btree.register_cursor(self)
  534. return self
  535. def __exit__(self, exc_type, exc_value, traceback):
  536. self.btree.deregister_cursor(self)
  537. return False
  538. class Immutable(Exception):
  539. """The BTree is immutable."""
  540. class BTree(Generic[KT, ET]):
  541. """An in-memory BTree with copy-on-write and cursors."""
  542. def __init__(self, *, t: int = DEFAULT_T, original: Optional["BTree"] = None):
  543. """Create a BTree.
  544. If *original* is not ``None``, then the BTree is shallow-cloned from
  545. *original* using copy-on-write. Otherwise a new BTree with the specified
  546. *t* value is created.
  547. The BTree is not thread-safe.
  548. """
  549. # We don't use a reference to ourselves as a creator as we don't want
  550. # to prevent GC of old btrees.
  551. self.creator = _Creator()
  552. self._immutable = False
  553. self.t: int
  554. self.root: _Node
  555. self.size: int
  556. self.cursors: set[Cursor] = set()
  557. if original is not None:
  558. if not original._immutable:
  559. raise ValueError("original BTree is not immutable")
  560. self.t = original.t
  561. self.root = original.root
  562. self.size = original.size
  563. else:
  564. if t < 3:
  565. raise ValueError("t must be >= 3")
  566. self.t = t
  567. self.root = _Node(self.t, self.creator, True)
  568. self.size = 0
  569. def make_immutable(self):
  570. """Make the BTree immutable.
  571. Attempts to alter the BTree after making it immutable will raise an
  572. Immutable exception. This operation cannot be undone.
  573. """
  574. if not self._immutable:
  575. self._immutable = True
  576. def _check_mutable_and_park(self) -> None:
  577. if self._immutable:
  578. raise Immutable
  579. for cursor in self.cursors:
  580. cursor.park()
  581. # Note that we don't use insert() and delete() but rather insert_element() and
  582. # delete_key() so that BTreeDict can be a proper MutableMapping and supply the
  583. # rest of the standard mapping API.
  584. def insert_element(self, elt: ET, in_order: bool = False) -> ET | None:
  585. """Insert the element into the BTree.
  586. If *in_order* is ``True``, then extra work will be done to make left siblings
  587. full, which optimizes storage space when the the elements are inserted in-order
  588. or close to it.
  589. Returns the previously existing element at the element's key or ``None``.
  590. """
  591. self._check_mutable_and_park()
  592. cloned = self.root.maybe_cow(self.creator)
  593. if cloned:
  594. self.root = cloned
  595. if self.root.is_maximal():
  596. old_root = self.root
  597. self.root = _Node(self.t, self.creator, False)
  598. self.root.adopt(*old_root.split())
  599. oelt = self.root.insert_nonfull(elt, in_order)
  600. if oelt is None:
  601. # We did not replace, so something was added.
  602. self.size += 1
  603. return oelt
  604. def get_element(self, key: KT) -> ET | None:
  605. """Get the element matching *key* from the BTree, or return ``None`` if it
  606. does not exist.
  607. """
  608. return self.root.get(key)
  609. def _delete(self, key: KT, exact: ET | None) -> ET | None:
  610. self._check_mutable_and_park()
  611. cloned = self.root.maybe_cow(self.creator)
  612. if cloned:
  613. self.root = cloned
  614. elt = self.root.delete(key, None, exact)
  615. if elt is not None:
  616. # We deleted something
  617. self.size -= 1
  618. if len(self.root.elts) == 0:
  619. # The root is now empty. If there is a child, then collapse this root
  620. # level and make the child the new root.
  621. if not self.root.is_leaf:
  622. assert len(self.root.children) == 1
  623. self.root = self.root.children[0]
  624. return elt
  625. def delete_key(self, key: KT) -> ET | None:
  626. """Delete the element matching *key* from the BTree.
  627. Returns the matching element or ``None`` if it does not exist.
  628. """
  629. return self._delete(key, None)
  630. def delete_exact(self, element: ET) -> ET | None:
  631. """Delete *element* from the BTree.
  632. Returns the matching element or ``None`` if it was not in the BTree.
  633. """
  634. delt = self._delete(element.key(), element)
  635. assert delt is element
  636. return delt
  637. def __len__(self):
  638. return self.size
  639. def visit_in_order(self, visit: Callable[[ET], None]) -> None:
  640. """Call *visit*(element) on all elements in the tree in sorted order."""
  641. self.root.visit_in_order(visit)
  642. def _visit_preorder_by_node(self, visit: Callable[[_Node], None]) -> None:
  643. self.root._visit_preorder_by_node(visit)
  644. def cursor(self) -> Cursor[KT, ET]:
  645. """Create a cursor."""
  646. return Cursor(self)
  647. def register_cursor(self, cursor: Cursor) -> None:
  648. """Register a cursor for the automatic parking service."""
  649. self.cursors.add(cursor)
  650. def deregister_cursor(self, cursor: Cursor) -> None:
  651. """Deregister a cursor from the automatic parking service."""
  652. self.cursors.discard(cursor)
  653. def __copy__(self):
  654. return self.__class__(original=self)
  655. def __iter__(self):
  656. with self.cursor() as cursor:
  657. while True:
  658. elt = cursor.next()
  659. if elt is None:
  660. break
  661. yield elt.key()
  662. VT = TypeVar("VT") # the type of a value in a BTreeDict
  663. class KV(Element, Generic[KT, VT]):
  664. """The BTree element type used in a ``BTreeDict``."""
  665. def __init__(self, key: KT, value: VT):
  666. self._key = key
  667. self._value = value
  668. def key(self) -> KT:
  669. return self._key
  670. def value(self) -> VT:
  671. return self._value
  672. def __str__(self): # pragma: no cover
  673. return f"KV({self._key}, {self._value})"
  674. def __repr__(self): # pragma: no cover
  675. return f"KV({self._key}, {self._value})"
  676. class BTreeDict(Generic[KT, VT], BTree[KT, KV[KT, VT]], MutableMapping[KT, VT]):
  677. """A MutableMapping implemented with a BTree.
  678. Unlike a normal Python dict, the BTreeDict may be mutated while iterating.
  679. """
  680. def __init__(
  681. self,
  682. *,
  683. t: int = DEFAULT_T,
  684. original: BTree | None = None,
  685. in_order: bool = False,
  686. ):
  687. super().__init__(t=t, original=original)
  688. self.in_order = in_order
  689. def __getitem__(self, key: KT) -> VT:
  690. elt = self.get_element(key)
  691. if elt is None:
  692. raise KeyError
  693. else:
  694. return cast(KV, elt).value()
  695. def __setitem__(self, key: KT, value: VT) -> None:
  696. elt = KV(key, value)
  697. self.insert_element(elt, self.in_order)
  698. def __delitem__(self, key: KT) -> None:
  699. if self.delete_key(key) is None:
  700. raise KeyError
  701. class Member(Element, Generic[KT]):
  702. """The BTree element type used in a ``BTreeSet``."""
  703. def __init__(self, key: KT):
  704. self._key = key
  705. def key(self) -> KT:
  706. return self._key
  707. class BTreeSet(BTree, Generic[KT], MutableSet[KT]):
  708. """A MutableSet implemented with a BTree.
  709. Unlike a normal Python set, the BTreeSet may be mutated while iterating.
  710. """
  711. def __init__(
  712. self,
  713. *,
  714. t: int = DEFAULT_T,
  715. original: BTree | None = None,
  716. in_order: bool = False,
  717. ):
  718. super().__init__(t=t, original=original)
  719. self.in_order = in_order
  720. def __contains__(self, key: Any) -> bool:
  721. return self.get_element(key) is not None
  722. def add(self, value: KT) -> None:
  723. elt = Member(value)
  724. self.insert_element(elt, self.in_order)
  725. def discard(self, value: KT) -> None:
  726. self.delete_key(value)