From 822d4f891cf943a9a32d02063dbfc6ba653c2d70 Mon Sep 17 00:00:00 2001
From: Mark Poliakov <mark@alnyan.me>
Date: Tue, 21 Jan 2025 18:50:26 +0200
Subject: [PATCH] rv64: implement address space dropping

---
 kernel/arch/riscv64/src/mem/process.rs | 32 +++++++------
 kernel/arch/riscv64/src/mem/table.rs   | 63 ++++++++++++++++++++++++--
 2 files changed, 78 insertions(+), 17 deletions(-)

diff --git a/kernel/arch/riscv64/src/mem/process.rs b/kernel/arch/riscv64/src/mem/process.rs
index 6e89e18f..0078684b 100644
--- a/kernel/arch/riscv64/src/mem/process.rs
+++ b/kernel/arch/riscv64/src/mem/process.rs
@@ -7,7 +7,9 @@ use libk_mm_interface::{
     address::{AsPhysicalAddress, PhysicalAddress},
     pointer::PhysicalRefMut,
     process::ProcessAddressSpaceManager,
-    table::{EntryLevel, EntryLevelExt, MapAttributes, NextPageTable, TableAllocator},
+    table::{
+        EntryLevel, EntryLevelDrop, EntryLevelExt, MapAttributes, NextPageTable, TableAllocator,
+    },
 };
 use memtables::riscv64::PageAttributes;
 use yggdrasil_abi::error::Error;
@@ -15,8 +17,8 @@ use yggdrasil_abi::error::Error;
 use crate::mem::{clone_kernel_tables, table::PageEntry};
 
 use super::{
-    table::{PageTable, L1, L2, L3},
-    tlb_flush_va_asid, KernelTableManagerImpl, USER_BOUNDARY,
+    table::{DroppableRange, PageTable, L1, L2, L3},
+    KernelTableManagerImpl, USER_BOUNDARY,
 };
 
 pub struct ProcessAddressSpaceImpl<TA: TableAllocator> {
@@ -81,7 +83,9 @@ impl<TA: TableAllocator> ProcessAddressSpaceManager<TA> for ProcessAddressSpaceI
         (physical, self.asid as u64)
     }
 
-    unsafe fn clear(&mut self) {}
+    unsafe fn clear(&mut self) {
+        unsafe { self.l1.drop_range::<TA>(L1::DROPPABLE_RANGE) };
+    }
 }
 
 impl<TA: TableAllocator> ProcessAddressSpaceImpl<TA> {
@@ -109,7 +113,7 @@ impl<TA: TableAllocator> ProcessAddressSpaceImpl<TA> {
         }
 
         l3[l3i] = entry;
-        tlb_flush_va_asid(virt, self.asid as usize);
+        super::tlb_flush_va_asid(virt, self.asid as usize);
 
         Ok(())
     }
@@ -126,7 +130,7 @@ impl<TA: TableAllocator> ProcessAddressSpaceImpl<TA> {
         let page = l3[l3i].as_page().ok_or(Error::DoesNotExist)?;
 
         l3[l3i] = PageEntry::INVALID;
-        tlb_flush_va_asid(virt, self.asid as usize);
+        super::tlb_flush_va_asid(virt, self.asid as usize);
 
         Ok(page)
     }
