From 4b3109a296b08c6b5f7204b0428700b414d11f70 Mon Sep 17 00:00:00 2001
From: Bryan Bernhart <bryan.bernhart@intel.com>
Date: Wed, 20 Apr 2022 10:27:32 -0700
Subject: [PATCH] Expose GetMemory directly to ResourceAllocation.

In a future change, UpdateResidency will be depreciated in favor of adding to the set directly.
---
 src/gpgmm/d3d12/ResourceAllocationD3D12.cpp   | 25 ++++++++-----------
 src/gpgmm/d3d12/ResourceAllocationD3D12.h     |  3 +++
 .../end2end/D3D12ResourceAllocatorTests.cpp   |  2 +-
 3 files changed, 14 insertions(+), 16 deletions(-)

diff --git a/src/gpgmm/d3d12/ResourceAllocationD3D12.cpp b/src/gpgmm/d3d12/ResourceAllocationD3D12.cpp
index 8be32b12e..b32905cf6 100644
--- a/src/gpgmm/d3d12/ResourceAllocationD3D12.cpp
+++ b/src/gpgmm/d3d12/ResourceAllocationD3D12.cpp
@@ -93,11 +93,8 @@ namespace gpgmm { namespace d3d12 {
             return E_INVALIDARG;
         }
 
-        Heap* resourceHeap = ToBackend(GetMemory());
-        ASSERT(resourceHeap != nullptr);
-
         if (mResidencyManager != nullptr) {
-            ReturnIfFailed(mResidencyManager->LockHeap(resourceHeap));
+            ReturnIfFailed(mResidencyManager->LockHeap(GetMemory()));
         }
 
         // Range coordinates are always subresource-relative so the range should only be
@@ -126,11 +123,8 @@ namespace gpgmm { namespace d3d12 {
         // subresource-relative coordinates.
         ASSERT(subresource == 0 || GetMethod() != AllocationMethod::kSubAllocatedWithin);
 
-        Heap* resourceHeap = ToBackend(GetMemory());
-        ASSERT(resourceHeap != nullptr);
-
         if (mResidencyManager != nullptr) {
-            mResidencyManager->UnlockHeap(resourceHeap);
+            mResidencyManager->UnlockHeap(GetMemory());
         }
 
         D3D12_RANGE newWrittenRange{};
@@ -145,13 +139,13 @@ namespace gpgmm { namespace d3d12 {
     }
 
     HRESULT ResourceAllocation::UpdateResidency(ResidencySet* residencySet) const {
-        Heap* resourceHeap = ToBackend(GetMemory());
+        Heap* resourceHeap = GetMemory();
         ASSERT(resourceHeap != nullptr);
         return resourceHeap->UpdateResidency(residencySet);
     }
 
     bool ResourceAllocation::IsResident() const {
-        const Heap* resourceHeap = ToBackend(GetMemory());
+        const Heap* resourceHeap = GetMemory();
         ASSERT(resourceHeap != nullptr);
         return resourceHeap->IsResident();
     }
@@ -166,15 +160,16 @@ namespace gpgmm { namespace d3d12 {
     }
 
     RESOURCE_ALLOCATION_INFO ResourceAllocation::GetInfo() const {
-        Heap* resourceHeap = ToBackend(GetMemory());
-        ASSERT(resourceHeap != nullptr);
-
-        return {GetSize(),   GetOffset(),  mOffsetFromResource,
-                GetMethod(), resourceHeap, mResource.Get()};
+        return {GetSize(),   GetOffset(), mOffsetFromResource,
+                GetMethod(), GetMemory(), mResource.Get()};
     }
 
     const char* ResourceAllocation::GetTypename() const {
         return "GPUMemoryAllocation";
     }
 
+    Heap* ResourceAllocation::GetMemory() const {
+        return ToBackend(MemoryAllocation::GetMemory());
+    }
+
 }}  // namespace gpgmm::d3d12
diff --git a/src/gpgmm/d3d12/ResourceAllocationD3D12.h b/src/gpgmm/d3d12/ResourceAllocationD3D12.h
index 16a8d6d8e..0add9cedd 100644
--- a/src/gpgmm/d3d12/ResourceAllocationD3D12.h
+++ b/src/gpgmm/d3d12/ResourceAllocationD3D12.h
@@ -91,6 +91,9 @@ namespace gpgmm { namespace d3d12 {
 
         const char* GetTypename() const;
 
+        // Returns the heap assigned to this resource allocation.
+        Heap* GetMemory() const;
+
       private:
         void DeleteThis() override;
 
diff --git a/src/tests/end2end/D3D12ResourceAllocatorTests.cpp b/src/tests/end2end/D3D12ResourceAllocatorTests.cpp
index 4603918a6..4763c0359 100644
--- a/src/tests/end2end/D3D12ResourceAllocatorTests.cpp
+++ b/src/tests/end2end/D3D12ResourceAllocatorTests.cpp
@@ -317,7 +317,7 @@ TEST_F(D3D12ResourceAllocatorTests, CreateBufferAlwaysCommitted) {
     EXPECT_EQ(allocation->GetSize(), kDefaultPreferredResourceHeapSize);
 
     // Commmitted resources cannot be backed by a D3D12 heap.
-    Heap* resourceHeap = ToBackend(allocation->GetMemory());
+    Heap* resourceHeap = allocation->GetMemory();
     ASSERT_NE(resourceHeap, nullptr);
     ASSERT_EQ(resourceHeap->GetHeap(), nullptr);