py_hamt

53def blake3_hashfn(input_bytes: bytes) -> bytes:
54    """
55    This is the default blake3 hash function used for the `HAMT`, with a 32 byte hash size.
56
57    """
58    # 32 bytes is the recommended byte size for blake3 and the default, but multihash forces us to explicitly specify
59    digest: bytes = b3.digest(input_bytes, size=32)
60    raw_bytes: bytes = b3.unwrap(digest)
61    return raw_bytes

289class HAMT:
290    """
291    An implementation of a Hash Array Mapped Trie for an arbitrary Content Addressed Storage (CAS) system, e.g. IPFS. This uses the IPLD data model.
292
293    Use this to store arbitrarily large key-value mappings in your CAS of choice.
294
295    For writing, this HAMT is async safe but NOT thread safe. Only write in an async event loop within the same thread.
296
297    When in read-only mode, the HAMT is both async and thread safe.
298
299    #### A note about memory management, read+write and read-only modes
300    The HAMT can be in either read+write mode or read-only mode. For either of these modes, the HAMT has some internal performance optimizations.
301
302    Note that in read+write, the real root node id IS NOT VALID. You should call `make_read_only()` to convert to read only mode and then read `root_node_id`.
303
304    These optimizations also trade off performance for memory use. Use `cache_size` to monitor the approximate memory usage. Be warned that for large key-value mapping sets this may take a bit to run. Use `cache_vacate` if you are over your memory limits.
305
306    #### IPFS HAMT Sample Code
307    ```python
308    kubo_cas = KuboCAS() # connects to a local kubo node with the default endpoints
309    hamt = await HAMT.build(cas=kubo_cas)
310    await hamt.set("foo", "bar")
311    assert (await hamt.get("foo")) == "bar"
312    await hamt.make_read_only()
313    cid = hamt.root_node_id # our root node CID
314    print(cid)
315    ```
316    """
317
318    def __init__(
319        self,
320        cas: ContentAddressedStore,
321        hash_fn: Callable[[bytes], bytes] = blake3_hashfn,
322        root_node_id: IPLDKind | None = None,
323        read_only: bool = False,
324        max_bucket_size: int = 4,
325        values_are_bytes: bool = False,
326    ):
327        """
328        Use `build` if you need to create a completely empty HAMT, as this requires some async operations with the CAS. For what each of the constructor input variables refer to, check the documentation with the matching names below.
329        """
330
331        self.cas: ContentAddressedStore = cas
332        """The backing storage system. py-hamt provides an implementation `KuboCAS` for IPFS."""
333
334        self.hash_fn: Callable[[bytes], bytes] = hash_fn
335        """
336        This is the hash function used to place a key-value within the HAMT.
337
338        To provide your own hash function, create a function that takes in arbitrarily long bytes and returns the hash bytes.
339
340        It's important to note that the resulting hash must must always be a multiple of 8 bits since python bytes object can only represent in segments of bytes, and thus 8 bits.
341
342        Theoretically your hash size must only be a minimum of 1 byte, and there can be less than or the same number of hash collisions as the bucket size. Any more and the HAMT will most likely throw errors.
343        """
344
345        self.lock: asyncio.Lock = asyncio.Lock()
346        """@private"""
347
348        self.values_are_bytes: bool = values_are_bytes
349        """Set this to true if you are only going to be storing python bytes objects into the hamt. This will improve performance by skipping a serialization step from IPLDKind.
350
351        This is theoretically safe to change in between operations, but this has not been verified in testing, so only do this at your own risk.
352        """
353
354        if max_bucket_size < 1:
355            raise ValueError("Bucket size maximum must be a positive integer")
356        self.max_bucket_size: int = max_bucket_size
357        """
358        This is only important for tuning performance when writing! For reading a HAMT that was written with a different max bucket size, this does not need to match and can be left unprovided.
359
360        This is an internal detail that has been exposed for performance tuning. The HAMT handles large key-value mapping sets even on a content addressed system by essentially sharding all the mappings across many smaller Nodes. The memory footprint of each of these Nodes footprint is a linear function of the maximum bucket size. Larger bucket sizes will result in larger Nodes, but more time taken to retrieve and decode these nodes from your backing CAS.
361
362        This must be a positive integer with a minimum of 1.
363        """
364
365        self.root_node_id: IPLDKind = root_node_id
366        """
367        This is type IPLDKind but the documentation generator pdoc mangles it a bit.
368
369        Read from this only when in read mode to get something valid!
370        """
371
372        self.read_only: bool = read_only
373        """Clients should NOT modify this.
374
375        This is here for checking whether the HAMT is in read only or read/write mode.
376
377        The distinction is made for performance and correctness reasons. In read only mode, the HAMT has an internal read cache that can speed up operations. In read/write mode, for reads the HAMT maintains strong consistency for reads by using async locks, and for writes the HAMT writes to an in memory buffer rather than performing (possibly) network calls to the underlying CAS.
378        """
379        self.node_store: NodeStore
380        """@private"""
381        if read_only:
382            self.node_store = ReadCacheStore(self)
383        else:
384            self.node_store = InMemoryTreeStore(self)
385
386    @classmethod
387    async def build(cls, *args: Any, **kwargs: Any) -> "HAMT":
388        """
389        Use this if you are initializing a completely empty HAMT! That means passing in None for the root_node_id. Method arguments are the exact same as `__init__`. If the root_node_id is not None, this will have no difference than creating a HAMT instance with __init__.
390
391        This separate async method is required since initializing an empty HAMT means sending some internal objects to the underlying CAS, which requires async operations. python does not allow for an async __init__, so this method is separately provided.
392        """
393        hamt = cls(*args, **kwargs)
394        if hamt.root_node_id is None:
395            hamt.root_node_id = await hamt.node_store.save(None, Node())
396        return hamt
397
398    # This is typically a massive blocking operation, you dont want to be running this concurrently with a bunch of other operations, so it's ok to have it not be async
399    async def make_read_only(self) -> None:
400        """
401        Makes the HAMT read only, which allows for more parallel read operations. The HAMT also needs to be in read only mode to get the real root node ID.
402
403        In read+write mode, the HAMT normally has to block separate get calls to enable strong consistency in case a set/delete operation falls in between.
404        """
405        async with self.lock:
406            inmemory_tree: InMemoryTreeStore = cast(InMemoryTreeStore, self.node_store)
407            await inmemory_tree.vacate()
408
409            self.read_only = True
410            self.node_store = ReadCacheStore(self)
411
412    async def enable_write(self) -> None:
413        """
414        Enable both reads and writes. This creates an internal structure for performance optimizations which will result in the root node ID no longer being valid, in order to read that at the end of your operations you must first use `make_read_only`.
415        """
416        async with self.lock:
417            # The read cache has no writes that need to be sent upstream so we can remove it without vacating
418            self.read_only = False
419            self.node_store = InMemoryTreeStore(self)
420
421    async def cache_size(self) -> int:
422        """
423        Returns the memory used by some internal performance optimization tools in bytes.
424
425        This is async concurrency safe, so call it whenever. This does mean it will block and wait for other writes to finish however.
426
427        Be warned that this may take a while to run for large HAMTs.
428
429        For more on memory management, see the `HAMT` class documentation.
430        """
431        if self.read_only:
432            return self.node_store.size()
433        async with self.lock:
434            return self.node_store.size()
435
436    async def cache_vacate(self) -> None:
437        """
438        Vacate and completely empty out the internal read/write cache.
439
440        Be warned that this may take a while if there have been a lot of write operations.
441
442        For more on memory management, see the `HAMT` class documentation.
443        """
444        if self.read_only:
445            await self.node_store.vacate()
446        else:
447            async with self.lock:
448                await self.node_store.vacate()
449
450    async def _reserialize_and_link(
451        self, node_stack: list[tuple[IPLDKind, Node]]
452    ) -> None:
453        """
454        This function starts from the node at the end of the list and reserializes so that each node holds valid new IDs after insertion into the store
455        Takes a stack of nodes, we represent a stack with a list where the first element is the root element and the last element is the top of the stack
456        Each element in the list is a tuple where the first element is the ID from the store and the second element is the Node in python
457        If a node ends up being empty, then it is deleted entirely, unless it is the root node
458        Modifies in place
459        """
460        # iterate in the reverse direction, this range goes from n-1 to 0, from the bottommost tree node to the root
461        for stack_index in range(len(node_stack) - 1, -1, -1):
462            old_id, node = node_stack[stack_index]
463
464            # If this node is empty, and it's not the root node, then we can delete it entirely from the list
465            is_root: bool = stack_index == 0
466            if node.is_empty() and not is_root:
467                # Unlink from the rest of the tree
468                _, prev_node = node_stack[stack_index - 1]
469                # When removing links, don't worry about two nodes having the same link since all nodes are guaranteed to be different by the removal of empty nodes after every single operation
470                for link_index in prev_node.iter_link_indices():
471                    link = prev_node.get_link(link_index)
472                    if link == old_id:
473                        # Delete the link by making it an empty bucket
474                        prev_node.data[link_index] = {}
475                        break
476
477                # Remove from our stack, continue reserializing up the tree
478                node_stack.pop(stack_index)
479                continue
480
481            # If not an empty node, just reserialize like normal and replace this one
482            new_store_id: IPLDKind = await self.node_store.save(old_id, node)
483            node_stack[stack_index] = (new_store_id, node)
484
485            # If this is not the last i.e. root node, we need to change the linking of the node prior in the list since we just reserialized
486            if not is_root:
487                _, prev_node = node_stack[stack_index - 1]
488                prev_node.replace_link(old_id, new_store_id)
489
490    # automatically skip encoding if the value provided is of the bytes variety
491    async def set(self, key: str, val: IPLDKind) -> None:
492        """Write a key-value mapping."""
493        if self.read_only:
494            raise Exception("Cannot call set on a read only HAMT")
495
496        data: bytes
497        if self.values_are_bytes:
498            data = cast(
499                bytes, val
500            )  # let users get an exception if they pass in a non bytes when they want to skip encoding
501        else:
502            data = dag_cbor.encode(val)
503
504        pointer: IPLDKind = await self.cas.save(data, codec="raw")
505        await self._set_pointer(key, pointer)
506
507    async def _set_pointer(self, key: str, val_ptr: IPLDKind) -> None:
508        async with self.lock:
509            node_stack: list[tuple[IPLDKind, Node]] = []
510            root_node: Node = await self.node_store.load(self.root_node_id)
511            node_stack.append((self.root_node_id, root_node))
512
513            # FIFO queue to keep track of all the KVs we need to insert
514            # This is needed if any buckets overflow and so we need to reinsert all those KVs
515            kvs_queue: list[tuple[str, IPLDKind]] = []
516            kvs_queue.append((key, val_ptr))
517
518            while len(kvs_queue) > 0:
519                _, top_node = node_stack[-1]
520                curr_key, curr_val_ptr = kvs_queue[0]
521
522                raw_hash: bytes = self.hash_fn(curr_key.encode())
523                map_key: int = extract_bits(raw_hash, len(node_stack) - 1, 8)
524
525                item = top_node.data[map_key]
526                if isinstance(item, list):
527                    next_node_id: IPLDKind = item[0]
528                    next_node: Node = await self.node_store.load(next_node_id)
529                    node_stack.append((next_node_id, next_node))
530                elif isinstance(item, dict):
531                    bucket: dict[str, IPLDKind] = item
532
533                    # If this bucket already has this same key, or has space, just rewrite the value and then go work on the others in the queue
534                    if curr_key in bucket or len(bucket) < self.max_bucket_size:
535                        bucket[curr_key] = curr_val_ptr
536                        kvs_queue.pop(0)
537                        continue
538
539                    # The current key is not in the bucket and the bucket is too full, so empty KVs from the bucket and restart insertion
540                    for k in bucket:
541                        v_ptr = bucket[k]
542                        kvs_queue.append((k, v_ptr))
543
544                    # Create a new link to a new node so that we can reflow these KVs into a new subtree
545                    new_node = Node()
546                    new_node_id: IPLDKind = await self.node_store.save(None, new_node)
547                    link: list[IPLDKind] = [new_node_id]
548                    top_node.data[map_key] = link
549
550            # Finally, reserialize and fix all links, deleting empty nodes as needed
551            await self._reserialize_and_link(node_stack)
552            self.root_node_id = node_stack[0][0]
553
554    async def delete(self, key: str) -> None:
555        """Delete a key-value mapping."""
556
557        # Also deletes the pointer at the same time so this doesn't have a _delete_pointer duo
558        if self.read_only:
559            raise Exception("Cannot call delete on a read only HAMT")
560
561        async with self.lock:
562            raw_hash: bytes = self.hash_fn(key.encode())
563
564            node_stack: list[tuple[IPLDKind, Node]] = []
565            root_node: Node = await self.node_store.load(self.root_node_id)
566            node_stack.append((self.root_node_id, root_node))
567
568            created_change: bool = False
569            while True:
570                _, top_node = node_stack[-1]
571                map_key: int = extract_bits(raw_hash, len(node_stack) - 1, 8)
572
573                item = top_node.data[map_key]
574                if isinstance(item, dict):
575                    bucket = item
576                    if key in bucket:
577                        del bucket[key]
578                        created_change = True
579                    # Break out since whether or not the key is in the bucket, it should have been here so either now reserialize or raise a KeyError
580                    break
581                elif isinstance(item, list):
582                    link: IPLDKind = item[0]
583                    next_node: Node = await self.node_store.load(link)
584                    node_stack.append((link, next_node))
585
586            # Finally, reserialize and fix all links, deleting empty nodes as needed
587            if created_change:
588                await self._reserialize_and_link(node_stack)
589                self.root_node_id = node_stack[0][0]
590            else:
591                # If we didn't make a change, then this key must not exist within the HAMT
592                raise KeyError
593
594    async def get(
595        self,
596        key: str,
597        offset: Optional[int] = None,
598        length: Optional[int] = None,
599        suffix: Optional[int] = None,
600    ) -> IPLDKind:
601        """Get a value."""
602        pointer: IPLDKind = await self.get_pointer(key)
603        data: bytes = await self.cas.load(
604            pointer, offset=offset, length=length, suffix=suffix
605        )
606        if self.values_are_bytes:
607            return data
608        else:
609            return dag_cbor.decode(data)
610
611    async def get_pointer(self, key: str) -> IPLDKind:
612        """
613        Get a store ID that points to the value for this key.
614
615        This is useful for some applications that want to implement a read cache. Due to the restrictions of `ContentAddressedStore` on IDs, pointers are regarded as immutable by python so they can be easily used as IDs for read caches. This is utilized in `ZarrHAMTStore` for example.
616        """
617        # If read only, no need to acquire a lock
618        pointer: IPLDKind
619        if self.read_only:
620            pointer = await self._get_pointer(key)
621        else:
622            async with self.lock:
623                pointer = await self._get_pointer(key)
624
625        return pointer
626
627    # Callers MUST handle acquiring a lock
628    async def _get_pointer(self, key: str) -> IPLDKind:
629        raw_hash: bytes = self.hash_fn(key.encode())
630
631        current_id: IPLDKind = self.root_node_id
632        current_depth: int = 0
633
634        # Don't check if result is none but use a boolean to indicate finding something, this is because None is a possible value of IPLDKind
635        result_ptr: IPLDKind = None
636        found_a_result: bool = False
637        while True:
638            top_id: IPLDKind = current_id
639            top_node: Node = await self.node_store.load(top_id)
640            map_key: int = extract_bits(raw_hash, current_depth, 8)
641
642            # Check if this key is in one of the buckets
643            item = top_node.data[map_key]
644            if isinstance(item, dict):
645                bucket = item
646                if key in bucket:
647                    result_ptr = bucket[key]
648                    found_a_result = True
649                    break
650
651            if isinstance(item, list):
652                link: IPLDKind = item[0]
653                current_id = link
654                current_depth += 1
655                continue
656
657            # Nowhere left to go, stop walking down the tree
658            break
659
660        if not found_a_result:
661            raise KeyError
662
663        return result_ptr
664
665    # Callers MUST handle locking or not on their own
666    async def _iter_nodes(self) -> AsyncIterator[tuple[IPLDKind, Node]]:
667        node_id_stack: list[IPLDKind] = [self.root_node_id]
668        while len(node_id_stack) > 0:
669            top_id: IPLDKind = node_id_stack.pop()
670            node: Node = await self.node_store.load(top_id)
671            yield (top_id, node)
672            node_id_stack.extend(list(node.iter_links()))
673
674    async def keys(self) -> AsyncIterator[str]:
675        """
676        AsyncIterator returning all keys in the HAMT.
677
678        If the HAMT is write enabled, to maintain strong consistency this will obtain an async lock and not allow any other operations to proceed.
679
680        When the HAMT is in read only mode however, this can be run concurrently with get operations.
681        """
682        if self.read_only:
683            async for k in self._keys_no_locking():
684                yield k
685        else:
686            async with self.lock:
687                async for k in self._keys_no_locking():
688                    yield k
689
690    async def _keys_no_locking(self) -> AsyncIterator[str]:
691        async for _, node in self._iter_nodes():
692            for bucket in node.iter_buckets():
693                for key in bucket:
694                    yield key
695
696    async def len(self) -> int:
697        """
698        Return the number of key value mappings in this HAMT.
699
700        When the HAMT is write enabled, to maintain strong consistency it will acquire a lock and thus not allow any other operations to proceed until the length is fully done being calculated. If read only, then this can be run concurrently with other operations.
701        """
702        count: int = 0
703        async for _ in self.keys():
704            count += 1
705
706        return count

