[vulkan] Promote vkAllocate/FreeMemory to fully virtual entry points

bug: 111137294
bug: 121420031
Change-Id: Id85c246a3b853d035934cb1803be8fe825777a73
diff --git a/system/vulkan/func_table.cpp b/system/vulkan/func_table.cpp
index 724190e..44faf95 100644
--- a/system/vulkan/func_table.cpp
+++ b/system/vulkan/func_table.cpp
@@ -27,6 +27,7 @@
 
 #include "VkEncoder.h"
 #include "HostConnection.h"
+#include "ResourceTracker.h"
 
 #include "goldfish_vk_private_defs.h"
 
@@ -239,7 +240,8 @@
 {
     auto vkEnc = HostConnection::get()->vkEncoder();
     VkResult vkAllocateMemory_VkResult_return = (VkResult)0;
-    vkAllocateMemory_VkResult_return = vkEnc->vkAllocateMemory(device, pAllocateInfo, pAllocator, pMemory);
+    auto resources = ResourceTracker::get();
+    vkAllocateMemory_VkResult_return = resources->on_vkAllocateMemory(vkEnc, VK_SUCCESS, device, pAllocateInfo, pAllocator, pMemory);
     return vkAllocateMemory_VkResult_return;
 }
 static void entry_vkFreeMemory(
@@ -248,7 +250,8 @@
     const VkAllocationCallbacks* pAllocator)
 {
     auto vkEnc = HostConnection::get()->vkEncoder();
-    vkEnc->vkFreeMemory(device, memory, pAllocator);
+    auto resources = ResourceTracker::get();
+    resources->on_vkFreeMemory(vkEnc, device, memory, pAllocator);
 }
 static VkResult entry_vkMapMemory(
     VkDevice device,
diff --git a/system/vulkan_enc/ResourceTracker.cpp b/system/vulkan_enc/ResourceTracker.cpp
index a3ff456..aac0fc6 100644
--- a/system/vulkan_enc/ResourceTracker.cpp
+++ b/system/vulkan_enc/ResourceTracker.cpp
@@ -425,11 +425,18 @@
         VkResult input_result,
         VkDevice device,
         const VkMemoryAllocateInfo* pAllocateInfo,
-        const VkAllocationCallbacks*,
+        const VkAllocationCallbacks* pAllocator,
         VkDeviceMemory* pMemory) {
 
         if (input_result != VK_SUCCESS) return input_result;
 
+        VkEncoder* enc = (VkEncoder*)context;
+
+        input_result =
+            enc->vkAllocateMemory(device, pAllocateInfo, pAllocator, pMemory);
+        
+        if (input_result != VK_SUCCESS) return input_result;
+
         VkDeviceSize allocationSize = pAllocateInfo->allocationSize;
         VkDeviceSize mappedSize = getNonCoherentExtendedSize(device, allocationSize);
         uint8_t* mappedPtr = nullptr;
@@ -452,7 +459,6 @@
             hostVisible && directMappingSupported;
 
         if (doDirectMap) {
-            VkEncoder* enc = (VkEncoder*)context;
 
             uint64_t directMappedAddr = 0;
 
@@ -484,6 +490,16 @@
         return input_result;
     }
 
+    void on_vkFreeMemory(
+        void* context,
+        VkDevice device,
+        VkDeviceMemory memory,
+        const VkAllocationCallbacks* pAllocateInfo) {
+
+        VkEncoder* enc = (VkEncoder*)context;
+        enc->vkFreeMemory(device, memory, pAllocateInfo);
+    }
+
     VkResult on_vkMapMemory(
         void*,
         VkResult host_result,
@@ -791,6 +807,15 @@
         context, input_result, device, pAllocateInfo, pAllocator, pMemory);
 }
 
+void ResourceTracker::on_vkFreeMemory(
+    void* context,
+    VkDevice device,
+    VkDeviceMemory memory,
+    const VkAllocationCallbacks* pAllocator) {
+    return mImpl->on_vkFreeMemory(
+        context, device, memory, pAllocator);
+}
+
 VkResult ResourceTracker::on_vkMapMemory(
     void* context,
     VkResult input_result,
diff --git a/system/vulkan_enc/ResourceTracker.h b/system/vulkan_enc/ResourceTracker.h
index d4299ea..a36a5d5 100644
--- a/system/vulkan_enc/ResourceTracker.h
+++ b/system/vulkan_enc/ResourceTracker.h
@@ -75,6 +75,11 @@
         const VkMemoryAllocateInfo* pAllocateInfo,
         const VkAllocationCallbacks* pAllocator,
         VkDeviceMemory* pMemory);
+    void on_vkFreeMemory(
+        void* context,
+        VkDevice device,
+        VkDeviceMemory memory,
+        const VkAllocationCallbacks* pAllocator);
 
     VkResult on_vkMapMemory(
         void* context,
diff --git a/system/vulkan_enc/VkEncoder.cpp b/system/vulkan_enc/VkEncoder.cpp
index 756c329..3822de8 100644
--- a/system/vulkan_enc/VkEncoder.cpp
+++ b/system/vulkan_enc/VkEncoder.cpp
@@ -1414,7 +1414,6 @@
     countingStream->clearPool();
     stream->clearPool();
     pool->freeAll();
-    mImpl->resources()->on_vkAllocateMemory(this, vkAllocateMemory_VkResult_return, device, pAllocateInfo, pAllocator, pMemory);
     return vkAllocateMemory_VkResult_return;
 }