allocation_bitmap.py

from mmfs_errors import BoundsError


class mmfsBitmapAllocator:
    """
    This is my implementation of free space tracking using a bitmap.

    The idea is 1 bit = 1 block, so each byte in the bitmap covers 8
    blocks. Bit set means the block is used, bit clear means its free.

    To get at any block we split the block number into which byte it
    lives in and which bit inside that byte:

        block 12 -> divmod(12, 8) -> byte 1, bit 4

        b[1] |= (1 << 4)     mark used
        b[1] &= ~(1 << 4)    mark free

    Everything before the data region (header, bitmap, catalog, MAT,
    MIT) gets marked protected as we format. After that, anything
    trying to flip one of those raises instead of quietly overwritign
    our metadata.
    """

    BLOCKS_PER_BYTE = 8

    def __init__(self, disk, protected=None, bitmap_start_block=1, bitmap_blocks=1, block_size=65536):
        self.DISK = disk
        self.BLOCK_SIZE = block_size
        self.PROTECTED_BLOCKS = protected if protected is not None else []
        self.BITMAP_START_BLOCK = bitmap_start_block
        self.TOTAL_BITMAP_BLOCKS = bitmap_blocks
        self.TOTAL_BITMAP_SIZE = bitmap_blocks * block_size
        self.BITMAP_OFFSET = bitmap_start_block * block_size

        # The bitmap always rounds up to whole
        # blocks, so it can track more blocks than the disk actually
        # has. Whoever is searching has to pass an end_block or we will
        # happily hand out a block past the end of the image.
        self.TOTAL_ALLOWED_TRACKED_BLOCKS = self.TOTAL_BITMAP_SIZE * self.BLOCKS_PER_BYTE

    def _calc_offset_positions(self, block, initialize):
        """
        Turns a block number into (byte_index, bit_index) so we know
        where in the bitmap to go. Also where the safety checks live,
        since every read and write goes through here.
        """
        if (block in self.PROTECTED_BLOCKS) and not initialize:
            raise BoundsError("[!!! CRITICAL !!!] Bitmap tried freeing/allocating a protected block")

        if block < 0 or block >= self.TOTAL_ALLOWED_TRACKED_BLOCKS:
            raise ValueError(
                f"block {block} out of range (0..{self.TOTAL_ALLOWED_TRACKED_BLOCKS - 1})"
            )

        return divmod(block, self.BLOCKS_PER_BYTE)

    def _read_byte(self, byte_index):
        # Every operation in here needs "seek to a byte and read it",
        # so its pulled out rather than repeated four times.
        
        self.DISK.seek(self.BITMAP_OFFSET + byte_index)
       
        return int.from_bytes(self.DISK.read(1), byteorder="big")

    def bit_flip(self, free, block=-1, initialize=False, blocks=None):
        """
        Mark blocks as used or free by updating their bitmap entries in place.

        Args:
            free: If True, free the selected blocks. Otherwise, mark them as
                allocated.

            block: A single block to update. Defaults to -1 so an invalid call
                fails instead of modifying block 0.

            initialize: Used only during formatting. Skips the protected-block
                check and marks each updated block as protected.

            blocks: A list of blocks to update. Takes priority over `block` when
                both are provided.

        The single-block and multiple-block cases share one implementation
        because the only difference is which collection is iterated.
        """
        targets = blocks if blocks else [block] # we can free multiple blocks or one block
                                                # casting block into a list eliminates
                                                # some duplicate code

        for b in targets:
            byte_index, bit_index = self._calc_offset_positions(block=b, initialize=initialize)

            if initialize:
                # If we are initializing a disk aka formatting it
                # we assume those blocks are protected and cannot
                # be freed, therefore I add them to a list on
                # initialization/formatting
                self.PROTECTED_BLOCKS.append(b)

            val = self._read_byte(byte_index)

            if free:
                val &= ~(1 << bit_index)
            else:
                val |= (1 << bit_index)

            # Reading moved us one byte forward, so step back before we
            # write or we would clobber the next byte instead.
            #
            # frustrated programmer note:
            #   this took me very long and a lot of fake disk images
            #   to understand how to seek backwards, it seems simple
            #   and it is, but when you are jumping all over a disk
            #   it can be a hastle figure out where the heck your
            #   cursor actually is and landed. it also took me
            #   a minute to understand if the cursor location
            #   would write at the current byte from the byte offset
            #   (think insert on a keybord) or if it would write it
            #   to the next byte 

            self.DISK.seek(-1, 1)
            self.DISK.write(val.to_bytes(1, byteorder="big"))

    def is_block_used(self, block):
        byte_index, bit_index = self._calc_offset_positions(block=block, initialize=False)
        return bool(self._read_byte(byte_index) & (1 << bit_index))

    def find_free_block(self):
        """
        Find the first free block in the bitmap.

        Will return the first available block number, or None if no free blocks remain.

        The current bitmap byte is cached so one read can check eight blocks.
        Without caching, each block check would require a separate seek and read.
        """
        current_byte = None
        current_byte_index = None

        for block in range(self.TOTAL_ALLOWED_TRACKED_BLOCKS):
            byte_index, bit_index = self._calc_offset_positions(block=block, initialize=False)

            if byte_index != current_byte_index:
                current_byte = self._read_byte(byte_index)
                current_byte_index = byte_index

            if not (current_byte & (1 << bit_index)):
                return block

        return None

    def find_free_extent(self, block_count, start_block=0, end_block=None):
        """
        Find a contiguous range of free blocks.

        Uses the same bitmap scan as find_free_block(), but searches for
        block_count consecutive free blocks instead of a single block.

        The consecutive block count resets whenever an allocated block is found.
        Returns the starting block of the range, or None if no large enough range
        is available.
        """
        if block_count <= 0:
            raise BoundsError("[!!! MEDIUM !!!] Bitmap allocation needs positive block count")

        if start_block < 0:
            raise BoundsError("[!!! MEDIUM !!!] Bitmap allocation needs positive start block")

        # Never search past what the bitmap can track, even if the
        # caller asks us to.
        limit = (
            self.TOTAL_ALLOWED_TRACKED_BLOCKS
            if end_block is None
            else min(end_block, self.TOTAL_ALLOWED_TRACKED_BLOCKS)
        )

        run_start = None
        run_length = 0

        current_byte = None
        current_byte_index = None

        for block in range(start_block, limit):
            byte_index, bit_index = self._calc_offset_positions(block=block, initialize=False)

            if byte_index != current_byte_index:
                current_byte = self._read_byte(byte_index)
                current_byte_index = byte_index

            if current_byte & (1 << bit_index):
                # Used block, so whatever run we had going is dead.
                run_start = None
                run_length = 0
                continue

            if run_start is None:
                run_start = block

            run_length += 1

            if run_length == block_count:
                return run_start

        return None