@@ -155,14 +159,14 @@ impl<TA: TableAllocator> ProcessAddressSpaceImpl<TA> {
 
 impl<TA: TableAllocator> Drop for ProcessAddressSpaceImpl<TA> {
     fn drop(&mut self) {
-        // TODO
-        // // SAFETY: with safe usage of the ProcessAddressSpaceImpl, clearing and dropping
-        // // is safe, no one refers to the memory
-        // unsafe {
-        //     self.clear();
-        //     let l1_phys = self.l1.as_physical_address();
-        //     TA::free_page_table(l1_phys);
-        // }
+        // SAFETY: with safe usage of the ProcessAddressSpaceImpl, clearing and dropping
+        // is safe, no one refers to the memory
+        unsafe {
+            self.clear();
+            let l1_phys = self.l1.as_physical_address();
+            TA::free_page_table(l1_phys);
+            super::tlb_flush_asid(self.asid as usize);
+        }
     }
 }
 
diff --git a/kernel/arch/riscv64/src/mem/table.rs b/kernel/arch/riscv64/src/mem/table.rs
index 1be91b4a..a3481a0c 100644
--- a/kernel/arch/riscv64/src/mem/table.rs
+++ b/kernel/arch/riscv64/src/mem/table.rs
@@ -1,16 +1,19 @@
 use core::{
     marker::PhantomData,
-    ops::{Index, IndexMut},
+    ops::{Index, IndexMut, Range},
 };
 
 use libk_mm_interface::{
     address::{AsPhysicalAddress, PhysicalAddress},
     pointer::{PhysicalRef, PhysicalRefMut},
-    table::{EntryLevel, NextPageTable, NonTerminalEntryLevel, TableAllocator},
+    table::{
+        page_index, EntryLevel, EntryLevelDrop, NextPageTable, NonTerminalEntryLevel,
+        TableAllocator,
+    },
 };
 use yggdrasil_abi::error::Error;
 
-use super::KernelTableManagerImpl;
+use super::{KernelTableManagerImpl, USER_BOUNDARY};
 
 pub use memtables::riscv64::PageAttributes;
 
@@ -44,6 +47,18 @@ pub struct PageTable<L: EntryLevel> {
 #[derive(Clone, Copy, Debug, PartialEq)]
 pub struct PageEntry<L: EntryLevel>(pub u64, PhantomData<L>);
 
+pub(super) trait DroppableRange {
+    const DROPPABLE_RANGE: Range<usize>;
+}
+
+impl DroppableRange for L1 {
+    const DROPPABLE_RANGE: Range<usize> = 0..page_index::<L1>(USER_BOUNDARY);
+}
+
+impl DroppableRange for L2 {
+    const DROPPABLE_RANGE: Range<usize> = 0..512;
+}
+
 impl NonTerminalEntryLevel for L1 {
     type NextLevel = L2;
 }
@@ -93,6 +108,48 @@ impl<L: EntryLevel> PageEntry<L> {
     }
 }
 
+impl<L: NonTerminalEntryLevel + DroppableRange> EntryLevelDrop for PageTable<L>
+where
+    PageTable<L::NextLevel>: EntryLevelDrop,
+{
+    const FULL_RANGE: Range<usize> = L::DROPPABLE_RANGE;
+
+    unsafe fn drop_range<TA: TableAllocator>(&mut self, range: Range<usize>) {
+        for index in range {
+            let entry = self[index];
+
+            if let Some(table) = entry.as_table() {
+                unsafe {
+                    let mut table_ref: PhysicalRefMut<
+                        PageTable<L::NextLevel>,
+                        KernelTableManagerImpl,
+                    > = PhysicalRefMut::map(table);
+
+                    table_ref.drop_all::<TA>();
+
+                    TA::free_page_table(table);
+                }
+            } else if entry.is_present() {
+                // Memory must've been cleared beforehand, so no non-table entries must be present
+                panic!(
+                    "Expected a table containing only tables, got table[{}] = {:#x?}",
+                    index, entry.0
+                );
+            }
+
+            self[index] = PageEntry::INVALID;
+            // dc_cvac((&raw const self[index]).addr());
+        }
+    }
+}
+
+impl EntryLevelDrop for PageTable<L3> {
+    const FULL_RANGE: Range<usize> = 0..512;
+
+    // Do nothing
+    unsafe fn drop_range<TA: TableAllocator>(&mut self, _range: Range<usize>) {}
+}
+
 impl<L: NonTerminalEntryLevel + 'static> NextPageTable for PageTable<L> {
     type NextLevel = PageTable<L::NextLevel>;
     type TableRef = PhysicalRef<'static, PageTable<L::NextLevel>, KernelTableManagerImpl>;