memory_allocator.py

from typing import List


class MemoryAllocator:
    # Constant defined for header size
    HEADER_SIZE = 3

    def __init__(self, total_size: int):
        """Initialize the memory data array with given size."""
        self.data: List = [None] * total_size

        # Header info
        self.data[0] = total_size
        self.data[1] = 0  # free chunk
        self.data[2] = -1  # prev chunk address

    def malloc(self, size: int):
        total_size = len(self.data)
        ptr = 0
        while ptr < total_size - MemoryAllocator.HEADER_SIZE:
            # Find a free chunk and can be allocated
            if self.data[ptr + 1] == 0 and self.data[ptr] - MemoryAllocator.HEADER_SIZE >= size:
                cur_chunk_size = self.data[ptr]
                updated_chunk_size = MemoryAllocator.HEADER_SIZE + size
                rem_chunk_size = cur_chunk_size - updated_chunk_size
                if rem_chunk_size <= MemoryAllocator.HEADER_SIZE:
                    updated_chunk_size += rem_chunk_size
                else:
                    new_ptr = ptr + MemoryAllocator.HEADER_SIZE + size
                    self.data[new_ptr] = rem_chunk_size
                    self.data[new_ptr + 1] = 0
                    self.data[new_ptr + 2] = ptr

                # Update current header info
                self.data[ptr] = updated_chunk_size
                self.data[ptr + 1] = 1

            # Check next chunk
            ptr += self.data[ptr]

        return ptr

    def free(self, ptr: int):
        # Check if current chunk is already free
        if self.data[ptr + 1] == 0:
            return
        
        prev_ptr, next_ptr = self.data[ptr + 2], ptr + self.data[ptr]

        # Free current chunk
        self.data[ptr + 1] = 0

        # The merged header ptr
        merged_head_ptr = ptr

        # Merge with prev free chunk if it exists
        if prev_ptr != -1 and self.data[prev_ptr + 1] == 0:
            merged_head_ptr = self._merge_chunks(prev_ptr, ptr)

        # Merge with next free chunk if it exists
        if self.data[next_ptr + 1] == 0:
            next_next_ptr = next_ptr + self.data[next_ptr]
            merged_head_ptr = self._merge_chunks(merged_head_ptr, next_ptr)
            if next_next_ptr < len(self.data):
                self.data[next_next_ptr + 2] = merged_head_ptr
        else:
            self.data[ptr + 2] = merged_head_ptr

    def _merge_chunks(self, ptr1: int, ptr2: int):
        """Merge two adjacent free chunks

        Returns ptr of the newly merged free chunk.
        """
        # Merge chunk size
        self.data[ptr1] += self.data[ptr2]

        # Reset header info for the second free chunk
        self.data[ptr2] = None
        self.data[ptr2 + 1] = None
        self.data[ptr2 + 2] = None

        return ptr1