14class ContentAddressedStore(ABC):
15    """
16    Abstract class that represents a content addressed storage that the `HAMT` can use for keeping data.
17
18    Note that the return type of save and input to load is really type `IPLDKind`, but the documentation generator pdoc mangles it unfortunately.
19
20    #### A note on the IPLDKind return types
21    Save and load return the type IPLDKind and not just a CID. As long as python regards the underlying type as immutable it can be used, allowing for more flexibility. There are two exceptions:
22    1. No lists or dicts, since python does not classify these as immutable.
23    2. No `None` values since this is used in HAMT's `__init__` to indicate that an empty HAMT needs to be initialized.
24    """
25
26    CodecInput = Literal["raw", "dag-cbor"]
27
28    @abstractmethod
29    async def save(self, data: bytes, codec: CodecInput) -> IPLDKind:
30        """Save data to a storage mechanism, and return an ID for the data in the IPLDKind type.
31
32        `codec` will be set to "dag-cbor" if this data should be marked as special linked data a la IPLD data model.
33        """
34
35    @abstractmethod
36    async def load(
37        self,
38        id: IPLDKind,
39        offset: Optional[int] = None,
40        length: Optional[int] = None,
41        suffix: Optional[int] = None,
42    ) -> bytes:
43        """Retrieve data."""
44
45    async def pin_cid(self, id: IPLDKind, target_rpc: str) -> None:
46        """Pin a CID in the storage."""
47        pass  # pragma: no cover
48
49    async def unpin_cid(self, id: IPLDKind, target_rpc: str) -> None:
50        """Unpin a CID in the storage."""
51        pass  # pragma: no cover
52
53    async def pin_update(
54        self, old_id: IPLDKind, new_id: IPLDKind, target_rpc: str
55    ) -> None:
56        """Update the pinned CID in the storage."""
57        pass  # pragma: no cover
58
59    async def pin_ls(self, target_rpc: str) -> list[Dict[str, Any]]:
60        """List all pinned CIDs in the storage."""
61        return []  # pragma: no cover

 64class InMemoryCAS(ContentAddressedStore):
 65    """Used mostly for faster testing, this is why this is not exported. It hashes all inputs and uses that as a key to an in-memory python dict, mimicking a content addressed storage system. The hash bytes are the ID that `save` returns and `load` takes in."""
 66
 67    store: dict[bytes, bytes]
 68    hash_alg: Multihash
 69
 70    def __init__(self):
 71        self.store = dict()
 72        self.hash_alg = multihash.get("blake3")
 73
 74    async def save(self, data: bytes, codec: ContentAddressedStore.CodecInput) -> bytes:
 75        hash: bytes = self.hash_alg.digest(data, size=32)
 76        self.store[hash] = data
 77        return hash
 78
 79    async def load(
 80        self,
 81        id: IPLDKind,
 82        offset: Optional[int] = None,
 83        length: Optional[int] = None,
 84        suffix: Optional[int] = None,
 85    ) -> bytes:
 86        """
 87        `ContentAddressedStore` allows any IPLD scalar key.  For the in-memory
 88        backend we *require* a `bytes` hash; anything else is rejected at run
 89        time. In OO type-checking, a subclass may widen (make more general) argument types,
 90        but it must never narrow them; otherwise callers that expect the base-class contract can break.
 91        Mypy enforces this contra-variance rule and emits the "violates Liskov substitution principle" error.
 92        This is why we use `cast` here, to tell mypy that we know what we are doing.
 93        h/t https://stackoverflow.com/questions/75209249/overriding-a-method-mypy-throws-an-incompatible-with-super-type-error-when-ch
 94        """
 95        key = cast(bytes, id)
 96        if not isinstance(key, (bytes, bytearray)):  # defensive guard
 97            raise TypeError(
 98                f"InMemoryCAS only supports byte‐hash keys; got {type(id).__name__}"
 99            )
100        data: bytes
101        try:
102            data = self.store[key]
103        except KeyError as exc:
104            raise KeyError("Object not found in in-memory store") from exc
105
106        if offset is not None:
107            start = offset
108            if length is not None:
109                end = start + length
110                return data[start:end]
111            else:
112                return data[start:]
113        elif suffix is not None:  # If only length is given, assume start from 0
114            return data[-suffix:]
115        else:  # Full load
116            return data