1 files changed, 60 insertions, 45 deletions

diff --git a/src/common/address_space.h b/src/common/address_space.h
index 5b3832f07..bf649018c 100644
--- a/src/common/address_space.h
+++ b/src/common/address_space.h
@@ -23,9 +23,29 @@ template <typename VaType, VaType UnmappedVa, typename PaType, PaType UnmappedPa
           bool PaContigSplit, size_t AddressSpaceBits, typename ExtraBlockInfo = EmptyStruct>
 requires AddressSpaceValid<VaType, AddressSpaceBits>
 class FlatAddressSpaceMap {
-private:
-    std::function<void(VaType, VaType)>
-        unmapCallback{}; //!< Callback called when the mappings in an region have changed
+public:
+    /// The maximum VA that this AS can technically reach
+    static constexpr VaType VaMaximum{(1ULL << (AddressSpaceBits - 1)) +
+                                      ((1ULL << (AddressSpaceBits - 1)) - 1)};
+
+    explicit FlatAddressSpaceMap(VaType va_limit,
+                                 std::function<void(VaType, VaType)> unmap_callback = {});
+
+    FlatAddressSpaceMap() = default;
+
+    void Map(VaType virt, PaType phys, VaType size, ExtraBlockInfo extra_info = {}) {
+        std::scoped_lock lock(block_mutex);
+        MapLocked(virt, phys, size, extra_info);
+    }
+
+    void Unmap(VaType virt, VaType size) {
+        std::scoped_lock lock(block_mutex);
+        UnmapLocked(virt, size);
+    }
+
+    VaType GetVALimit() const {
+        return va_limit;
+    }
 
 protected:
     /**
@@ -33,68 +53,55 @@ protected:
      * another block with a different phys address is hit
      */
     struct Block {
-        VaType virt{UnmappedVa}; //!< VA of the block
-        PaType phys{UnmappedPa}; //!< PA of the block, will increase 1-1 with VA until a new block
-                                 //!< is encountered
-        [[no_unique_address]] ExtraBlockInfo extraInfo;
+        /// VA of the block
+        VaType virt{UnmappedVa};
+        /// PA of the block, will increase 1-1 with VA until a new block is encountered
+        PaType phys{UnmappedPa};
+        [[no_unique_address]] ExtraBlockInfo extra_info;
 
         Block() = default;
 
-        Block(VaType virt_, PaType phys_, ExtraBlockInfo extraInfo_)
-            : virt(virt_), phys(phys_), extraInfo(extraInfo_) {}
+        Block(VaType virt_, PaType phys_, ExtraBlockInfo extra_info_)
+            : virt(virt_), phys(phys_), extra_info(extra_info_) {}
 
-        constexpr bool Valid() {
+        bool Valid() const {
             return virt != UnmappedVa;
         }
 
-        constexpr bool Mapped() {
+        bool Mapped() const {
             return phys != UnmappedPa;
         }
 
-        constexpr bool Unmapped() {
+        bool Unmapped() const {
             return phys == UnmappedPa;
         }
 
-        bool operator<(const VaType& pVirt) const {
-            return virt < pVirt;
+        bool operator<(const VaType& p_virt) const {
+            return virt < p_virt;
         }
     };
 
-    std::mutex blockMutex;
-    std::vector<Block> blocks{Block{}};
-
     /**
      * @brief Maps a PA range into the given AS region
-     * @note blockMutex MUST be locked when calling this
+     * @note block_mutex MUST be locked when calling this
      */
-    void MapLocked(VaType virt, PaType phys, VaType size, ExtraBlockInfo extraInfo);
+    void MapLocked(VaType virt, PaType phys, VaType size, ExtraBlockInfo extra_info);
 
     /**
      * @brief Unmaps the given range and merges it with other unmapped regions
-     * @note blockMutex MUST be locked when calling this
+     * @note block_mutex MUST be locked when calling this
      */
     void UnmapLocked(VaType virt, VaType size);
 
-public:
-    static constexpr VaType VaMaximum{(1ULL << (AddressSpaceBits - 1)) +
-                                      ((1ULL << (AddressSpaceBits - 1)) -
-                                       1)}; //!< The maximum VA that this AS can technically reach
-
-    VaType vaLimit{VaMaximum}; //!< A soft limit on the maximum VA of the AS
-
-    FlatAddressSpaceMap(VaType vaLimit, std::function<void(VaType, VaType)> unmapCallback = {});
-
-    FlatAddressSpaceMap() = default;
+    std::mutex block_mutex;
+    std::vector<Block> blocks{Block{}};
 
-    void Map(VaType virt, PaType phys, VaType size, ExtraBlockInfo extraInfo = {}) {
-        std::scoped_lock lock(blockMutex);
-        MapLocked(virt, phys, size, extraInfo);
-    }
+    /// a soft limit on the maximum VA of the AS
+    VaType va_limit{VaMaximum};
 
-    void Unmap(VaType virt, VaType size) {
-        std::scoped_lock lock(blockMutex);
-        UnmapLocked(virt, size);
-    }
+private:
+    /// Callback called when the mappings in an region have changed
+    std::function<void(VaType, VaType)> unmap_callback{};
 };
 
 /**
@@ -108,14 +115,8 @@ class FlatAllocator
 private:
     using Base = FlatAddressSpaceMap<VaType, UnmappedVa, bool, false, false, AddressSpaceBits>;
 
-    VaType currentLinearAllocEnd; //!< The end address for the initial linear allocation pass, once
-                                  //!< this reaches the AS limit the slower allocation path will be
-                                  //!< used
-
 public:
-    VaType vaStart; //!< The base VA of the allocator, no allocations will be below this
-
-    FlatAllocator(VaType vaStart, VaType vaLimit = Base::VaMaximum);
+    explicit FlatAllocator(VaType va_start, VaType va_limit = Base::VaMaximum);
 
     /**
      * @brief Allocates a region in the AS of the given size and returns its address
@@ -131,5 +132,19 @@ public:
      * @brief Frees an AS region so it can be used again
      */
     void Free(VaType virt, VaType size);
+
+    VaType GetVAStart() const {
+        return va_start;
+    }
+
+private:
+    /// The base VA of the allocator, no allocations will be below this
+    VaType va_start;
+
+    /**
+     * The end address for the initial linear allocation pass
+     * Once this reaches the AS limit the slower allocation path will be used
+     */
+    VaType current_linear_alloc_end;
 };
 } // namespace Common