layers: Clean up acceleration state objects

Remove unneeded data members and make many others const.

The NV and KHR state objects cannot be combined because the NV version
has a VkDeviceMemory bound to by vkBindAccelerationStructureMemoryNV(),
but the KHR version gets its memory from a VkBuffer included in
the VkAccelerationStructureCreateInfoKHR struct.

4 files changed, 72 insertions, 77 deletions

diff --git a/layers/core_validation.cpp b/layers/core_validation.cpp
index f04a7aeff..887ff1a04 100644
--- a/layers/core_validation.cpp
+++ b/layers/core_validation.cpp
@@ -213,7 +213,7 @@ bool CoreChecks::ValidateMemoryIsBoundToAccelerationStructure(const ACCELERATION
 // Check to see if memory was bound to this acceleration structure
 bool CoreChecks::ValidateMemoryIsBoundToAccelerationStructure(const ACCELERATION_STRUCTURE_STATE_KHR *as_state,
                                                               const char *api_name, const char *error_code) const {
-    return VerifyBoundMemoryIsValid(as_state->MemState(), as_state->acceleration_structure(), as_state->Handle(), api_name,
+    return VerifyBoundMemoryIsValid(as_state->buffer_state->MemState(), as_state->acceleration_structure(), as_state->Handle(), api_name,
                                     error_code);
 }
 
@@ -8188,33 +8188,32 @@ bool CoreChecks::ValidateBindAccelerationStructureMemory(VkDevice device,
     if (mem_info) {
         skip |= ValidateInsertAccelerationStructureMemoryRange(info.accelerationStructure, mem_info.get(), info.memoryOffset,
                                                                "vkBindAccelerationStructureMemoryNV()");
-        skip |= ValidateMemoryTypes(mem_info.get(), as_state->memory_requirements.memoryRequirements.memoryTypeBits,
+        skip |= ValidateMemoryTypes(mem_info.get(), as_state->memory_requirements.memoryTypeBits,
                                     "vkBindAccelerationStructureMemoryNV()",
                                     "VUID-VkBindAccelerationStructureMemoryInfoNV-memory-03622");
     }
 
     // Validate memory requirements alignment
-    if (SafeModulo(info.memoryOffset, as_state->memory_requirements.memoryRequirements.alignment) != 0) {
+    if (SafeModulo(info.memoryOffset, as_state->memory_requirements.alignment) != 0) {
         skip |= LogError(info.accelerationStructure, "VUID-VkBindAccelerationStructureMemoryInfoNV-memoryOffset-03623",
                          "vkBindAccelerationStructureMemoryNV(): memoryOffset  0x%" PRIxLEAST64
                          " must be an integer multiple of the alignment 0x%" PRIxLEAST64
                          " member of the VkMemoryRequirements structure returned from "
                          "a call to vkGetAccelerationStructureMemoryRequirementsNV with accelerationStructure and type of "
                          "VK_ACCELERATION_STRUCTURE_MEMORY_REQUIREMENTS_TYPE_OBJECT_NV",
-                         info.memoryOffset, as_state->memory_requirements.memoryRequirements.alignment);
+                         info.memoryOffset, as_state->memory_requirements.alignment);
     }
 
     if (mem_info) {
         // Validate memory requirements size
-        if (as_state->memory_requirements.memoryRequirements.size > (mem_info->alloc_info.allocationSize - info.memoryOffset)) {
+        if (as_state->memory_requirements.size > (mem_info->alloc_info.allocationSize - info.memoryOffset)) {
             skip |= LogError(info.accelerationStructure, "VUID-VkBindAccelerationStructureMemoryInfoNV-size-03624",
                              "vkBindAccelerationStructureMemoryNV(): The size 0x%" PRIxLEAST64
                              " member of the VkMemoryRequirements structure returned from a call to "
                              "vkGetAccelerationStructureMemoryRequirementsNV with accelerationStructure and type of "
                              "VK_ACCELERATION_STRUCTURE_MEMORY_REQUIREMENTS_TYPE_OBJECT_NV must be less than or equal to the size "
                              "of memory minus memoryOffset 0x%" PRIxLEAST64 ".",
-                             as_state->memory_requirements.memoryRequirements.size,
-                             mem_info->alloc_info.allocationSize - info.memoryOffset);
+                             as_state->memory_requirements.size, mem_info->alloc_info.allocationSize - info.memoryOffset);
         }
     }
 
@@ -8559,8 +8558,7 @@ bool CoreChecks::PreCallValidateCmdBuildAccelerationStructureNV(VkCommandBuffer
             // Use requirements fetched at create time
         }
         if (scratch_buffer_state != nullptr && dst_as_state != nullptr &&
-            dst_as_state->update_scratch_memory_requirements.memoryRequirements.size >
-                (scratch_buffer_state->createInfo.size - scratchOffset)) {
+            dst_as_state->update_scratch_memory_requirements.size > (scratch_buffer_state->createInfo.size - scratchOffset)) {
             skip |= LogError(commandBuffer, "VUID-vkCmdBuildAccelerationStructureNV-update-02492",
                              "vkCmdBuildAccelerationStructureNV(): If update is VK_TRUE, The size member of the "
                              "VkMemoryRequirements structure returned from a call to "
@@ -8579,8 +8577,7 @@ bool CoreChecks::PreCallValidateCmdBuildAccelerationStructureNV(VkCommandBuffer
             // Use requirements fetched at create time
         }
         if (scratch_buffer_state != nullptr && dst_as_state != nullptr &&
-            dst_as_state->build_scratch_memory_requirements.memoryRequirements.size >
-                (scratch_buffer_state->createInfo.size - scratchOffset)) {
+            dst_as_state->build_scratch_memory_requirements.size > (scratch_buffer_state->createInfo.size - scratchOffset)) {
             skip |= LogError(commandBuffer, "VUID-vkCmdBuildAccelerationStructureNV-update-02491",
                              "vkCmdBuildAccelerationStructureNV(): If update is VK_FALSE, The size member of the "
                              "VkMemoryRequirements structure returned from a call to "
diff --git a/layers/descriptor_sets.cpp b/layers/descriptor_sets.cpp
index 6e55afd94..c1d629098 100644
--- a/layers/descriptor_sets.cpp
+++ b/layers/descriptor_sets.cpp
@@ -1715,7 +1715,7 @@ bool CoreChecks::ValidateAccelerationDescriptor(const char *caller, const DrawDi
                                 report_data->FormatHandle(acc).c_str());
             }
         } else {
-            for (const auto &item: acc_node->GetBoundMemory()) {
+            for (const auto &item: acc_node->buffer_state->GetBoundMemory()) {
                 auto &mem_binding = item.second;
                 if (mem_binding.mem_state->Destroyed()) {
                     auto set = descriptor_set->GetSet();
diff --git a/layers/ray_tracing_state.h b/layers/ray_tracing_state.h
index 9b7be65bb..d0746474c 100644
--- a/layers/ray_tracing_state.h
+++ b/layers/ray_tracing_state.h
@@ -1,7 +1,7 @@
-/* Copyright (c) 2015-2021 The Khronos Group Inc.
- * Copyright (c) 2015-2021 Valve Corporation
- * Copyright (c) 2015-2021 LunarG, Inc.
- * Copyright (C) 2015-2021 Google Inc.
+/* Copyright (c) 2015-2022 The Khronos Group Inc.
+ * Copyright (c) 2015-2022 Valve Corporation
+ * Copyright (c) 2015-2022 LunarG, Inc.
+ * Copyright (C) 2015-2022 Google Inc.
  * Modifications Copyright (C) 2020 Advanced Micro Devices, Inc. All rights reserved.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
@@ -27,28 +27,19 @@
  */
 #pragma once
 #include "device_memory_state.h"
+#include "buffer_state.h"
+#include "layer_chassis_dispatch.h"
 
 class ACCELERATION_STRUCTURE_STATE : public BINDABLE {
   public:
-    safe_VkAccelerationStructureCreateInfoNV create_infoNV = {};
-    safe_VkAccelerationStructureInfoNV build_info;
-    bool memory_requirements_checked = false;
-    VkMemoryRequirements2 memory_requirements;
-    bool build_scratch_memory_requirements_checked = false;
-    VkMemoryRequirements2 build_scratch_memory_requirements;
-    bool update_scratch_memory_requirements_checked = false;
-    VkMemoryRequirements2 update_scratch_memory_requirements;
-    bool built = false;
-    uint64_t opaque_handle = 0;
-    const VkAllocationCallbacks *allocator = NULL;
-    ACCELERATION_STRUCTURE_STATE(VkAccelerationStructureNV as, const VkAccelerationStructureCreateInfoNV *ci)
+    ACCELERATION_STRUCTURE_STATE(VkDevice device, VkAccelerationStructureNV as, const VkAccelerationStructureCreateInfoNV *ci)
         : BINDABLE(as, kVulkanObjectTypeAccelerationStructureNV, false, false, 0),
           create_infoNV(ci),
-          memory_requirements{},
-          build_scratch_memory_requirements_checked{},
-          build_scratch_memory_requirements{},
-          update_scratch_memory_requirements_checked{},
-          update_scratch_memory_requirements{} {}
+          memory_requirements(GetMemReqs(device, as, VK_ACCELERATION_STRUCTURE_MEMORY_REQUIREMENTS_TYPE_OBJECT_NV)),
+          build_scratch_memory_requirements(
+              GetMemReqs(device, as, VK_ACCELERATION_STRUCTURE_MEMORY_REQUIREMENTS_TYPE_BUILD_SCRATCH_NV)),
+          update_scratch_memory_requirements(
+              GetMemReqs(device, as, VK_ACCELERATION_STRUCTURE_MEMORY_REQUIREMENTS_TYPE_UPDATE_SCRATCH_NV)) {}
     ACCELERATION_STRUCTURE_STATE(const ACCELERATION_STRUCTURE_STATE &rh_obj) = delete;
 
     VkAccelerationStructureNV acceleration_structure() const { return handle_.Cast<VkAccelerationStructureNV>(); }
@@ -57,37 +48,68 @@ class ACCELERATION_STRUCTURE_STATE : public BINDABLE {
         built = true;
         build_info.initialize(pInfo);
     };
-};
 
-class ACCELERATION_STRUCTURE_STATE_KHR : public BINDABLE {
-  public:
-    safe_VkAccelerationStructureCreateInfoKHR create_infoKHR = {};
-    safe_VkAccelerationStructureBuildGeometryInfoKHR build_info_khr;
+    const safe_VkAccelerationStructureCreateInfoNV create_infoNV = {};
+    safe_VkAccelerationStructureInfoNV build_info;
     bool memory_requirements_checked = false;
-    VkMemoryRequirements2 memory_requirements;
+    const VkMemoryRequirements memory_requirements;
     bool build_scratch_memory_requirements_checked = false;
-    VkMemoryRequirements2 build_scratch_memory_requirements;
+    const VkMemoryRequirements build_scratch_memory_requirements;
     bool update_scratch_memory_requirements_checked = false;
-    VkMemoryRequirements2 update_scratch_memory_requirements;
+    const VkMemoryRequirements update_scratch_memory_requirements;
     bool built = false;
     uint64_t opaque_handle = 0;
-    const VkAllocationCallbacks *allocator = NULL;
-    ACCELERATION_STRUCTURE_STATE_KHR(VkAccelerationStructureKHR as, const VkAccelerationStructureCreateInfoKHR *ci)
-        : BINDABLE(as, kVulkanObjectTypeAccelerationStructureKHR, false, false, 0),
-          create_infoKHR(ci),
-          memory_requirements{},
-          build_scratch_memory_requirements_checked{},
-          build_scratch_memory_requirements{},
-          update_scratch_memory_requirements_checked{},
-          update_scratch_memory_requirements{} {}
+
+  private:
+    static VkMemoryRequirements GetMemReqs(VkDevice device, VkAccelerationStructureNV as,
+                                           VkAccelerationStructureMemoryRequirementsTypeNV mem_type) {
+        auto req_info = LvlInitStruct<VkAccelerationStructureMemoryRequirementsInfoNV>();
+        req_info.type = mem_type;
+        req_info.accelerationStructure = as;
+        auto requirements = LvlInitStruct<VkMemoryRequirements2>();
+        DispatchGetAccelerationStructureMemoryRequirementsNV(device, &req_info, &requirements);
+        return requirements.memoryRequirements;
+    }
+};
+
+class ACCELERATION_STRUCTURE_STATE_KHR : public BASE_NODE {
+  public:
+    ACCELERATION_STRUCTURE_STATE_KHR(VkAccelerationStructureKHR as, const VkAccelerationStructureCreateInfoKHR *ci,
+                                     std::shared_ptr<BUFFER_STATE> &&buf_state)
+        : BASE_NODE(as, kVulkanObjectTypeAccelerationStructureKHR), create_infoKHR(ci), buffer_state(buf_state) {}
     ACCELERATION_STRUCTURE_STATE_KHR(const ACCELERATION_STRUCTURE_STATE_KHR &rh_obj) = delete;
 
     VkAccelerationStructureKHR acceleration_structure() const { return handle_.Cast<VkAccelerationStructureKHR>(); }
 
+    void LinkChildNodes() override {
+        // Connect child node(s), which cannot safely be done in the constructor.
+        buffer_state->AddParent(this);
+    }
+
+    void Destroy() override {
+        if (buffer_state) {
+            buffer_state->RemoveParent(this);
+            buffer_state = nullptr;
+        }
+        BASE_NODE::Destroy();
+    }
+
+    virtual ~ACCELERATION_STRUCTURE_STATE_KHR() {
+        if (!Destroyed()) {
+            Destroy();
+        }
+    }
+
     void Build(const VkAccelerationStructureBuildGeometryInfoKHR *pInfo) {
         built = true;
         build_info_khr.initialize(pInfo);
     };
+
+    const safe_VkAccelerationStructureCreateInfoKHR create_infoKHR = {};
+    safe_VkAccelerationStructureBuildGeometryInfoKHR build_info_khr;
+    bool built = false;
+    uint64_t opaque_handle = 0;
+    std::shared_ptr<BUFFER_STATE> buffer_state;
 };
 
 // Safe struct that spans NV and KHR VkRayTracingPipelineCreateInfo structures.
diff --git a/layers/state_tracker.cpp b/layers/state_tracker.cpp
index 4a2220562..dabe1209c 100644
--- a/layers/state_tracker.cpp
+++ b/layers/state_tracker.cpp
@@ -2397,27 +2397,7 @@ void ValidationStateTracker::PostCallRecordCreateAccelerationStructureNV(VkDevic
                                                                          VkAccelerationStructureNV *pAccelerationStructure,
                                                                          VkResult result) {
     if (VK_SUCCESS != result) return;
-    auto as_state = std::make_shared<ACCELERATION_STRUCTURE_STATE>(*pAccelerationStructure, pCreateInfo);
-
-    // Query the requirements in case the application doesn't (to avoid bind/validation time query)
-    auto as_memory_requirements_info = LvlInitStruct<VkAccelerationStructureMemoryRequirementsInfoNV>();
-    as_memory_requirements_info.type = VK_ACCELERATION_STRUCTURE_MEMORY_REQUIREMENTS_TYPE_OBJECT_NV;
-    as_memory_requirements_info.accelerationStructure = as_state->acceleration_structure();
-    DispatchGetAccelerationStructureMemoryRequirementsNV(device, &as_memory_requirements_info, &as_state->memory_requirements);
-
-    auto scratch_memory_req_info = LvlInitStruct<VkAccelerationStructureMemoryRequirementsInfoNV>();
-    scratch_memory_req_info.type = VK_ACCELERATION_STRUCTURE_MEMORY_REQUIREMENTS_TYPE_BUILD_SCRATCH_NV;
-    scratch_memory_req_info.accelerationStructure = as_state->acceleration_structure();
-    DispatchGetAccelerationStructureMemoryRequirementsNV(device, &scratch_memory_req_info,
-                                                         &as_state->build_scratch_memory_requirements);
-
-    auto update_memory_req_info = LvlInitStruct<VkAccelerationStructureMemoryRequirementsInfoNV>();
-    update_memory_req_info.type = VK_ACCELERATION_STRUCTURE_MEMORY_REQUIREMENTS_TYPE_UPDATE_SCRATCH_NV;
-    update_memory_req_info.accelerationStructure = as_state->acceleration_structure();
-    DispatchGetAccelerationStructureMemoryRequirementsNV(device, &update_memory_req_info,
-                                                         &as_state->update_scratch_memory_requirements);
-    as_state->allocator = pAllocator;
-    Add(std::move(as_state));
+    Add(std::make_shared<ACCELERATION_STRUCTURE_STATE>(device, *pAccelerationStructure, pCreateInfo));
 }
 
 void ValidationStateTracker::PostCallRecordCreateAccelerationStructureKHR(VkDevice device,
@@ -2426,9 +2406,8 @@ void ValidationStateTracker::PostCallRecordCreateAccelerationStructureKHR(VkDevi
                                                                           VkAccelerationStructureKHR *pAccelerationStructure,
                                                                           VkResult result) {
     if (VK_SUCCESS != result) return;
-    auto as_state = std::make_shared<ACCELERATION_STRUCTURE_STATE_KHR>(*pAccelerationStructure, pCreateInfo);
-    as_state->allocator = pAllocator;
-    Add(std::move(as_state));
+    auto buffer_state = Get<BUFFER_STATE>(pCreateInfo->buffer);
+    Add(std::make_shared<ACCELERATION_STRUCTURE_STATE_KHR>(*pAccelerationStructure, pCreateInfo, std::move(buffer_state)));
 }
 
 void ValidationStateTracker::PostCallRecordBuildAccelerationStructuresKHR(
@@ -2500,13 +2479,10 @@ void ValidationStateTracker::PostCallRecordGetAccelerationStructureMemoryRequire
     auto as_state = Get<ACCELERATION_STRUCTURE_STATE>(pInfo->accelerationStructure);
     if (as_state != nullptr) {
         if (pInfo->type == VK_ACCELERATION_STRUCTURE_MEMORY_REQUIREMENTS_TYPE_OBJECT_NV) {
-            as_state->memory_requirements = *pMemoryRequirements;
             as_state->memory_requirements_checked = true;
         } else if (pInfo->type == VK_ACCELERATION_STRUCTURE_MEMORY_REQUIREMENTS_TYPE_BUILD_SCRATCH_NV) {
-            as_state->build_scratch_memory_requirements = *pMemoryRequirements;
             as_state->build_scratch_memory_requirements_checked = true;
         } else if (pInfo->type == VK_ACCELERATION_STRUCTURE_MEMORY_REQUIREMENTS_TYPE_UPDATE_SCRATCH_NV) {
-            as_state->update_scratch_memory_requirements = *pMemoryRequirements;
             as_state->update_scratch_memory_requirements_checked = true;
         }
     }