From c57e39d52ba97aac00fdc1b78f7ea0b49d0bb694 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?S=C3=B6nke=20Holz?= <sholz8530@gmail.com>
Date: Sun, 21 Apr 2024 12:52:25 +0200
Subject: [PATCH] Kernel/riscv64: Don't flush the entire TLB in
 Processor::flush_tlb_local

---
 Kernel/Arch/riscv64/Processor.cpp | 16 +++++++++++++---
 1 file changed, 13 insertions(+), 3 deletions(-)

diff --git a/Kernel/Arch/riscv64/Processor.cpp b/Kernel/Arch/riscv64/Processor.cpp
index 17f69d13a38..b67b9c2436e 100644
--- a/Kernel/Arch/riscv64/Processor.cpp
+++ b/Kernel/Arch/riscv64/Processor.cpp
@@ -137,10 +137,19 @@ template<typename T>
 }
 
 template<typename T>
-void ProcessorBase<T>::flush_tlb_local(VirtualAddress, size_t)
+void ProcessorBase<T>::flush_tlb_local(VirtualAddress vaddr, size_t page_count)
 {
-    // FIXME: Don't flush all pages
-    flush_entire_tlb_local();
+    auto addr = vaddr.get();
+    while (page_count > 0) {
+        // clang-format off
+        asm volatile("sfence.vma %0"
+             :
+             : "r"(addr)
+             : "memory");
+        // clang-format on
+        addr += PAGE_SIZE;
+        page_count--;
+    }
 }
 
 template<typename T>
@@ -152,6 +161,7 @@ void ProcessorBase<T>::flush_entire_tlb_local()
 template<typename T>
 void ProcessorBase<T>::flush_tlb(Memory::PageDirectory const*, VirtualAddress vaddr, size_t page_count)
 {
+    // FIXME: Use the SBI RFENCE extension to flush the TLB of other harts when we support SMP on riscv64.
     flush_tlb_local(vaddr, page_count);
 }