Fixed incorrect accounting of memory budget.

Code by @medranSolus
Fixes #241
diff --git a/include/vk_mem_alloc.h b/include/vk_mem_alloc.h
index ce3279a..af82f3b 100644
--- a/include/vk_mem_alloc.h
+++ b/include/vk_mem_alloc.h
@@ -11889,7 +11889,7 @@
     VMA_ASSERT(m_MapCount == 0 && "Allocation was not unmapped before destruction.");

 

     // Check if owned string was freed.

-    VMA_ASSERT(m_pUserData == VMA_NULL);

+    VMA_ASSERT((IsUserDataString() && m_pUserData == VMA_NULL) || !IsUserDataString());

 }

 

 void VmaAllocation_T::InitBlockAllocation(

@@ -12326,14 +12326,8 @@
     if (res != VK_SUCCESS)

     {

         // Free all already created allocations.

-        const uint32_t heapIndex = m_hAllocator->MemoryTypeIndexToHeapIndex(m_MemoryTypeIndex);

         while (allocIndex--)

-        {

-            VmaAllocation_T* const alloc = pAllocations[allocIndex];

-            const VkDeviceSize allocSize = alloc->GetSize();

-            Free(alloc);

-            m_hAllocator->m_Budget.RemoveAllocation(heapIndex, allocSize);

-        }

+            Free(pAllocations[allocIndex]);

         memset(pAllocations, 0, sizeof(VmaAllocation) * allocationCount);

     }

 

@@ -12618,6 +12612,9 @@
         pBlockToDelete->Destroy(m_hAllocator);

         vma_delete(m_hAllocator, pBlockToDelete);

     }

+

+    m_hAllocator->m_Budget.RemoveAllocation(m_hAllocator->MemoryTypeIndexToHeapIndex(m_MemoryTypeIndex), hAllocation->GetSize());

+    m_hAllocator->m_AllocationObjectAllocator.Free(hAllocation);

 }

 

 VkDeviceSize VmaBlockVector::CalcMaxBlockSize() const

@@ -14467,7 +14464,6 @@
 

             FreeVulkanMemory(memTypeIndex, currAlloc->GetSize(), hMemory);

             m_Budget.RemoveAllocation(MemoryTypeIndexToHeapIndex(memTypeIndex), currAlloc->GetSize());

-            currAlloc->SetUserData(this, VMA_NULL);

             m_AllocationObjectAllocator.Free(currAlloc);

         }

 

@@ -14874,10 +14870,6 @@
             default:

                 VMA_ASSERT(0);

             }

-

-            m_Budget.RemoveAllocation(MemoryTypeIndexToHeapIndex(allocation->GetMemoryTypeIndex()), allocation->GetSize());

-            allocation->SetUserData(this, VMA_NULL);

-            m_AllocationObjectAllocator.Free(allocation);

         }

     }

 }

@@ -15478,6 +15470,9 @@
 

     FreeVulkanMemory(memTypeIndex, allocation->GetSize(), hMemory);

 

+    m_Budget.RemoveAllocation(MemoryTypeIndexToHeapIndex(allocation->GetMemoryTypeIndex()), allocation->GetSize());

+    m_AllocationObjectAllocator.Free(allocation);

+

     VMA_DEBUG_LOG("    Freed DedicatedMemory MemoryTypeIndex=%u", memTypeIndex);

 }

 

diff --git a/src/Tests.cpp b/src/Tests.cpp
index 9937691..302a762 100644
--- a/src/Tests.cpp
+++ b/src/Tests.cpp
@@ -7731,11 +7731,11 @@
     
     WriteMainTestResultHeader(file);
     PerformMainTests(file);
-    //PerformCustomMainTest(file);
+    PerformCustomMainTest(file);
 
     WritePoolTestResultHeader(file);
     PerformPoolTests(file);
-    //PerformCustomPoolTest(file);
+    PerformCustomPoolTest(file);
     
     fclose(file);
 #endif // #if defined(VMA_DEBUG_MARGIN) && VMA_DEBUG_MARGIN > 0