From d6c1e5a5a6229b52c20c643ff626a948d4850865 Mon Sep 17 00:00:00 2001
From: John Zulauf <jzulauf@lunarg.com>
Date: Thu, 22 Jun 2023 17:29:21 -0600
Subject: layers: New merge/split in buffer_address_map

---
 layers/containers/custom_containers.h  |  1 +
 layers/state_tracker/state_tracker.cpp | 43 +++++++++++++++++++++-------------
 layers/state_tracker/state_tracker.h   |  5 +++-
 3 files changed, 32 insertions(+), 17 deletions(-)

diff --git a/layers/containers/custom_containers.h b/layers/containers/custom_containers.h
index 5ea43f5e3..bf3ec84aa 100644
--- a/layers/containers/custom_containers.h
+++ b/layers/containers/custom_containers.h
@@ -357,6 +357,7 @@ class small_vector {
     inline const_iterator cend() const { return GetWorkingStore() + size_; }
     inline const_iterator end() const { return GetWorkingStore() + size_; }
     inline size_type size() const { return size_; }
+    auto capacity() const { return capacity_; }
 
     inline pointer data() { return GetWorkingStore(); }
     inline const_pointer data() const { return GetWorkingStore(); }
diff --git a/layers/state_tracker/state_tracker.cpp b/layers/state_tracker/state_tracker.cpp
index 89ef11cfa..c965b1d10 100644
--- a/layers/state_tracker/state_tracker.cpp
+++ b/layers/state_tracker/state_tracker.cpp
@@ -342,6 +342,29 @@ void ValidationStateTracker::PreCallRecordCmdBlitImage2(VkCommandBuffer commandB
                                 Get<IMAGE_STATE>(pBlitImageInfo->dstImage));
 }
 
+struct BufferAddressInfillUpdateOps {
+    using Map = typename ValidationStateTracker::BufferAddressRangeMap;
+    using Iterator = typename Map::iterator;
+    using Value = typename Map::value_type;
+    using Mapped = typename Map::mapped_type;
+    using Range = typename Map::key_type;
+    void infill(Map &map, const Iterator &pos, const Range &infill_range) const {
+        map.insert(pos, Value(infill_range, insert_value));
+    }
+    void update(const Iterator &pos) const {
+        auto &current_buffer_list = pos->second;
+        assert(!current_buffer_list.empty());
+        const auto buffer_found_it = std::find(current_buffer_list.begin(), current_buffer_list.end(), insert_value[0]);
+        if (buffer_found_it == current_buffer_list.end()) {
+            if (current_buffer_list.capacity() <= (current_buffer_list.size() + 1)) {
+                current_buffer_list.reserve(current_buffer_list.capacity() * 2);
+            }
+            current_buffer_list.emplace_back(insert_value[0]);
+        }
+    }
+    const Mapped &insert_value;
+};
+
 void ValidationStateTracker::PostCallRecordCreateBuffer(VkDevice device, const VkBufferCreateInfo *pCreateInfo,
                                                         const VkAllocationCallbacks *pAllocator, VkBuffer *pBuffer,
                                                         VkResult result) {
@@ -366,14 +389,8 @@ void ValidationStateTracker::PostCallRecordCreateBuffer(VkDevice device, const V
             buffer_state->deviceAddress = opaque_capture_address->opaqueCaptureAddress;
             const auto address_range = buffer_state->DeviceAddressRange();
 
-            buffer_address_map_.split_and_merge_insert(
-                {address_range, {buffer_state}}, [](auto &current_buffer_list, const auto &new_buffer) {
-                    assert(!current_buffer_list.empty());
-                    const auto buffer_found_it = std::find(current_buffer_list.begin(), current_buffer_list.end(), new_buffer[0]);
-                    if (buffer_found_it == current_buffer_list.end()) {
-                        current_buffer_list.emplace_back(new_buffer[0]);
-                    }
-                });
+            BufferAddressInfillUpdateOps ops{{buffer_state}};
+            sparse_container::infill_update_range(buffer_address_map_, address_range, ops);
         }
 
         const VkBufferUsageFlags descriptor_buffer_usages =
@@ -5707,14 +5724,8 @@ void ValidationStateTracker::RecordGetBufferDeviceAddress(const VkBufferDeviceAd
         buffer_state->deviceAddress = address;
         const auto address_range = buffer_state->DeviceAddressRange();
 
-        buffer_address_map_.split_and_merge_insert(
-            {address_range, {buffer_state}}, [](auto &current_buffer_list, const auto &new_buffer) {
-                assert(!current_buffer_list.empty());
-                const auto buffer_found_it = std::find(current_buffer_list.begin(), current_buffer_list.end(), new_buffer[0]);
-                if (buffer_found_it == current_buffer_list.end()) {
-                    current_buffer_list.emplace_back(new_buffer[0]);
-                }
-            });
+        BufferAddressInfillUpdateOps ops{{buffer_state}};
+        sparse_container::infill_update_range(buffer_address_map_, address_range, ops);
     }
 }
 
diff --git a/layers/state_tracker/state_tracker.h b/layers/state_tracker/state_tracker.h
index 662771341..588146b46 100644
--- a/layers/state_tracker/state_tracker.h
+++ b/layers/state_tracker/state_tracker.h
@@ -1596,6 +1596,9 @@ class ValidationStateTracker : public ValidationObject {
 
     mutable VideoProfileDesc::Cache video_profile_cache_;
 
+    using BufferAddressMapStore = small_vector<std::shared_ptr<BUFFER_STATE>, 1, size_t>;
+    using BufferAddressRangeMap = sparse_container::range_map<VkDeviceAddress, BufferAddressMapStore>;
+
   protected:
     // tracks which queue family index were used when creating the device for quick lookup
     vvl::unordered_set<uint32_t> queue_family_index_set;
@@ -1608,7 +1611,7 @@ class ValidationStateTracker : public ValidationObject {
     };
     std::vector<DeviceQueueInfo> device_queue_info_list;
     // If vkGetBufferDeviceAddress is called, keep track of buffer <-> address mapping.
-    sparse_container::range_map<VkDeviceAddress, small_vector<std::shared_ptr<BUFFER_STATE>, 1, size_t>> buffer_address_map_;
+    BufferAddressRangeMap buffer_address_map_;
     mutable std::shared_mutex buffer_address_lock_;
 
     vl_concurrent_unordered_map<uint64_t, VkFormatFeatureFlags2KHR> ahb_ext_formats_map;
-- 
cgit v1.2.3