Add a TLB module

2 files changed, 22 insertions, 1 deletions

diff --git a/src/lib.rs b/src/lib.rs
index 8f799cd..a6283db 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -41,4 +41,5 @@ pub mod syscall;
 pub mod perfcnt;
 pub mod cpuid {
     pub use raw_cpuid::*;
-}
\ No newline at end of file
+}
+pub mod tlb;
diff --git a/src/tlb.rs b/src/tlb.rs
new file mode 100644
index 0000000..f20e37e
--- /dev/null
+++ b/src/tlb.rs
@@ -0,0 +1,20 @@
+//! Functions to flush the translation lookaside buffer (TLB).
+
+/// Invalidate the given address in the TLB using the `invlpg` instruction.
+///
+/// # Safety
+/// This function is unsafe as it causes a general protection fault (GP) if the current privilege
+/// level is not 0.
+pub unsafe fn flush(addr: usize) {
+    asm!("invlpg ($0)" :: "r" (addr) : "memory" : "volatile");
+}
+
+/// Invalidate the TLB completely by reloading the CR3 register.
+///
+/// # Safety
+/// This function is unsafe as it causes a general protection fault (GP) if the current privilege
+/// level is not 0.
+pub unsafe fn flush_all() {
+    use controlregs::{cr3, cr3_write};
+    cr3_write(cr3())
+}