3 files changed, 65 insertions, 23 deletions

diff --git a/src/common/page_table.cpp b/src/common/page_table.cpp
index bccea0894..8fd8620fd 100644
--- a/src/common/page_table.cpp
+++ b/src/common/page_table.cpp
@@ -10,16 +10,10 @@ PageTable::PageTable() = default;
 
 PageTable::~PageTable() noexcept = default;
 
-void PageTable::Resize(std::size_t address_space_width_in_bits, std::size_t page_size_in_bits,
-                       bool has_attribute) {
-    const std::size_t num_page_table_entries{1ULL
-                                             << (address_space_width_in_bits - page_size_in_bits)};
+void PageTable::Resize(size_t address_space_width_in_bits, size_t page_size_in_bits) {
+    const size_t num_page_table_entries{1ULL << (address_space_width_in_bits - page_size_in_bits)};
     pointers.resize(num_page_table_entries);
     backing_addr.resize(num_page_table_entries);
-
-    if (has_attribute) {
-        attributes.resize(num_page_table_entries);
-    }
 }
 
 } // namespace Common
diff --git a/src/common/page_table.h b/src/common/page_table.h
index 9754fabf9..8d4ee9249 100644
--- a/src/common/page_table.h
+++ b/src/common/page_table.h
@@ -4,6 +4,7 @@
 
 #pragma once
 
+#include <atomic>
 #include <tuple>
 
 #include "common/common_types.h"
@@ -20,10 +21,6 @@ enum class PageType : u8 {
     /// Page is mapped to regular memory, but also needs to check for rasterizer cache flushing and
     /// invalidation
     RasterizerCachedMemory,
-    /// Page is mapped to a I/O region. Writing and reading to this page is handled by functions.
-    Special,
-    /// Page is allocated for use.
-    Allocated,
 };
 
 struct SpecialRegion {
@@ -48,6 +45,59 @@ struct SpecialRegion {
  * mimics the way a real CPU page table works.
  */
 struct PageTable {
+    /// Number of bits reserved for attribute tagging.
+    /// This can be at most the guaranteed alignment of the pointers in the page table.
+    static constexpr int ATTRIBUTE_BITS = 2;
+
+    /**
+     * Pair of host pointer and page type attribute.
+     * This uses the lower bits of a given pointer to store the attribute tag.
+     * Writing and reading the pointer attribute pair is guaranteed to be atomic for the same method
+     * call. In other words, they are guaranteed to be synchronized at all times.
+     */
+    class PageInfo {
+    public:
+        /// Returns the page pointer
+        [[nodiscard]] u8* Pointer() const noexcept {
+            return ExtractPointer(raw.load(std::memory_order_relaxed));
+        }
+
+        /// Returns the page type attribute
+        [[nodiscard]] PageType Type() const noexcept {
+            return ExtractType(raw.load(std::memory_order_relaxed));
+        }
+
+        /// Returns the page pointer and attribute pair, extracted from the same atomic read
+        [[nodiscard]] std::pair<u8*, PageType> PointerType() const noexcept {
+            const uintptr_t non_atomic_raw = raw.load(std::memory_order_relaxed);
+            return {ExtractPointer(non_atomic_raw), ExtractType(non_atomic_raw)};
+        }
+
+        /// Returns the raw representation of the page information.
+        /// Use ExtractPointer and ExtractType to unpack the value.
+        [[nodiscard]] uintptr_t Raw() const noexcept {
+            return raw.load(std::memory_order_relaxed);
+        }
+
+        /// Write a page pointer and type pair atomically
+        void Store(u8* pointer, PageType type) noexcept {
+            raw.store(reinterpret_cast<uintptr_t>(pointer) | static_cast<uintptr_t>(type));
+        }
+
+        /// Unpack a pointer from a page info raw representation
+        [[nodiscard]] static u8* ExtractPointer(uintptr_t raw) noexcept {
+            return reinterpret_cast<u8*>(raw & (~uintptr_t{0} << ATTRIBUTE_BITS));
+        }
+
+        /// Unpack a page type from a page info raw representation
+        [[nodiscard]] static PageType ExtractType(uintptr_t raw) noexcept {
+            return static_cast<PageType>(raw & ((uintptr_t{1} << ATTRIBUTE_BITS) - 1));
+        }
+
+    private:
+        std::atomic<uintptr_t> raw;
+    };
+
     PageTable();
     ~PageTable() noexcept;
 
@@ -63,20 +113,16 @@ struct PageTable {
      *
      * @param address_space_width_in_bits The address size width in bits.
      * @param page_size_in_bits           The page size in bits.
-     * @param has_attribute               Whether or not this page has any backing attributes.
      */
-    void Resize(std::size_t address_space_width_in_bits, std::size_t page_size_in_bits,
-                bool has_attribute);
+    void Resize(size_t address_space_width_in_bits, size_t page_size_in_bits);
 
     /**
      * Vector of memory pointers backing each page. An entry can only be non-null if the
-     * corresponding entry in the `attributes` vector is of type `Memory`.
+     * corresponding attribute element is of type `Memory`.
      */
-    VirtualBuffer<u8*> pointers;
+    VirtualBuffer<PageInfo> pointers;
 
     VirtualBuffer<u64> backing_addr;
-
-    VirtualBuffer<PageType> attributes;
 };
 
 } // namespace Common
diff --git a/src/common/virtual_buffer.h b/src/common/virtual_buffer.h
index 91d430036..fb1a6f81f 100644
--- a/src/common/virtual_buffer.h
+++ b/src/common/virtual_buffer.h
@@ -15,10 +15,12 @@ void FreeMemoryPages(void* base, std::size_t size) noexcept;
 template <typename T>
 class VirtualBuffer final {
 public:
-    static_assert(
-        std::is_trivially_constructible_v<T>,
-        "T must be trivially constructible, as non-trivial constructors will not be executed "
-        "with the current allocator");
+    // TODO: Uncomment this and change Common::PageTable::PageInfo to be trivially constructible
+    // using std::atomic_ref once libc++ has support for it
+    // static_assert(
+    //     std::is_trivially_constructible_v<T>,
+    //     "T must be trivially constructible, as non-trivial constructors will not be executed "
+    //     "with the current allocator");
 
     constexpr VirtualBuffer() = default;
     explicit VirtualBuffer(std::size_t count) : alloc_size{count * sizeof(T)} {