2 files changed, 83 insertions, 3 deletions

diff --git a/src/window.cpp b/src/window.cpp
index 6e285b5..480175a 100644
--- a/src/window.cpp
+++ b/src/window.cpp
@@ -18,6 +18,12 @@
 #include <cstdint>
 #include <fstream>
 
+MemBucket *MemBkt_Vulkan = nullptr;
+struct pke_vkAllocData {
+	void *data;
+	std::size_t size;
+};
+DynArray<pke_vkAllocData> *vulkanAllocs = nullptr;
 
 #define NELEMS(x)  (sizeof(x) / sizeof((x)[0]))
 
@@ -234,7 +240,72 @@ unsigned int FindQueueFamilyIndex(VkPhysicalDevice device, short hasPresentSuppo
 	return 0xFFFFFFFF;
 }
 
+void PKE_vkFreeFunction(void *pUserData, void *pMemory) {
+	pke_vkAllocData *chunk = nullptr;
+	size_t index = -1;
+	size_t count = vulkanAllocs->Count();
+	for (long i = 0; i < count; ++i) {
+		if ((*vulkanAllocs)[i].data == pMemory) {
+			chunk = &(*vulkanAllocs)[i];
+			index = i;
+			break;
+		}
+	}
+	assert(chunk != nullptr);
+	Pke_Delete(chunk->data, chunk->size, MemBkt_Vulkan);
+	vulkanAllocs->Remove(index);
+}
+
+void *PKE_vkAllocateFunction(void *pUserData, size_t size, size_t alignment, VkSystemAllocationScope allocScope) {
+	size_t alignedSize = size + ((size % alignment) == 0 ? 0 : (alignment - (size % alignment)));
+	void *ptr = Pke_New(alignedSize, alignment, MemBkt_Vulkan);
+	vulkanAllocs->Push({
+		.data = ptr,
+		.size = alignedSize,
+	});
+	return ptr;
+}
+
+void *PKE_vkReallocationFunction(void *pUserData, void *pOriginal, size_t size, size_t alignment, VkSystemAllocationScope allocScope) {
+	pke_vkAllocData *chunk = nullptr;
+	size_t index = -1;
+	if (pOriginal) {
+		size_t count = vulkanAllocs->Count();
+		for (long i = 0; i < count; ++i) {
+			if ((*vulkanAllocs)[i].data == pOriginal) {
+				chunk = &(*vulkanAllocs)[i];
+				index = i;
+				break;
+			}
+		}
+	}
+	assert(chunk != nullptr);
+	if (chunk->size > size) {
+		return pOriginal;
+	}
+	void *newPtr = PKE_vkAllocateFunction(pUserData, size, alignment, allocScope);
+	memcpy(newPtr, chunk->data, chunk->size);
+	Pke_Delete(chunk->data, chunk->size, MemBkt_Vulkan);
+	vulkanAllocs->Remove(index);
+	return newPtr;
+}
+
+void PKE_vkInternalAllocationNotification(void *pUserData, size_t size, VkInternalAllocationType allocationType, VkSystemAllocationScope allocationScope) {
+	return;
+}
+
+void PKE_vkInternalFreeNotification(void *pUserData, size_t size, VkInternalAllocationType allocationType, VkSystemAllocationScope allocationScope) {
+	return;
+}
+
 void InitVulkan() {
+	vkAllocator = Pke_New<VkAllocationCallbacks>(MemBkt_Vulkan);
+	vkAllocator->pUserData = nullptr;
+	vkAllocator->pfnAllocation = PKE_vkAllocateFunction;
+	vkAllocator->pfnReallocation = PKE_vkReallocationFunction;
+	vkAllocator->pfnFree = PKE_vkFreeFunction;
+	vkAllocator->pfnInternalAllocation = PKE_vkInternalAllocationNotification;
+	vkAllocator->pfnInternalFree = PKE_vkInternalFreeNotification;
 
 	VkResult vkResult;
 	if (ENABLE_VALIDATION_LAYERS) {
@@ -328,7 +399,7 @@ void InitVulkan() {
 		Pke_Delete<VkExtensionProperties>(extensions, extensionCount);
 	}
 
-	vkResult = vkCreateInstance(&createInfo, nullptr, &vkInstance);
+	vkResult = vkCreateInstance(&createInfo, vkAllocator, &vkInstance);
 	if (vkResult != VK_SUCCESS) {
 		printf("Failed to create VkInstance! : %d\n", vkResult);
 		throw vkResult;
@@ -348,7 +419,7 @@ void InitVulkan() {
 		debugReportCallbackCreateInfo.pfnCallback = DebugReport;
 		debugReportCallbackCreateInfo.pUserData = nullptr;
 
-		vkResult = func(vkInstance, &debugReportCallbackCreateInfo, nullptr, &vkDebugReport);
+		vkResult = func(vkInstance, &debugReportCallbackCreateInfo, vkAllocator, &vkDebugReport);
 
 		if (vkResult != VK_SUCCESS) {
 			fprintf(stderr, "%s\n", "Failed to create debug report!");
@@ -356,7 +427,7 @@ void InitVulkan() {
 	}
 
 	// create surface
-	if (glfwCreateWindowSurface(vkInstance, window, nullptr, &vkSurfaceKHR) != VK_SUCCESS) {
+	if (glfwCreateWindowSurface(vkInstance, window, vkAllocator, &vkSurfaceKHR) != VK_SUCCESS) {
 		throw "Failed to create window surface";
 	}
 
@@ -2460,6 +2531,10 @@ void FramebufferResizeCallback(GLFWwindow *window, int width, int height) {
 
 void CreateWindow(PKEWindowProperties *wp) {
 	if (vkInstance != nullptr) return;
+	MemBkt_Vulkan = Pke_BeginTransientBucket();
+	vulkanAllocs = Pke_New<DynArray<pke_vkAllocData>>(MemBkt_Vulkan);
+	new (vulkanAllocs) DynArray<pke_vkAllocData>(MemBkt_Vulkan);
+	vulkanAllocs->Reserve(2048);
 	glfwInit();
 	glfwWindowHint(GLFW_CLIENT_API, GLFW_NO_API);
 	window = glfwCreateWindow(wp->width, wp->height, "Pikul", nullptr, nullptr);
@@ -2540,6 +2615,9 @@ void DestroyWindow() {
 	vkDestroyInstance(vkInstance, vkAllocator);
 	glfwDestroyWindow(window);
 	glfwTerminate();
+
+	vulkanAllocs->~DynArray();
+	Pke_EndTransientBucket(MemBkt_Vulkan);
 }
 
 VkShaderModule UploadShader(AssetHandle handle) {
diff --git a/src/window.hpp b/src/window.hpp
index 7f42766..83f7e68 100644
--- a/src/window.hpp
+++ b/src/window.hpp
@@ -20,6 +20,8 @@
 // TODO replace me with something more elegant
 const unsigned int MAX_FRAMES_IN_FLIGHT = 3;
 
+extern MemBucket *MemBkt_Vulkan;
+
 extern bool shouldRecreateSwapchain;
 extern GLFWwindow *window;
 extern GLFWmonitor *monitor;