Improve page table creation in python file

wally-pipelined/testgen/virtual_memory_util.py

@@ -27,9 +27,13 @@ PTE_W = 1 << 2
 PTE_R = 1 << 1
 PTE_V = 1 << 0
 
+PTE_PTR_MASK = ~(PTE_W | PTE_R | PTE_X)
 
 pgdir = []
 
+pages = {}
+
+
 testcase_num = 0
 signature_len = 2000
 signature = [0xff for _ in range(signature_len)]
@@ -41,6 +45,7 @@ class Architecture:
   def __init__(self, xlen):
     if (xlen == 32):
       self.PTESIZE = 4
+      self.PTE_BITS = 32
 
       self.VPN_BITS = 20
       self.VPN_SEGMENT_BITS = 10
@@ -50,6 +55,7 @@ class Architecture:
       self.LEVELS = 2
     elif (xlen == 64):
       self.PTESIZE = 8
+      self.PTE_BITS = 54
 
       self.VPN_BITS = 27
       self.VPN_SEGMENT_BITS = 9
@@ -62,9 +68,9 @@ class Architecture:
 
     self.PGSIZE = 2**12
     self.NPTENTRIES = self.PGSIZE // self.PTESIZE
-    self.PTE_BITS = 8 * self.PTESIZE
     self.OFFSET_BITS = 12
     self.FLAG_BITS = 8
+    self.VA_BITS = self.VPN_BITS + self.OFFSET_BITS
 
 class PageTableEntry:
   def __init__(self, ppn, flags, arch):
@@ -85,41 +91,127 @@ class PageTableEntry:
 
 class PageTable:
   """
-  Represents a single level of the page table, with  
+  Represents a single level of the page table, located at some physical page
+  number `ppn` with symbol `name`, using a specified architecture `arch`.
   """
-  def __init__(self, name, arch):
+  def __init__(self, name, ppn, arch):
     self.table = {}
     self.name = name
+    self.ppn = ppn
     self.arch = arch
 
-  def add_entry(self, vpn_segment, ppn_segment, flags, linked_table = None):
+    self.children = 0
+
+    pages[ppn] = self
+
+  def add_entry(self, vpn_segment, ppn, flags):
     if not (0 <= vpn_segment < 2**self.arch.VPN_SEGMENT_BITS):
       raise ValueError("Invalid virtual page segment number")
-    self.table[vpn_segment] = (PageTableEntry(ppn_segment, flags, self.arch), linked_table)
+    self.table[vpn_segment] = PageTableEntry(ppn, flags, self.arch)
 
   def add_mapping(self, va, pa, flags):
-    if not (0 <= va < 2**self.arch.VPN_BITS):
+    """
+    Maps a virtual address `va` to a physical address `pa` with given `flags`,
+    creating missing page table levels as needed.
+    """
+    if not (0 <= va < 2**self.arch.VA_BITS):
       raise ValueError("Invalid virtual page number")
+
+    vpn = virtual_to_vpn(va, self.arch)
+    ppn = pa >> self.arch.OFFSET_BITS
+    current_level = self
+
+    pathname = self.name
+
     for level in range(self.arch.LEVELS - 1, -1, -1):
-      
+      if level == 0:
-
+        current_level.add_entry(vpn[level], ppn, flags)
-    
+      elif vpn[level] in current_level.table:
+        current_level = pages[current_level.table[vpn[level]].ppn]
+        pathname += f"_{current_level.name}"
+      else:
+        next_level_ppn = next_ppn()
+        current_level.add_entry(vpn[level], next_level_ppn, flags & PTE_PTR_MASK)
+        pathname += f"_t{current_level.children}"
+        current_level.children += 1
+        pages[next_level_ppn] = PageTable(pathname, next_level_ppn, self.arch)
+        current_level = pages[next_level_ppn]
 
   def assembly(self):
+    # Sort the page table
     entries = list(sorted(self.table.items(), key=lambda item: item[0]))
     current_index = 0
+
+    # Align the table
     asm = f".balign {self.arch.PGSIZE}\n{self.name}:\n"
     for entry in entries:
-      vpn_index, (pte, _) = entry
+      vpn_index, pte = entry
       if current_index < vpn_index:
         asm += f"  .fill {vpn_index - current_index}, {self.arch.PTESIZE}, 0\n"
-      asm += f"  .4byte {str(pte)}\n"
+      asm += f"  .{self.arch.PTESIZE}byte {str(pte)}\n"
       current_index = vpn_index + 1
     if current_index < self.arch.NPTENTRIES:
       asm += f"  .fill {self.arch.NPTENTRIES - current_index}, {self.arch.PTESIZE}, 0\n"
     return asm
   
+  def __str__(self):
+    return self.assembly()
+
+  def __repr__(self):
+    return f"<table: {self.table}>"
+
+
 ##################################
 # functions
 ##################################
 
+def virtual_to_vpn(vaddr, arch):
+  if not (0 <= vaddr < 2**arch.VA_BITS):
+    raise ValueError("Invalid physical address")
+
+  page_number = [0 for _ in range(arch.LEVELS)]
+
+  vaddr = vaddr >> arch.OFFSET_BITS
+  mask = 2**arch.VPN_SEGMENT_BITS - 1
+  for level in range(arch.LEVELS):
+    page_number[level] = vaddr & mask
+    vaddr = vaddr >> arch.VPN_SEGMENT_BITS
+
+  return page_number
+
+INITIAL_PPN = 0x80002
+next_free_ppn = INITIAL_PPN
+def next_ppn():
+  global next_free_ppn
+  ppn = next_free_ppn
+  next_free_ppn += 1
+  return ppn
+
+def print_pages():
+  for page in pages:
+    print(pages[page])
+
+##################################
+# helper variables
+##################################
+rv32 = Architecture(32)
+rv64 = Architecture(64)
+
+if __name__ == "__main__":
+  arch = rv64
+  pgdir = PageTable("page_directory", next_ppn(), arch)
+
+  # Directly map the first 20 pages of RAM
+  for page in range(20):
+    vaddr = 0x80000000 + (arch.PGSIZE * page)
+    paddr = 0x80000000 + (arch.PGSIZE * page)
+    pgdir.add_mapping(vaddr, paddr, PTE_R | PTE_W | PTE_U | PTE_X | PTE_V)
+  """
+  supervisor_pgdir = PageTable("sdir", next_ppn(), rv64)
+  supervisor_pgdir.add_mapping(0x80000000, 0x80000000, PTE_R | PTE_W | PTE_X)
+  supervisor_pgdir.add_mapping(0x80000001, 0x80000001, PTE_R | PTE_W | PTE_X)
+  supervisor_pgdir.add_mapping(0x80001000, 0x80000000, PTE_R | PTE_W | PTE_X)
+  supervisor_pgdir.add_mapping(0xffff0000, 0x80000000, PTE_R | PTE_W | PTE_X)
+  """
+
+  print_pages()