this repo has no description
Refactor MST class into NodeWrangler
+39
-40
+39
-40
mst.py
+39
-40
mst.py
···
172
172
if cid is None, returns an empty MST node
173
173
"""
174
174
if cid is None:
175
-
return MSTNode.empty_root()
175
+
return self.put(MSTNode.empty_root())
176
176
177
177
return MSTNode.deserialise(self.bs.get(bytes(cid)))
178
178
···
182
182
return node # this is convenient
183
183
184
184
185
-
186
-
class MST:
185
+
class NodeWrangler:
186
+
"""
187
+
NodeWrangler is where core MST transformation ops are implemented, backed
188
+
by a NodeStore
189
+
"""
187
190
ns: NodeStore
188
-
root: CID
189
191
190
-
def __init__(self, ns: NodeStore, root: Optional[CID]=None) -> None:
192
+
def __init__(self, ns: NodeStore) -> None:
191
193
self.ns = ns
192
-
if root is None:
193
-
root = ns.put(MSTNode.empty_root()).cid
194
-
self.root = root
195
-
196
-
def put(self, key: str, val: CID):
197
-
self.root = self._put(key, val)
198
-
199
-
def delete(self, key: str):
200
-
self.root = self._delete(key)
201
194
202
-
def _put(self, key: str, val: CID) -> CID:
203
-
root = ns.get(self.root)
195
+
def put(self, root_cid: CID, key: str, val: CID) -> CID:
196
+
root = ns.get(root_cid)
204
197
if root.is_empty(): # special case for empty tree
205
198
return self._put_here(root, key, val)
206
199
return self._put_recursive(root, key, val, MSTNode.key_height(key), root.height)
200
+
201
+
def delete(self, root_cid: CID, key: str) -> CID:
202
+
root = ns.get(root_cid)
203
+
204
+
# Note: the seemingly redundant outer .get().cid is required to transform
205
+
# a None cid into the cid representing an empty node (we could maybe find a more elegant
206
+
# way of doing this...)
207
+
return self.ns.get(self._squash_top(self._delete_recursive(root, key, MSTNode.key_height(key), root.height))).cid
208
+
209
+
207
210
208
211
def _put_here(self, node: MSTNode, key: str, val: CID) -> CID:
209
212
i = node.gte_index(key)
···
278
281
return node_cid
279
282
return self._squash_top(node.subtrees[0])
280
283
281
-
def _delete(self, key: str) -> CID:
282
-
root = ns.get(self.root)
283
-
# XXX: handle empty tree result case
284
-
return self._squash_top(self._delete_recursive(root, key, MSTNode.key_height(key), root.height))
285
-
286
-
287
284
def _delete_recursive(self, node: MSTNode, key: str, key_height: int, tree_height: int) -> Optional[CID]:
288
285
if key_height > tree_height: # the key cannot possibly be in this tree, no change needed
289
-
return node.cid
286
+
return node._to_optional()
290
287
291
288
i = node.gte_index(key)
292
289
if key_height < tree_height: # the key must be deleted from a subtree
293
290
if node.subtrees[i] is None:
294
-
return node.cid # the key cannot be in this subtree, no change needed
291
+
return node._to_optional() # the key cannot be in this subtree, no change needed
295
292
return self.ns.put(MSTNode(
296
293
keys=node.keys,
297
294
vals=node.vals,
···
304
301
305
302
i = node.gte_index(key)
306
303
if i == len(node.keys) or node.keys[i] != key:
307
-
return node.cid # key already not present
304
+
return node._to_optional() # key already not present
308
305
309
306
assert(node.keys[i] == key) # sanity check (should always be true)
310
307
···
333
330
),
334
331
) + right.subtrees[1:]
335
332
))._to_optional()
336
-
337
-
def __repr__(self):
338
-
return self.pretty(self.root)
339
333
340
334
def pretty(self, node_cid: Optional[CID]) -> str:
341
335
if node_cid is None:
···
357
351
commit_obj = dag_cbor.decode(bs.get(bytes(bs.car_roots[0])))
358
352
mst_root: CID = commit_obj["data"]
359
353
ns = NodeStore(bs)
360
-
mst = MST(ns, mst_root)
354
+
mst = NodeWrangler(ns, mst_root)
361
355
print(mst)
362
356
else:
363
357
bs = MemoryBlockStore()
364
358
ns = NodeStore(bs)
365
-
mst = MST(ns)
366
-
print(mst)
367
-
mst.root = mst._put("hello", hash_to_cid(b"blah"))
368
-
print(mst)
369
-
mst.root = mst._put("foo", hash_to_cid(b"bar"))
370
-
print(mst)
371
-
mst.root = mst._put("bar", hash_to_cid(b"bat"))
372
-
print(mst)
373
-
mst.root = mst._delete("foo")
374
-
mst.root = mst._delete("hello")
375
-
print(mst)
359
+
mst = NodeWrangler(ns)
360
+
root = ns.get(None).cid
361
+
print(mst.pretty(root))
362
+
root = mst.put(root, "hello", hash_to_cid(b"blah"))
363
+
print(mst.pretty(root))
364
+
root = mst.put(root, "foo", hash_to_cid(b"bar"))
365
+
print(mst.pretty(root))
366
+
root = mst.put(root, "bar", hash_to_cid(b"bat"))
367
+
print(mst.pretty(root))
368
+
root = mst.delete(root, "foo")
369
+
root = mst.delete(root, "hello")
370
+
print(mst.pretty(root))
371
+
root = mst.delete(root, "bar")
372
+
print(mst.pretty(root))
373
+
root = mst.delete(root, "bar")
374
+
print(mst.pretty(root))