From 192dd8283fcc45ebfa94500695bb3fabdc089887 Mon Sep 17 00:00:00 2001
From: Andrey Turkin <andrey.turkin@gmail.com>
Date: Fri, 12 Jul 2024 09:01:23 +0300
Subject: [PATCH] RISC-V: implement TLB flush

---
 src/arch/riscv64/sv39.rs | 11 ++++++++---
 src/arch/riscv64/sv48.rs | 12 +++++++++---
 2 files changed, 17 insertions(+), 6 deletions(-)

diff --git a/src/arch/riscv64/sv39.rs b/src/arch/riscv64/sv39.rs
index 607f4a2b1e..81dd8583f1 100644
--- a/src/arch/riscv64/sv39.rs
+++ b/src/arch/riscv64/sv39.rs
@@ -32,9 +32,13 @@ impl Arch for RiscV64Sv39Arch {
     }
 
     #[inline(always)]
-    unsafe fn invalidate(_address: VirtualAddress) {
-        //TODO: can one address be invalidated?
-        Self::invalidate_all();
+    unsafe fn invalidate(address: VirtualAddress) {
+        asm!("sfence.vma {}", in(reg) address.data());
+    }
+
+    #[inline(always)]
+    unsafe fn invalidate_all() {
+        asm!("sfence.vma");
     }
 
     #[inline(always)]
@@ -51,6 +55,7 @@ impl Arch for RiscV64Sv39Arch {
         let satp = (8 << 60) | // Sv39 MODE
             (address.data() >> Self::PAGE_SHIFT); // Convert to PPN (TODO: ensure alignment)
         asm!("csrw satp, {0}", in(reg) satp);
+        Self::invalidate_all();
     }
 
     fn virt_is_valid(address: VirtualAddress) -> bool {
diff --git a/src/arch/riscv64/sv48.rs b/src/arch/riscv64/sv48.rs
index e80e52a000..b6796ab3d8 100644
--- a/src/arch/riscv64/sv48.rs
+++ b/src/arch/riscv64/sv48.rs
@@ -32,9 +32,13 @@ impl Arch for RiscV64Sv48Arch {
     }
 
     #[inline(always)]
-    unsafe fn invalidate(_address: VirtualAddress) {
-        //TODO: can one address be invalidated?
-        Self::invalidate_all();
+    unsafe fn invalidate(address: VirtualAddress) {
+        asm!("sfence.vma {}", in(reg) address.data());
+    }
+
+    #[inline(always)]
+    unsafe fn invalidate_all() {
+        asm!("sfence.vma");
     }
 
     #[inline(always)]
@@ -51,7 +55,9 @@ impl Arch for RiscV64Sv48Arch {
         let satp = (9 << 60) | // Sv48 MODE
             (address.data() >> Self::PAGE_SHIFT); // Convert to PPN (TODO: ensure alignment)
         asm!("csrw satp, {0}", in(reg) satp);
+        Self::invalidate_all();
     }
+
     fn virt_is_valid(address: VirtualAddress) -> bool {
         // RISC-V SV48 uses 48-bit sign-extended addresses, identical to 4-level paging on x86_64.
         let mask = !((Self::PAGE_ADDRESS_SIZE as usize - 1) >> 1);