Kernel: Add PageTable API to invalidate full address space

kernel/arch/i686/PageTable.cpp

@@ -26,8 +26,6 @@ namespace Kernel
 	constexpr uint64_t s_page_flag_mask = 0x8000000000000FFF;
 	constexpr uint64_t s_page_addr_mask = ~s_page_flag_mask;
 
-	static bool s_is_initialized = false;
-
 	static PageTable* s_kernel = nullptr;
 	static bool s_has_nxe = false;
 	static bool s_has_pge = false;
@@ -341,32 +339,10 @@ namespace Kernel
 		const bool is_userspace = (vaddr < KERNEL_OFFSET);
 		if (is_userspace && this != &PageTable::current())
 			;
-		else if (pages <= 32 || !s_is_initialized)
+		else if (pages >= full_tlb_flush_threshold)
-		{
+			invalidate_full_address_space(!is_userspace);
-			for (size_t i = 0; i < pages; i++)
+		else for (size_t i = 0; i < pages; i++)
 			asm volatile("invlpg (%0)" :: "r"(vaddr + i * PAGE_SIZE));
-		}
-		else if (is_userspace || !s_has_pge)
-		{
-			asm volatile("movl %0, %%cr3" :: "r"(static_cast<uint32_t>(m_highest_paging_struct)));
-		}
-		else
-		{
-			asm volatile(
-				"movl %%cr4, %%eax;"
-
-				"andl $~0x80, %%eax;"
-				"movl %%eax, %%cr4;"
-
-				"movl %0, %%cr3;"
-
-				"orl $0x80, %%eax;"
-				"movl %%eax, %%cr4;"
-				:
-				: "r"(static_cast<uint32_t>(m_highest_paging_struct))
-				: "eax"
-			);
-		}
 
 		if (send_smp_message)
 		{
@@ -381,6 +357,34 @@ namespace Kernel
 		}
 	}
 
+	void PageTable::invalidate_full_address_space(bool global)
+	{
+		if (!global || !s_has_pge)
+		{
+			asm volatile(
+				"movl %%cr3, %%eax;"
+				"movl %%eax, %%cr3;"
+				::: "eax"
+			);
+		}
+		else
+		{
+			asm volatile(
+				"movl %%cr4, %%eax;"
+
+				"andl $~0x80, %%eax;"
+				"movl %%eax, %%cr4;"
+
+				"movl %%cr3, %%ecx;"
+				"movl %%ecx, %%cr3;"
+
+				"orl $0x80, %%eax;"
+				"movl %%eax, %%cr4;"
+				::: "eax", "ecx"
+			);
+		}
+	}
+
 	void PageTable::unmap_page(vaddr_t vaddr, bool invalidate)
 	{
 		ASSERT(vaddr);

kernel/arch/x86_64/PageTable.cpp

@@ -24,7 +24,6 @@ namespace Kernel
 	SpinLock PageTable::s_fast_page_lock;
 
 	static constexpr vaddr_t s_hhdm_offset = 0xFFFF800000000000;
-	static bool s_is_initialized = false;
 
 	constexpr uint64_t s_page_flag_mask = 0x8000000000000FFF;
 	constexpr uint64_t s_page_addr_mask = ~s_page_flag_mask;
@@ -470,32 +469,10 @@ namespace Kernel
 		const bool is_userspace = (vaddr < KERNEL_OFFSET);
 		if (is_userspace && this != &PageTable::current())
 			;
-		else if (pages <= 32 || !s_is_initialized)
+		else if (pages >= full_tlb_flush_threshold)
-		{
+			invalidate_full_address_space(!is_userspace);
-			for (size_t i = 0; i < pages; i++)
+		else for (size_t i = 0; i < pages; i++)
 			asm volatile("invlpg (%0)" :: "r"(vaddr + i * PAGE_SIZE));
-		}
-		else if (is_userspace || !s_has_pge)
-		{
-			asm volatile("movq %0, %%cr3" :: "r"(m_highest_paging_struct));
-		}
-		else
-		{
-			asm volatile(
-				"movq %%cr4, %%rax;"
-
-				"andq $~0x80, %%rax;"
-				"movq %%rax, %%cr4;"
-
-				"movq %0, %%cr3;"
-
-				"orq $0x80, %%rax;"
-				"movq %%rax, %%cr4;"
-				:
-				: "r"(m_highest_paging_struct)
-				: "rax"
-			);
-		}
 
 		if (send_smp_message)
 		{
@@ -510,6 +487,34 @@ namespace Kernel
 		}
 	}
 
+	void PageTable::invalidate_full_address_space(bool global)
+	{
+		if (!global || !s_has_pge)
+		{
+			asm volatile(
+				"movq %%cr3, %%rax;"
+				"movq %%rax, %%cr3;"
+				::: "rax"
+			);
+		}
+		else
+		{
+			asm volatile(
+				"movq %%cr4, %%rax;"
+
+				"andq $~0x80, %%rax;"
+				"movq %%rax, %%cr4;"
+
+				"movq %%cr3, %%rcx;"
+				"movq %%rcx, %%cr3;"
+
+				"orq $0x80, %%rax;"
+				"movq %%rax, %%cr4;"
+				::: "rax", "rcx"
+			);
+		}
+	}
+
 	void PageTable::unmap_page(vaddr_t vaddr, bool invalidate)
 	{
 		ASSERT(vaddr);

kernel/include/kernel/Memory/PageTable.h

@@ -45,6 +45,8 @@ namespace Kernel
 			WriteThrough,
 		};
 
+		static constexpr bool full_tlb_flush_threshold = 32;
+
 	public:
 		static void initialize_fast_page();
 		static void initialize_and_load();
@@ -133,6 +135,7 @@ namespace Kernel
 
 		void invalidate_page(vaddr_t addr, bool send_smp_message) { invalidate_range(addr, 1, send_smp_message); }
 		void invalidate_range(vaddr_t addr, size_t pages, bool send_smp_message);
+		void invalidate_full_address_space(bool global);
 
 		InterruptState lock() const { return m_lock.lock(); }
 		void unlock(InterruptState state) const { m_lock.unlock(state); }