PCID support instructions

src/asm/asm.s

@@ -94,6 +94,12 @@ _x86_64_asm_invlpg:
     invlpg (%rdi)
     retq
 
+.global _x86_64_asm_invpcid
+.p2align 4
+_x86_64_asm_invpcid:
+    invpcid (%rsi), %rdi
+    retq
+
 .global _x86_64_asm_ltr
 .p2align 4
 _x86_64_asm_ltr:

src/asm/mod.rs

@@ -138,6 +138,12 @@ extern "C" {
     )]
     pub(crate) fn x86_64_asm_invlpg(addr: u64);
 
+    #[cfg_attr(
+        any(target_env = "gnu", target_env = "musl"),
+        link_name = "_x86_64_asm_invpcid"
+    )]
+    pub(crate) fn x86_64_asm_invpcid(kind: u64, desc: u64);
+
     #[cfg_attr(
         any(target_env = "gnu", target_env = "musl"),
         link_name = "_x86_64_asm_read_cr0"

src/instructions/tlb.rs

@@ -23,3 +23,84 @@ pub fn flush_all() {
     let (frame, flags) = Cr3::read();
     unsafe { Cr3::write(frame, flags) }
 }
+
+/// The Invalidate PCID Command to execute.
+#[derive(Debug)]
+pub enum InvPicdCommand {
+    /// The logical processor invalidates mappings—except global translations—for the linear address and PCID specified.
+    IndividualAddressInvalidation(VirtAddr, Pcid),
+
+    /// The logical processor invalidates all mappings—except global translations—associated with the PCID.
+    SingleContextInvalidation(Pcid),
+
+    /// The logical processor invalidates all mappings—including global translations—associated with any PCID.
+    AllContextInvalidationIncludeGlobal,
+
+    /// The logical processor invalidates all mappings—except global translations—associated with any PCID.
+    AllContextInvalidationExcludeGlobal,
+}
+
+/// The INVPCID descriptor comprises 128 bits and consists of a PCID and a linear address.
+/// For INVPCID type 0, the processor uses the full 64 bits of the linear address even outside 64-bit mode; the linear address is not used for other INVPCID types.
+#[repr(u128)]
+#[derive(Debug)]
+struct InvpcidDescriptor {
+    address: u64,
+    pcid: u64,
+}
+
+/// Structure of a PCID. A PCID has to be <= 4096 for x86_64.
+#[repr(transparent)]
+#[derive(Debug)]
+pub struct Pcid(u64);
+
+impl Pcid {
+    /// Create a new PCID. Will result in a failure if the value of
+    /// PCID is out of expected bounds.
+    pub fn new(pcid: u16) -> Result<Pcid, &'static str> {
+        if pcid >= 4096 {
+            Err("PCID should be < 4096.")
+        } else {
+            Ok(Pcid(pcid as u64))
+        }
+    }
+
+    /// Get the value of the current PCID.
+    pub fn value(&self) -> u16 {
+        self.0 as u16
+    }
+}
+
+/// Invalidate the given address in the TLB using the `invpcid` instruction.
+#[inline]
+pub fn flush_pcid(command: InvPicdCommand) {
+    let mut desc = InvpcidDescriptor {
+        address: 0,
+        pcid: 0,
+    };
+
+    let kind;
+    match command {
+        InvPicdCommand::IndividualAddressInvalidation(addr, pcid) => {
+            kind = 0;
+            desc.pcid = pcid.value() as u64;
+            desc.address = addr.as_u64()
+        }
+        InvPicdCommand::SingleContextInvalidation(pcid) => {
+            kind = 1;
+            desc.pcid = pcid.0
+        }
+        InvPicdCommand::AllContextInvalidationIncludeGlobal => kind = 2,
+        InvPicdCommand::AllContextInvalidationExcludeGlobal => kind = 3,
+    }
+
+    #[cfg(feature = "inline_asm")]
+    unsafe {
+        llvm_asm!("invpcid ($0), $1" :: "r" (&desc as *const InvpcidDescriptor as u64), "r" (kind) : "memory")
+    };
+
+    #[cfg(not(feature = "inline_asm"))]
+    unsafe {
+        crate::asm::x86_64_asm_invpcid(kind, &desc as *const InvpcidDescriptor as u64)
+    };
+}

src/registers/control.rs

@@ -129,7 +129,7 @@ bitflags! {
 mod x86_64 {
     use super::*;
     use crate::structures::paging::PhysFrame;
-    use crate::{PhysAddr, VirtAddr};
+    use crate::{instructions::tlb::Pcid, PhysAddr, VirtAddr};
 
     impl Cr0 {
         /// Read the current set of CR0 flags.
@@ -251,6 +251,28 @@ mod x86_64 {
             (frame, flags)
         }
 
+        /// Read the current P4 table address from the CR3 register along with PCID.
+        /// The correct functioning of this requires CR4.PCIDE = 1.
+        #[inline]
+        pub fn read_as_pcid() -> (PhysFrame, u16) {
+            let value: u64;
+
+            #[cfg(feature = "inline_asm")]
+            unsafe {
+                llvm_asm!("mov %cr3, $0" : "=r" (value));
+            }
+
+            #[cfg(not(feature = "inline_asm"))]
+            unsafe {
+                value = crate::asm::x86_64_asm_read_cr3();
+            }
+
+            let addr = PhysAddr::new(value & 0x_000f_ffff_ffff_f000);
+            let frame = PhysFrame::containing_address(addr);
+            let pcid = value & 0xFFF;
+            (frame, pcid as u16)
+        }
+
         /// Write a new P4 table address into the CR3 register.
         ///
         /// ## Safety
@@ -267,6 +289,24 @@ mod x86_64 {
             #[cfg(not(feature = "inline_asm"))]
             crate::asm::x86_64_asm_write_cr3(value)
         }
+
+        /// Write a new P4 table address into the CR3 register.
+        ///
+        /// ## Safety
+        /// Changing the level 4 page table is unsafe, because it's possible to violate memory safety by
+        /// changing the page mapping.
+        /// Also, using this requires CR4.PCIDE flag to be set.
+        #[inline]
+        pub unsafe fn write_pcid(frame: PhysFrame, pcid: Pcid) {
+            let addr = frame.start_address();
+            let value = addr.as_u64() | pcid.value() as u64;
+
+            #[cfg(feature = "inline_asm")]
+            llvm_asm!("mov $0, %cr3" :: "r" (value) : "memory");
+
+            #[cfg(not(feature = "inline_asm"))]
+            crate::asm::x86_64_asm_write_cr3(value)
+        }
     }
 
     impl Cr4 {