page_table.py

import math
from dataclasses import dataclass
from typing import Iterator, Optional


@dataclass
class Page:
    """Page"""

    address: int  # The virtual address of the page
    # Physical frame number (None if it's not in RAM)
    frame: Optional[int] = None

    modified: bool = False  # Dirty bit
    # For clock algorithm (None if not using clock)
    referenced: Optional[bool] = None

    def evict(self) -> None:
        """Reset the page when it's evicted from RAM"""
        self.frame = None
        self.modified = False
        if self.referenced is not None:  # Only reset for clock
            self.referenced = False


@dataclass
class PageTableEntry:
    """Page table entry"""

    page: Optional[Page] = None
    valid: bool = False  # Wheher the page table entry points to a valid page
    # Kind of redundant since I can just check if page is None but it's at least a little clearer this way


class PageTable:
    """Two-level page table

    Sources:
        https://www.geeksforgeeks.org/two-level-paging-and-multi-level-paging-in-os/
        https://docs.python.org/3/library/platform.html#platform.architecture
    """

    # Constants
    PAGE_SIZE: int = 2048  # 2KB
    ROOT_ENTRIES: int = 2048  # 2048 entries
    ENTRY_SIZE: int = 4  # 4 bytes (32 bit)

    # Derived constants
    PAGE_OFFSET_BITS: int = math.floor(math.log2(PAGE_SIZE))  # 11 bits
    ROOT_BITS: int = math.floor(math.log2(ROOT_ENTRIES))  # 11 bits
    LEAF_BITS: int = 32 - PAGE_OFFSET_BITS - ROOT_BITS  # 10 bits
    LEAF_ENTRIES: int = 1 << LEAF_BITS  # 1024 entries

    def __init__(self, num_frames: int):
        """Initialize the page table with pre-allocated root array"""
        self.num_frames: int = num_frames

        # Each root entry points to a leaf array (or empty list if not allocated)
        self.root: list[list[PageTableEntry]] = [None] * self.ROOT_ENTRIES
        self.num_leaves: int = 0

        # Creates a set of frames, which automatically does deduplication and sorting
        # We'll use it to keep track of which frames are free
        # Could have also used a dictionary, but we don't really need to store values, just keys
        self.free_frames: set[int] = {i for i in range(num_frames)}

    @property
    def total_size(self) -> int:
        """Calculate total size of page table in bytes

        Returns:
            int: Total size of page table in bytes
        """
        root_size = self.ROOT_ENTRIES * self.ENTRY_SIZE  # 2048 * 4
        leaf_size = self.LEAF_ENTRIES * self.ENTRY_SIZE  # 1024 * 4
        total_leaf_size = self.num_leaves * leaf_size

        return root_size + total_leaf_size

    def allocate_frame(self, page: Page) -> int:
        """Allocate a free frame to a page

        Time complexity: O(1)

        Args:
            page (Page): Page to allocate frame to

        Returns:
            int: Frame number if allocation successful

        Raises:
            MemoryError: If no free frames are available
        """
        if len(self.free_frames) == 0:  # No free frames
            raise MemoryError("No free frames available")

        page.frame = self.free_frames.pop()  # Gets the frame with the lowest number
        return page.frame

    def add_page(self, address: int) -> Page:
        """Add a page to the page table

        Time complexity: O(1) average case, O(LEAF_ENTRIES) worst case

        Args:
            address: Virtual address
        """
        entry = self.__get_pte(address)
        entry.valid = True
        entry.page = Page(address)

        return entry.page

    def del_page(self, page: Page) -> None:
        """Free a page's frame, invalidate its page table entry, and delete it

        Time complexity: O(1) average case, O(LEAF_ENTRIES) worst case

        Args:
            page (Page): Page to remove
        """
        if page.frame is not None:  # Put the frame back in pool of free frames
            self.free_frames.add(page.frame)

        entry = self.__get_pte(page.address)
        if entry is None or entry.valid is False:  # The leaf entry is already invalid
            return

        entry.valid = False  # Invalidate the leaf entry
        page.evict()  # Reset the page to its initial state

    def get_page(self, address: int) -> Optional[Page]:
        """Get a page from the page table

        Time complexity: O(1) average case, O(LEAF_ENTRIES) worst case

        Args:
            address: Virtual address

        Returns:
            The page at the given address, or None if it doesn't exist
        """
        entry = self.__get_pte(address)

        if entry is None or entry.valid is False:  # Check if the leaf entry is valid
            return None

        return entry.page

    def __get_pte(self, address: int) -> Optional[PageTableEntry]:
        """Retrieve the page table entry for a given address

        Time complexity: O(1) average case, O(LEAF_ENTRIES) worst case

        Args:
            address: Virtual address

        Returns:
            Optional[PageTableEntry]: The page table entry at the given address (if it exists)
        """
        # [root index (31:21)][leaf index (20:11)][page offset (10:0)]

        # By sifting right by 21 bits, we get the root index (31:21)
        # Masking with 2047 (11111111111) ensures we only get the 11 bits we need
        root_index = (address >> (self.PAGE_OFFSET_BITS + self.LEAF_BITS)) & (
            self.ROOT_ENTRIES - 1
        )  # address >> 21 & 2047

        # By sifting right by 11 bits, we get the leaf index (20:11)
        # Masking with 1023 (1111111111) ensures we only get the 10 bits we need
        leaf_index = (address >> self.PAGE_OFFSET_BITS) & (
            self.LEAF_ENTRIES - 1
        )  # address >> 11 & 1023

        if self.root[root_index] is None:  # The root entry doesn't exist yet
            # Allocate the array of leaf entries
            self.root[root_index] = [
                PageTableEntry(valid=False) for _ in range(self.LEAF_ENTRIES)
            ]
            self.num_leaves += 1

        return self.root[root_index][leaf_index]

    def __iter__(self) -> Iterator[tuple[str, int]]:
        """Makes the page table iterable so I can put it into a dict with `dict(page_table)`

        Yields:
            Iterator[tuple[str, int]]: Iterable of the page table stats
        """
        yield "num_leaves", self.num_leaves
        yield "total_size", self.total_size

    def __str__(self) -> str:
        """Makes PageTable passable to `str()`. So I can do `print(page_table)`

        Returns:
            str: The number of leaves and total size of the page table in bytes
        """
        return (
            f"Number of page table leaves:\t{self.num_leaves}\n"
            f"Total size of page table:\t{self.total_size}"
        )