|
| 1 | +import math |
| 2 | + |
| 3 | + |
| 4 | +def extract_bits(hash_obj: bytes, depth: int, nbits: int) -> int: |
| 5 | + """Extract `nbits` bits from `hash_obj`, beginning at position `depth * nbits`, |
| 6 | + and convert them into an unsigned integer value. |
| 7 | +
|
| 8 | + Args: |
| 9 | + hash_obj (bytes): binary hash to extract bit sequence from |
| 10 | + depth (int): depth of the node containing the hash |
| 11 | + nbits (int): bit width of hash |
| 12 | +
|
| 13 | + Returns: |
| 14 | + int: An unsigned integer version of the bit sequence |
| 15 | + """ |
| 16 | + start = depth * nbits |
| 17 | + start_offset = start % 8 |
| 18 | + |
| 19 | + byte_count = math.ceil((start_offset + nbits) / 8) |
| 20 | + byte_start = start >> 3 |
| 21 | + end_offset = byte_count * 8 - nbits - start_offset |
| 22 | + |
| 23 | + result = 0 |
| 24 | + |
| 25 | + for i in range(byte_count): |
| 26 | + local = hash_obj[byte_start + i] |
| 27 | + shift = 0 |
| 28 | + local_bit_length = 8 |
| 29 | + |
| 30 | + if i == 0: |
| 31 | + local_bit_length -= start_offset |
| 32 | + |
| 33 | + if i == byte_count - 1: |
| 34 | + local_bit_length -= start_offset |
| 35 | + shift = end_offset |
| 36 | + local >>= shift |
| 37 | + |
| 38 | + if local_bit_length < 8: |
| 39 | + m = (1 << local_bit_length) - 1 |
| 40 | + local &= m |
| 41 | + |
| 42 | + if shift < 8: |
| 43 | + result = result << (8 - shift) |
| 44 | + result |= local |
| 45 | + |
| 46 | + return result |
| 47 | + |
| 48 | + |
| 49 | +def set_bit(bitmap: bytes, position: int, to_set: bool) -> bytes: |
| 50 | + """set the `position` bit in the given `bitmap` to be `to_set` (truthy=1, falsey=0) |
| 51 | +
|
| 52 | + Args: |
| 53 | + bitmap (bytes): bitmap to modify |
| 54 | + position (int): location in the bitmap to modify |
| 55 | + to_set (bool): whether to set true or false |
| 56 | +
|
| 57 | + Returns: |
| 58 | + bytes: Modified bitmap |
| 59 | + """ |
| 60 | + has = bitmap_has(bitmap, position) |
| 61 | + byte = math.floor(position / 8) |
| 62 | + offset = position % 8 |
| 63 | + # if we assume that `bitmap` is already the opposite of `set`, we could skip this check |
| 64 | + if (to_set and not has) or (not to_set and has): |
| 65 | + new_bit_map = bytearray(bitmap) |
| 66 | + b = bitmap[byte] |
| 67 | + if to_set: |
| 68 | + b |= 1 << offset |
| 69 | + else: |
| 70 | + b ^= 1 << offset |
| 71 | + |
| 72 | + # since bytes are immutable, we need to change bytes to bytearrays |
| 73 | + new_bit_map[byte] = b |
| 74 | + return bytes(new_bit_map) |
| 75 | + return bitmap |
| 76 | + |
| 77 | + |
| 78 | +def bitmap_has( |
| 79 | + bitmap: bytes, |
| 80 | + position: int, |
| 81 | +) -> bool: |
| 82 | + """check whether `bitmap` has a `1` at the given `position` bit. |
| 83 | +
|
| 84 | + Args: |
| 85 | + bitmap (bytes): bytes to check |
| 86 | + position (int): Position of bit to read |
| 87 | +
|
| 88 | + Returns: |
| 89 | + bool: whether the `bitmap` has a 1 value at the `position` bit |
| 90 | + """ |
| 91 | + byte = math.floor(position / 8) |
| 92 | + offset = position % 8 |
| 93 | + return ((bitmap[byte] >> offset) & 1) == 1 |
| 94 | + |
| 95 | + |
| 96 | +def rank(bitmap: bytes, position: int) -> int: |
| 97 | + """count how many `1` bits are in `bitmap` up until `position` |
| 98 | + tells us where in the compacted element array an element should live |
| 99 | + Args: |
| 100 | + bitmap (bytes): bitmap to count truthy bits on |
| 101 | + position (int): where to stop counting |
| 102 | +
|
| 103 | + Returns: |
| 104 | + int: how many bits are `1` in `bitmap` |
| 105 | + """ |
| 106 | + t = 0 |
| 107 | + for i in range(position): |
| 108 | + if bitmap_has(bitmap, i): |
| 109 | + t += 1 |
| 110 | + return t |
0 commit comments