From d24d49b3a5b9ba002451f1cee06c284d4f833414 Mon Sep 17 00:00:00 2001
From: David Allemang <david.allemang@kitware.com>
Date: Wed, 3 Apr 2024 17:05:39 -0400
Subject: [PATCH] VkAllocator for vkAllocateMemory

---
 src/gfx.zig  | 70 ++++++++++++++++++++++++++++++----------------------
 src/main.zig |  8 ++++--
 2 files changed, 46 insertions(+), 32 deletions(-)

diff --git a/src/gfx.zig b/src/gfx.zig
index b780178..30085bb 100644
--- a/src/gfx.zig
+++ b/src/gfx.zig
@@ -296,11 +296,14 @@ pub fn uploadData(
     }, null);
     defer vkd.destroyBuffer(dev, staging_buffer, null);
 
+    const vally = VkAllocator.init(pdev, vki);
+
     const mem_reqs = vkd.getBufferMemoryRequirements(dev, staging_buffer);
-    const staging_memory = try allocate(pdev, vki, dev, vkd, mem_reqs, .{
-        .host_visible_bit = true,
-        .host_coherent_bit = true,
-    });
+    const staging_memory = try vally.alloc(dev, vkd, mem_reqs, .{ .host_visible_bit = true, .host_coherent_bit = true });
+    // const staging_memory = try allocate(pdev, vki, dev, vkd, mem_reqs, .{
+    //     .host_visible_bit = true,
+    //     .host_coherent_bit = true,
+    // });
     defer vkd.freeMemory(dev, staging_memory, null);
 
     try vkd.bindBufferMemory(dev, staging_buffer, staging_memory, 0);
@@ -362,36 +365,43 @@ pub fn copyBuffer(
     try vkd.queueWaitIdle(queue);
 }
 
-pub fn findMemoryTypeIndex(
-    pdev: vk.PhysicalDevice,
-    memory_type_bits: u32,
-    flags: vk.MemoryPropertyFlags,
-    vki: InstanceDispatch,
-) !u32 {
-    const mem_props = vki.getPhysicalDeviceMemoryProperties(pdev);
+pub const VkAllocator = struct {
+    memory_types: [vk.MAX_MEMORY_TYPES]vk.MemoryType,
+    memory_type_count: u32,
 
-    for (mem_props.memory_types[0..mem_props.memory_type_count], 0..) |mem_type, i| {
-        if (memory_type_bits & (@as(u32, 1) << @truncate(i)) != 0 and mem_type.property_flags.contains(flags)) {
-            return @truncate(i);
-        }
+    pub fn init(
+        pdev: vk.PhysicalDevice,
+        vki: InstanceDispatch,
+    ) VkAllocator {
+        const props = vki.getPhysicalDeviceMemoryProperties(pdev);
+
+        return VkAllocator{
+            .memory_types = props.memory_types,
+            .memory_type_count = props.memory_type_count,
+        };
     }
 
-    return error.NoSuitableMemoryType;
-}
+    pub fn alloc(
+        self: VkAllocator,
+        dev: vk.Device,
+        vkd: DeviceDispatch,
+        reqs: vk.MemoryRequirements,
+        flags: vk.MemoryPropertyFlags,
+    ) !vk.DeviceMemory {
+        const memory_type_bits = reqs.memory_type_bits;
 
-pub fn allocate(
-    pdev: vk.PhysicalDevice,
-    vki: InstanceDispatch,
-    dev: vk.Device,
-    vkd: DeviceDispatch,
-    requirements: vk.MemoryRequirements,
-    flags: vk.MemoryPropertyFlags,
-) !vk.DeviceMemory {
-    return try vkd.allocateMemory(dev, &.{
-        .allocation_size = requirements.size,
-        .memory_type_index = try findMemoryTypeIndex(pdev, requirements.memory_type_bits, flags, vki),
-    }, null);
-}
+        for (self.memory_types[0..self.memory_type_count], 0..) |mem_type, idx| {
+            if (memory_type_bits & (@as(u32, 1) << @truncate(idx)) != 0 and mem_type.property_flags.contains(flags)) {
+                return try vkd.allocateMemory(dev, &.{
+                    .allocation_size = reqs.size,
+                    .memory_type_index = @intCast(idx),
+                }, null);
+            }
+        }
+
+        return error.NoSuitableMemoryType;
+    }
+};
 
 pub const BaseDispatch = vk.BaseWrapper(.{
     .createInstance = true,
diff --git a/src/main.zig b/src/main.zig
index e07847b..7f5d3a1 100644
--- a/src/main.zig
+++ b/src/main.zig
@@ -245,6 +245,8 @@ pub fn main() !void {
     defer chain.deinit(ally);
     defer deinit_chain(chain, dev, vkd, pool);
 
+    const device_local = gfx.VkAllocator.init(pdev, vki);
+
     const pipeline_layout = try vkd.createPipelineLayout(dev, &.{
         .flags = .{},
         .set_layout_count = 0,
@@ -264,7 +266,8 @@ pub fn main() !void {
     }, null);
     defer vkd.destroyBuffer(dev, vertex_buffer, null);
     const vertex_mem_reqs = vkd.getBufferMemoryRequirements(dev, vertex_buffer);
-    const vertex_memory = try gfx.allocate(pdev, vki, dev, vkd, vertex_mem_reqs, .{ .device_local_bit = true });
+    // const vertex_memory = try gfx.allocate(pdev, vki, dev, vkd, vertex_mem_reqs, .{ .device_local_bit = true });
+    const vertex_memory = try device_local.alloc(dev, vkd, vertex_mem_reqs, .{ .device_local_bit = true });
     defer vkd.freeMemory(dev, vertex_memory, null);
     try vkd.bindBufferMemory(dev, vertex_buffer, vertex_memory, 0);
 
@@ -277,7 +280,8 @@ pub fn main() !void {
     }, null);
     defer vkd.destroyBuffer(dev, index_buffer, null);
     const index_mem_reqs = vkd.getBufferMemoryRequirements(dev, index_buffer);
-    const index_memory = try gfx.allocate(pdev, vki, dev, vkd, index_mem_reqs, .{ .device_local_bit = true });
+    // const index_memory = try gfx.allocate(pdev, vki, dev, vkd, index_mem_reqs, .{ .device_local_bit = true });
+    const index_memory = try device_local.alloc(dev, vkd, index_mem_reqs, .{ .device_local_bit = true });
     defer vkd.freeMemory(dev, index_memory, null);
     try vkd.bindBufferMemory(dev, index_buffer, index_memory, 0);