/*
 * Copyright © 2018 Valve Corporation
 *
 * Permission is hereby granted, free of charge, to any person obtaining a
 * copy of this software and associated documentation files (the "Software"),
 * to deal in the Software without restriction, including without limitation
 * the rights to use, copy, modify, merge, publish, distribute, sublicense,
 * and/or sell copies of the Software, and to permit persons to whom the
 * Software is furnished to do so, subject to the following conditions:
 *
 * The above copyright notice and this permission notice (including the next
 * paragraph) shall be included in all copies or substantial portions of the
 * Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL
 * THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
 * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS
 * IN THE SOFTWARE.
 */

#include <vulkan/vulkan.h>
#include <vulkan/vk_layer.h>
#include "vk_layer_dispatch_table.h"

#include <sys/stat.h>
#include <fcntl.h>
#include <unistd.h>
#include <string.h>

#include <mutex>
#include <unordered_map>
#include <atomic>

#include "serialize.h"

#undef VK_LAYER_EXPORT
#define VK_LAYER_EXPORT extern "C"

#include "blob.h"

// global lock
std::mutex global_lock;
typedef std::lock_guard<std::mutex> lock_guard_t;
typedef std::unique_lock<std::mutex> unique_lock_t;

// atomic counter for pipeline tests
std::atomic<unsigned> pipeline_count;

// use the loader's dispatch table pointer as a key for dispatch map lookups
template<typename DispatchableType>
void *GetKey(DispatchableType inst)
{
  return *(void **)inst;
}

// instance and device dispatch tables
std::unordered_map<void *, VkLayerInstanceDispatchTable> instance_dispatch;
std::unordered_map<void *, VkLayerDispatchTable> device_dispatch;

// create info objects
std::unordered_map<VkSampler, VkSamplerCreateInfo *> samplers;
std::unordered_map<VkShaderModule, VkShaderModuleCreateInfo *> shader_modules;
std::unordered_map<VkDescriptorSetLayout, VkDescriptorSetLayoutCreateInfo *> descriptor_set_layouts;
std::unordered_map<VkPipelineLayout, VkPipelineLayoutCreateInfo *> pipeline_layouts;
std::unordered_map<VkRenderPass, VkRenderPassCreateInfo *> render_passes;

// instance chain
VK_LAYER_EXPORT VkResult VKAPI_CALL
vkpipeline_db_CreateInstance(
    const VkInstanceCreateInfo*                 pCreateInfo,
    const VkAllocationCallbacks*                pAllocator,
    VkInstance*                                 pInstance)
{
    VkResult result;

    VkLayerInstanceCreateInfo *layerCreateInfo =
        (VkLayerInstanceCreateInfo *)pCreateInfo->pNext;

    while (layerCreateInfo &&
           (layerCreateInfo->sType != VK_STRUCTURE_TYPE_LOADER_INSTANCE_CREATE_INFO ||
            layerCreateInfo->function != VK_LAYER_LINK_INFO)) {
        layerCreateInfo = (VkLayerInstanceCreateInfo *)layerCreateInfo->pNext;
    }

    if (!layerCreateInfo)
        return VK_ERROR_INITIALIZATION_FAILED;

    PFN_vkGetInstanceProcAddr getInstanceProcAddr =
        layerCreateInfo->u.pLayerInfo->pfnNextGetInstanceProcAddr;

    layerCreateInfo->u.pLayerInfo = layerCreateInfo->u.pLayerInfo->pNext;

    PFN_vkCreateInstance createInstance = (PFN_vkCreateInstance)
        getInstanceProcAddr(VK_NULL_HANDLE, "vkCreateInstance");

    result = createInstance(pCreateInfo, pAllocator, pInstance);
    if (result != VK_SUCCESS)
        return result;

    // Create the dispatch instance table.
    VkLayerInstanceDispatchTable dispatchTable;
    dispatchTable.GetInstanceProcAddr = (PFN_vkGetInstanceProcAddr)
        getInstanceProcAddr(*pInstance, "vkGetInstanceProcAddr");
    dispatchTable.DestroyInstance = (PFN_vkDestroyInstance)
        getInstanceProcAddr(*pInstance, "vkDestroyInstance");
    dispatchTable.EnumerateDeviceExtensionProperties = (PFN_vkEnumerateDeviceExtensionProperties)
        getInstanceProcAddr(*pInstance, "vkEnumerateDeviceExtensionProperties");

    lock_guard_t l(global_lock);
    instance_dispatch[GetKey(*pInstance)] = dispatchTable;

    return VK_SUCCESS;
}

VK_LAYER_EXPORT void VKAPI_CALL
vkpipeline_db_DestroyInstance(
    VkInstance                                  instance,
    const VkAllocationCallbacks                 *pAllocator)
{
    lock_guard_t l(global_lock);
    instance_dispatch[GetKey(instance)].DestroyInstance(instance, pAllocator);
    instance_dispatch.erase(GetKey(instance));
}

VK_LAYER_EXPORT VkResult VKAPI_CALL
vkpipeline_db_EnumerateInstanceLayerProperties(
    uint32_t                                    *pPropertyCount,
    VkLayerProperties                           *pProperties)
{
    if (pPropertyCount)
        *pPropertyCount = 1;

    if (pProperties) {
        strcpy(pProperties->layerName, "VK_LAYER_vkpipeline_db");
        strcpy(pProperties->description, "vkpipeline-db capture layer");
        pProperties->implementationVersion = 1;
        pProperties->specVersion = VK_MAKE_VERSION(1, 0, VK_HEADER_VERSION);
    }

    return VK_SUCCESS;
}

VK_LAYER_EXPORT VkResult VKAPI_CALL
vkpipeline_db_EnumerateInstanceExtensionProperties(
    const char                                  *pLayerName,
    uint32_t                                    *pPropertyCount,
    VkExtensionProperties                       *pProperties)
{
    if (pLayerName == NULL || strcmp(pLayerName, "VK_LAYER_vkpipeline_db"))
        return VK_ERROR_LAYER_NOT_PRESENT;

    if (pPropertyCount)
        *pPropertyCount = 0;

    return VK_SUCCESS;
}

// device chain
VK_LAYER_EXPORT VkResult VKAPI_CALL
vkpipeline_db_CreateDevice(
    VkPhysicalDevice                            physicalDevice,
    const VkDeviceCreateInfo*                   pCreateInfo,
    const VkAllocationCallbacks*                pAllocator,
    VkDevice*                                   pDevice)
{
    VkResult result;

    VkLayerDeviceCreateInfo *layerCreateInfo =
        (VkLayerDeviceCreateInfo *)pCreateInfo->pNext;

    while (layerCreateInfo &&
           (layerCreateInfo->sType != VK_STRUCTURE_TYPE_LOADER_DEVICE_CREATE_INFO ||
            layerCreateInfo->function != VK_LAYER_LINK_INFO)) {
        layerCreateInfo = (VkLayerDeviceCreateInfo *)layerCreateInfo->pNext;
    }

    if (!layerCreateInfo)
        return VK_ERROR_INITIALIZATION_FAILED;
  
    PFN_vkGetInstanceProcAddr getInstanceProcAddr =
        layerCreateInfo->u.pLayerInfo->pfnNextGetInstanceProcAddr;
    PFN_vkGetDeviceProcAddr getDeviceProcAddr =
        layerCreateInfo->u.pLayerInfo->pfnNextGetDeviceProcAddr;

    layerCreateInfo->u.pLayerInfo = layerCreateInfo->u.pLayerInfo->pNext;

    PFN_vkCreateDevice createDevice = (PFN_vkCreateDevice)
        getInstanceProcAddr(VK_NULL_HANDLE, "vkCreateDevice");

    result = createDevice(physicalDevice, pCreateInfo, pAllocator, pDevice);
    if (result != VK_SUCCESS)
        return result;

    // create the device dispatch table
    VkLayerDispatchTable dispatchTable;
    dispatchTable.GetDeviceProcAddr = (PFN_vkGetDeviceProcAddr)
        getDeviceProcAddr(*pDevice, "vkGetDeviceProcAddr");
    dispatchTable.DestroyDevice = (PFN_vkDestroyDevice)
        getDeviceProcAddr(*pDevice, "vkDestroyDevice");
    dispatchTable.CreateSampler =
        (PFN_vkCreateSampler)getDeviceProcAddr(*pDevice, "vkCreateSampler");
    dispatchTable.DestroySampler =
        (PFN_vkDestroySampler)getDeviceProcAddr(*pDevice, "vkDestroySampler");
    dispatchTable.CreateDescriptorSetLayout =
        (PFN_vkCreateDescriptorSetLayout)getDeviceProcAddr(*pDevice, "vkCreateDescriptorSetLayout");
    dispatchTable.DestroyDescriptorSetLayout =
        (PFN_vkDestroyDescriptorSetLayout)getDeviceProcAddr(*pDevice, "vkDestroyDescriptorSetLayout");
    dispatchTable.CreatePipelineLayout =
        (PFN_vkCreatePipelineLayout)getDeviceProcAddr(*pDevice, "vkCreatePipelineLayout");
    dispatchTable.DestroyPipelineLayout =
        (PFN_vkDestroyPipelineLayout)getDeviceProcAddr(*pDevice, "vkDestroyPipelineLayout");
    dispatchTable.CreateShaderModule =
        (PFN_vkCreateShaderModule)getDeviceProcAddr(*pDevice, "vkCreateShaderModule");
    dispatchTable.DestroyShaderModule =
        (PFN_vkDestroyShaderModule)getDeviceProcAddr(*pDevice, "vkDestroyShaderModule");
    dispatchTable.CreateRenderPass =
        (PFN_vkCreateRenderPass)getDeviceProcAddr(*pDevice, "vkCreateRenderPass");
    dispatchTable.DestroyRenderPass =
        (PFN_vkDestroyRenderPass)getDeviceProcAddr(*pDevice, "vkDestroyRenderPass");
    dispatchTable.CreateGraphicsPipelines =
        (PFN_vkCreateGraphicsPipelines)getDeviceProcAddr(*pDevice, "vkCreateGraphicsPipelines");
    dispatchTable.CreateComputePipelines =
        (PFN_vkCreateComputePipelines)getDeviceProcAddr(*pDevice, "vkCreateComputePipelines");

    lock_guard_t l(global_lock);
    device_dispatch[GetKey(*pDevice)] = dispatchTable;

    return VK_SUCCESS;
}

VK_LAYER_EXPORT void VKAPI_CALL
vkpipeline_db_DestroyDevice(
    VkDevice                                    device,
    const VkAllocationCallbacks                 *pAllocator)
{
    lock_guard_t l(global_lock);
    device_dispatch[GetKey(device)].DestroyDevice(device, pAllocator);
    device_dispatch.erase(GetKey(device));
}

VK_LAYER_EXPORT VkResult VKAPI_CALL
vkpipeline_db_EnumerateDeviceLayerProperties(
    VkPhysicalDevice                            physicalDevice,
    uint32_t                                    *pPropertyCount,
    VkLayerProperties                           *pProperties)
{
    return vkpipeline_db_EnumerateInstanceLayerProperties(pPropertyCount,
                                                        pProperties);
}

VK_LAYER_EXPORT VkResult VKAPI_CALL
vkpipeline_db_EnumerateDeviceExtensionProperties(
    VkPhysicalDevice                            physicalDevice,
    const char                                  *pLayerName,
    uint32_t                                    *pPropertyCount,
    VkExtensionProperties                       *pProperties)
{
    if (pLayerName == NULL || strcmp(pLayerName, "VK_LAYER_vkpipeline_db")) {
        if (physicalDevice == VK_NULL_HANDLE)
            return VK_SUCCESS;

        lock_guard_t l(global_lock);
        return instance_dispatch[GetKey(physicalDevice)].EnumerateDeviceExtensionProperties(physicalDevice, pLayerName, pPropertyCount, pProperties);
    }

    if (pPropertyCount)
        *pPropertyCount = 0;

    return VK_SUCCESS;
}

// sampler
static const VkSamplerCreateInfo *
get_sampler_info(VkSampler sampler)
{
    unique_lock_t lock(global_lock);
    auto it = samplers.find(sampler);
    lock.unlock();

    if (it == samplers.end())
        return nullptr;
    return it->second;
}

static void
destroy_sampler_info(VkSampler sampler)
{
    unique_lock_t lock(global_lock);
    auto it = samplers.find(sampler);
    if (it != samplers.end())
        free(it->second);
    samplers.erase(sampler);
    lock.unlock();
}

static void
record_sampler_info(VkSampler *pSampler, const VkSamplerCreateInfo *pCreateInfo)
{
    VkSamplerCreateInfo *pSamplerInfo;

    pSamplerInfo = (VkSamplerCreateInfo *)malloc(sizeof(*pSamplerInfo));

    memcpy(pSamplerInfo, pCreateInfo, sizeof(*pSamplerInfo));

    lock_guard_t lock(global_lock);
    samplers[*pSampler] = pSamplerInfo;
}

VK_LAYER_EXPORT VkResult VKAPI_CALL
vkpipeline_db_CreateSampler(
    VkDevice                                    device,
    const VkSamplerCreateInfo*                  pCreateInfo,
    const VkAllocationCallbacks*                pAllocator,
    VkSampler*                                  pSampler)
{
    VkResult result;

    result = device_dispatch[GetKey(device)].CreateSampler(device,
                                                           pCreateInfo,
                                                           pAllocator,
                                                           pSampler);
    if (result == VK_SUCCESS)
        record_sampler_info(pSampler, pCreateInfo);
    return result;
}

VK_LAYER_EXPORT void VKAPI_CALL
vkpipeline_db_DestroySampler(
    VkDevice                                    device,
    VkSampler                                   sampler,
    const VkAllocationCallbacks*                pAllocator)
{
    destroy_sampler_info(sampler);
    device_dispatch[GetKey(device)].DestroySampler(device, sampler, pAllocator);
}

// descriptor set layout
static const VkDescriptorSetLayoutCreateInfo *
get_descriptor_set_layout_info(VkDescriptorSetLayout descriptorSetLayout)
{
    unique_lock_t lock(global_lock);
    auto it = descriptor_set_layouts.find(descriptorSetLayout);
    lock.unlock();

    if (it == descriptor_set_layouts.end())
        return nullptr;
    return it->second;
}

static void
destroy_descriptor_set_layout_info(VkDescriptorSetLayout descriptorSetLayout)
{
    unique_lock_t lock(global_lock);
    auto it = descriptor_set_layouts.find(descriptorSetLayout);
    if (it != descriptor_set_layouts.end()) {
        free((void *)it->second->pBindings);
        free(it->second);
    }
    descriptor_set_layouts.erase(descriptorSetLayout);
    lock.unlock();
}

static void
record_descriptor_set_layout_info(VkDescriptorSetLayout *pSetLayout,
                                  const VkDescriptorSetLayoutCreateInfo *pCreateInfo)
{
    VkDescriptorSetLayoutCreateInfo *pSetLayoutInfo;

    pSetLayoutInfo = (VkDescriptorSetLayoutCreateInfo *)
        malloc(sizeof(*pSetLayoutInfo));

    memcpy(pSetLayoutInfo, pCreateInfo, sizeof(*pSetLayoutInfo));

    if (pCreateInfo->bindingCount) {
        VkDescriptorSetLayoutBinding *pBindings;

        pBindings = (VkDescriptorSetLayoutBinding *)
            malloc(sizeof(*pBindings) * pCreateInfo->bindingCount);

        memcpy(pBindings, pCreateInfo->pBindings,
               sizeof(*pBindings) * pCreateInfo->bindingCount);
        pSetLayoutInfo->pBindings = pBindings;
    }

    lock_guard_t lock(global_lock);
    descriptor_set_layouts[*pSetLayout] = pSetLayoutInfo;
}

VK_LAYER_EXPORT VkResult VKAPI_CALL
vkpipeline_db_CreateDescriptorSetLayout(
    VkDevice                                    device,
    const VkDescriptorSetLayoutCreateInfo*      pCreateInfo,
    const VkAllocationCallbacks*                pAllocator,
    VkDescriptorSetLayout*                      pSetLayout)
{
    VkResult result;

    result = device_dispatch[GetKey(device)].CreateDescriptorSetLayout(device,
                                                                       pCreateInfo,
                                                                       pAllocator,
                                                                       pSetLayout);
    if (result == VK_SUCCESS)
        record_descriptor_set_layout_info(pSetLayout, pCreateInfo);
    return result;
}

VK_LAYER_EXPORT void VKAPI_CALL
vkpipeline_db_DestroyDescriptorSetLayout(
    VkDevice                                    device,
    VkDescriptorSetLayout                       descriptorSetLayout,
    const VkAllocationCallbacks*                pAllocator)
{
    destroy_descriptor_set_layout_info(descriptorSetLayout);
    device_dispatch[GetKey(device)].DestroyDescriptorSetLayout(device,
                                                               descriptorSetLayout,
                                                               pAllocator);
}

// pipeline layout
static const VkPipelineLayoutCreateInfo *
get_pipeline_layout_info(VkPipelineLayout pipelineLayout)
{
    unique_lock_t lock(global_lock);
    auto it = pipeline_layouts.find(pipelineLayout);
    lock.unlock();

    if (it == pipeline_layouts.end())
        return nullptr;
    return it->second;
}

static void
destroy_pipeline_layout_info(VkPipelineLayout pipelineLayout)
{
    unique_lock_t lock(global_lock);
    auto it = pipeline_layouts.find(pipelineLayout);
    if (it != pipeline_layouts.end()) {
        free((void *)it->second->pSetLayouts);
        free((void *)it->second->pPushConstantRanges);
        free(it->second);
    }
    pipeline_layouts.erase(pipelineLayout);
    lock.unlock();
}

static void
record_pipeline_layout_info(VkPipelineLayout *pPipelineLayout,
                            const VkPipelineLayoutCreateInfo *pCreateInfo)
{
    VkPipelineLayoutCreateInfo *pPipelineLayoutInfo;

    pPipelineLayoutInfo = (VkPipelineLayoutCreateInfo *)
        malloc(sizeof(*pPipelineLayoutInfo));

    memcpy(pPipelineLayoutInfo, pCreateInfo, sizeof(*pPipelineLayoutInfo));

    if (pCreateInfo->setLayoutCount) {
        VkDescriptorSetLayout *pSetLayouts;

        pSetLayouts = (VkDescriptorSetLayout *)
            malloc(sizeof(*pSetLayouts) *
                   pCreateInfo->setLayoutCount);

        memcpy(pSetLayouts, pCreateInfo->pSetLayouts,
               sizeof(*pSetLayouts) *
                      pCreateInfo->setLayoutCount);
        pPipelineLayoutInfo->pSetLayouts = pSetLayouts;
    }

    if (pCreateInfo->pushConstantRangeCount) {
        VkPushConstantRange *pPushConstantRanges;

        pPushConstantRanges = (VkPushConstantRange *)
            malloc(sizeof(*pPushConstantRanges) *
                   pCreateInfo->pushConstantRangeCount);

        memcpy(pPushConstantRanges, pCreateInfo->pPushConstantRanges,
               sizeof(*pPushConstantRanges) *
                      pCreateInfo->pushConstantRangeCount);
        pPipelineLayoutInfo->pPushConstantRanges = pPushConstantRanges;
    }

    lock_guard_t lock(global_lock);
    pipeline_layouts[*pPipelineLayout] = pPipelineLayoutInfo;
}

VK_LAYER_EXPORT VkResult VKAPI_CALL
vkpipeline_db_CreatePipelineLayout(
    VkDevice                                    device,
    const VkPipelineLayoutCreateInfo*           pCreateInfo,
    const VkAllocationCallbacks*                pAllocator,
    VkPipelineLayout*                           pPipelineLayout)
{
    VkResult result;

    result = device_dispatch[GetKey(device)].CreatePipelineLayout(device,
                                                                  pCreateInfo,
                                                                  pAllocator,
                                                                  pPipelineLayout);
    if (result == VK_SUCCESS)
        record_pipeline_layout_info(pPipelineLayout, pCreateInfo);
    return result;
}

VK_LAYER_EXPORT void VKAPI_CALL
vkpipeline_db_DestroyPipelineLayout(
    VkDevice                                    device,
    VkPipelineLayout                            pipelineLayout,
    const VkAllocationCallbacks*                pAllocator)
{
    destroy_pipeline_layout_info(pipelineLayout);
    device_dispatch[GetKey(device)].DestroyPipelineLayout(device,
                                                          pipelineLayout,
                                                          pAllocator);
}

// render pass
static const VkRenderPassCreateInfo *
get_render_pass_info(VkRenderPass renderPass)
{
    unique_lock_t lock(global_lock);
    auto it = render_passes.find(renderPass);
    lock.unlock();

    if (it == render_passes.end())
        return nullptr;
    return it->second;
}

static void
destroy_render_pass_info(VkRenderPass renderPass)
{
    unique_lock_t lock(global_lock);
    auto it = render_passes.find(renderPass);
    if (it != render_passes.end()) {
        for (uint32_t i = 0; i < it->second->subpassCount; i++) {
            const VkSubpassDescription *subpass = &it->second->pSubpasses[i];

            free((void *)subpass->pInputAttachments);
            free((void *)subpass->pColorAttachments);
            free((void *)subpass->pDepthStencilAttachment);
            free((void *)subpass->pPreserveAttachments);
        }
        free((void *)it->second->pAttachments);
        free((void *)it->second->pSubpasses);
        free((void *)it->second->pDependencies);
        free(it->second);
    }
    render_passes.erase(renderPass);
    lock.unlock();
}

static void
record_render_pass_info(VkRenderPass *pRenderPass,
                        const VkRenderPassCreateInfo *pCreateInfo)
{
    VkRenderPassCreateInfo *pRenderPassInfo;

    pRenderPassInfo =
        (VkRenderPassCreateInfo *)malloc(sizeof(*pRenderPassInfo));

    memcpy(pRenderPassInfo, pCreateInfo, sizeof(*pRenderPassInfo));

    if (pCreateInfo->attachmentCount) {
        VkAttachmentDescription *pAttachments = (VkAttachmentDescription *)
            malloc(sizeof(*pAttachments) * pCreateInfo->attachmentCount);

        memcpy(pAttachments, pCreateInfo->pAttachments,
               sizeof(*pAttachments) * pCreateInfo->attachmentCount);
        pRenderPassInfo->pAttachments = pAttachments;
    }

    if (pCreateInfo->subpassCount) {
        VkSubpassDescription *pSubpasses = (VkSubpassDescription *)
            calloc(pCreateInfo->subpassCount, sizeof(*pSubpasses));

        for (uint32_t i = 0; i < pCreateInfo->subpassCount; i++) {
            const VkSubpassDescription *subpass = &pCreateInfo->pSubpasses[i];

            pSubpasses[i].flags = subpass->flags;
            pSubpasses[i].pipelineBindPoint = subpass->pipelineBindPoint;
            pSubpasses[i].inputAttachmentCount = subpass->inputAttachmentCount;
            pSubpasses[i].colorAttachmentCount = subpass->colorAttachmentCount;
            pSubpasses[i].preserveAttachmentCount = subpass->preserveAttachmentCount;

            if (subpass->inputAttachmentCount) {
                VkAttachmentReference *pInputAttachments;

                pInputAttachments = (VkAttachmentReference *)
                    malloc(sizeof(*pInputAttachments) * subpass->inputAttachmentCount);

                memcpy(pInputAttachments, subpass->pInputAttachments,
                       sizeof(*pInputAttachments) * subpass->inputAttachmentCount);

                pSubpasses[i].pInputAttachments = pInputAttachments;
            }

            if (subpass->colorAttachmentCount) {
                VkAttachmentReference *pColorAttachments;

                pColorAttachments = (VkAttachmentReference *)
                    malloc(sizeof(*pColorAttachments) * subpass->colorAttachmentCount);

                memcpy(pColorAttachments, subpass->pColorAttachments,
                       sizeof(*pColorAttachments) * subpass->colorAttachmentCount);

                pSubpasses[i].pColorAttachments = pColorAttachments;
            }

            if (subpass->colorAttachmentCount && subpass->pResolveAttachments) {
                VkAttachmentReference *pResolveAttachments;

                pResolveAttachments = (VkAttachmentReference *)
                    malloc(sizeof(*pResolveAttachments) * subpass->colorAttachmentCount);

                memcpy(pResolveAttachments, subpass->pResolveAttachments,
                       sizeof(*pResolveAttachments) * subpass->colorAttachmentCount);

                pSubpasses[i].pResolveAttachments = pResolveAttachments;
            }

            if (subpass->pDepthStencilAttachment) {
                VkAttachmentReference *pDepthStencilAttachment;

                pDepthStencilAttachment = (VkAttachmentReference *)
                    malloc(sizeof(*pDepthStencilAttachment));

                memcpy(pDepthStencilAttachment, subpass->pDepthStencilAttachment,
                       sizeof(*pDepthStencilAttachment));

                pSubpasses[i].pDepthStencilAttachment = pDepthStencilAttachment;
            }
            if (subpass->preserveAttachmentCount) {
                uint32_t *pPreserveAttachments;

                pPreserveAttachments = (uint32_t *)
                    malloc(sizeof(*pPreserveAttachments) * subpass->preserveAttachmentCount);

                memcpy(pPreserveAttachments, subpass->pPreserveAttachments,
                       sizeof(*pPreserveAttachments) * subpass->preserveAttachmentCount);

                pSubpasses[i].pPreserveAttachments = pPreserveAttachments;
            }
        }

        pRenderPassInfo->pSubpasses = pSubpasses;
    }

    if (pCreateInfo->dependencyCount) {
        VkSubpassDependency *pDependencies = (VkSubpassDependency *)
            malloc(sizeof(*pDependencies) * pCreateInfo->dependencyCount);

        memcpy(pDependencies, pCreateInfo->pDependencies,
               sizeof(*pDependencies) * pCreateInfo->dependencyCount);
        pRenderPassInfo->pDependencies = pDependencies;
    }

    lock_guard_t lock(global_lock);
    render_passes[*pRenderPass] = pRenderPassInfo;
}

VK_LAYER_EXPORT VkResult VKAPI_CALL
vkpipeline_db_CreateRenderPass(
    VkDevice                                    device,
    const VkRenderPassCreateInfo*               pCreateInfo,
    const VkAllocationCallbacks*                pAllocator,
    VkRenderPass*                               pRenderPass)
{
    VkResult result;

    result = device_dispatch[GetKey(device)].CreateRenderPass(device,
                                                              pCreateInfo,
                                                              pAllocator,
                                                              pRenderPass);
    if (result == VK_SUCCESS)
        record_render_pass_info(pRenderPass, pCreateInfo);
    return result;
}

VK_LAYER_EXPORT void VKAPI_CALL
vkpipeline_db_DestroyRenderPass(
    VkDevice                                    device,
    VkRenderPass                                renderPass,
    const VkAllocationCallbacks*                pAllocator)
{
    destroy_render_pass_info(renderPass);
    device_dispatch[GetKey(device)].DestroyRenderPass(device,
                                                      renderPass,
                                                      pAllocator);
}

// shader module
static const VkShaderModuleCreateInfo *
get_shader_module_info(VkShaderModule shaderModule)
{
    unique_lock_t lock(global_lock);
    auto it = shader_modules.find(shaderModule);
    lock.unlock();

    if (it == shader_modules.end())
        return nullptr;
    return it->second;
}

static void
destroy_shader_module_info(VkShaderModule shaderModule)
{
    unique_lock_t lock(global_lock);
    auto it = shader_modules.find(shaderModule);
    if (it != shader_modules.end()) {
        free((void *)it->second->pCode);
        free(it->second);
    }
    shader_modules.erase(shaderModule);
    lock.unlock();
}

static void
record_shader_module_info(VkShaderModule *pShaderModule,
                          const VkShaderModuleCreateInfo *pCreateInfo)
{
    VkShaderModuleCreateInfo *pShaderModuleInfo;
    uint32_t *pCode;

    pShaderModuleInfo =
        (VkShaderModuleCreateInfo *)malloc(sizeof(*pShaderModuleInfo));
    
    memcpy(pShaderModuleInfo, pCreateInfo, sizeof(*pShaderModuleInfo));

    pCode = (uint32_t *)malloc(pCreateInfo->codeSize);
   
    memcpy(pCode, pCreateInfo->pCode, pCreateInfo->codeSize);
    pShaderModuleInfo->pCode = pCode;

    lock_guard_t lock(global_lock);
    shader_modules[*pShaderModule] = pShaderModuleInfo;
}

VK_LAYER_EXPORT VkResult VKAPI_CALL
vkpipeline_db_CreateShaderModule(
    VkDevice                                    device,
    const VkShaderModuleCreateInfo*             pCreateInfo,
    const VkAllocationCallbacks*                pAllocator,
    VkShaderModule*                             pShaderModule)
{
    VkResult result;

    result = device_dispatch[GetKey(device)].CreateShaderModule(device,
                                                               pCreateInfo,
                                                               pAllocator,
                                                               pShaderModule);
    if (result == VK_SUCCESS)
        record_shader_module_info(pShaderModule, pCreateInfo);
    return result;
}

VK_LAYER_EXPORT void VKAPI_CALL
vkpipeline_db_DestroyShaderModule(
    VkDevice                                    device,
    VkShaderModule                              shaderModule,
    const VkAllocationCallbacks*                pAllocator)
{
    destroy_shader_module_info(shaderModule);
    device_dispatch[GetKey(device)].DestroyShaderModule(device,
                                                        shaderModule,
                                                        pAllocator);
}

// graphics/compute pipeline
static const char *
get_capture_path()
{
    static const char *capture_path = NULL;
    static bool read_env_var = false;
    static bool print_error = false;

    if (!read_env_var) {
        capture_path = getenv("VKPIPELINE_DB_CAPTURE_PATH");
        read_env_var = true;
    }

    if (!capture_path) {
        if (!print_error) {
            fprintf(stderr, "Failed to capture because "
                    "VKPIPELINE_DB_CAPTURE_PATH is not set, aborted!\n");
            print_error = 1;
        }
    }

    return capture_path;
}

static void
free_pipeline(struct pipeline_info *pipeline)
{
    free(pipeline->pShaderStagesInfo);
    free(pipeline->pShaderModulesInfo);
    free(pipeline->pSetLayoutsInfo);
    free(pipeline);
}

static int
capture_pipeline(struct pipeline_info *pipeline)
{
    const char *capture_path;
    struct blob metadata;
    char filename[1024];
    int ret = 0;
    int fd;

    if (!pipeline)
        return -1;

    capture_path = get_capture_path();
    if (!capture_path)
        return -1;

    blob_init(&metadata);
    serialize_pipeline(pipeline, &metadata);
    free_pipeline(pipeline);

    snprintf(filename, sizeof(filename),
             "%s/%d.pipeline_test", capture_path, pipeline_count++);

    fd = open(filename, O_WRONLY | O_CREAT, 0644);
    if (fd == -1) {
        perror("open");
        ret = -1;
        goto fail_open;
    }

    if (write(fd, metadata.data, metadata.size) == -1) {
        perror("write");
        ret = -1;
        goto fail_write;
    }

fail_write:
    close(fd);
fail_open:
    free(metadata.data);

    return ret;
}

static struct pipeline_info *
get_graphics_pipeline_info(const VkGraphicsPipelineCreateInfo *pCreateInfo)
{
    struct pipeline_info *pipeline;

    pipeline = (struct pipeline_info *)calloc(1, sizeof(*pipeline));
   
    pipeline->bindPoint = VK_PIPELINE_BIND_POINT_GRAPHICS;
    pipeline->stageCount = pCreateInfo->stageCount;

    // shader stages and modules
    pipeline->pShaderStagesInfo = (VkPipelineShaderStageCreateInfo *)
        calloc(pipeline->stageCount, sizeof(*pipeline->pShaderStagesInfo));
    pipeline->pShaderModulesInfo = (VkShaderModuleCreateInfo *)
        calloc(pipeline->stageCount, sizeof (*pipeline->pShaderModulesInfo));

    for (uint32_t i = 0; i < pCreateInfo->stageCount; i++) {
        const VkPipelineShaderStageCreateInfo *stage = &pCreateInfo->pStages[i];
        const VkShaderModuleCreateInfo *pShaderModuleInfo =
            get_shader_module_info(stage->module);

        pipeline->pShaderStagesInfo[i] = *stage;
        pipeline->pShaderModulesInfo[i] = *pShaderModuleInfo;
    }

    // graphics states
    pipeline->vertexInputState = *pCreateInfo->pVertexInputState;
    pipeline->inputAssemblyState = *pCreateInfo->pInputAssemblyState;
    if (pCreateInfo->pTessellationState)
        pipeline->tessellationState = *pCreateInfo->pTessellationState;
    if (pCreateInfo->pViewportState)
        pipeline->viewportState = *pCreateInfo->pViewportState;
    pipeline->rasterizationState = *pCreateInfo->pRasterizationState;
    if (pCreateInfo->pMultisampleState)
        pipeline->multisampleState = *pCreateInfo->pMultisampleState;
    if (pCreateInfo->pDepthStencilState)
        pipeline->depthStencilState = *pCreateInfo->pDepthStencilState;
    if (pCreateInfo->pColorBlendState)
        pipeline->colorBlendState = *pCreateInfo->pColorBlendState;
    if (pCreateInfo->pDynamicState)
        pipeline->dynamicState = *pCreateInfo->pDynamicState;

    // pipeline layout
    const VkPipelineLayoutCreateInfo *pPipelineLayoutInfo =
        get_pipeline_layout_info(pCreateInfo->layout);
    pipeline->pipelineLayoutInfo = *pPipelineLayoutInfo;

    pipeline->pSetLayoutsInfo = (VkDescriptorSetLayoutCreateInfo *)
        malloc(sizeof(*pipeline->pSetLayoutsInfo) * pipeline->pipelineLayoutInfo.setLayoutCount);

    for (uint32_t i = 0; i < pPipelineLayoutInfo->setLayoutCount; i++) {
        const VkDescriptorSetLayout *layout =
            &pPipelineLayoutInfo->pSetLayouts[i];
        const VkDescriptorSetLayoutCreateInfo *pSetLayoutInfo =
            get_descriptor_set_layout_info(*layout);

        pipeline->pSetLayoutsInfo[i] = *pSetLayoutInfo;
    }

    // render pass
    const VkRenderPassCreateInfo *pRenderPassInfo =
        get_render_pass_info(pCreateInfo->renderPass);
    pipeline->renderPassInfo = *pRenderPassInfo;

    return pipeline;
}

VK_LAYER_EXPORT VkResult VKAPI_CALL
vkpipeline_db_CreateGraphicsPipelines(
    VkDevice                                    device,
    VkPipelineCache                             pipelineCache,
    uint32_t                                    createInfoCount,
    const VkGraphicsPipelineCreateInfo*         pCreateInfos,
    const VkAllocationCallbacks*                pAllocator,
    VkPipeline*                                 pPipelines)
{
    VkResult result;

    result = device_dispatch[GetKey(device)].CreateGraphicsPipelines(device,
                                                                     pipelineCache,
                                                                     createInfoCount,
                                                                     pCreateInfos,
                                                                     pAllocator,
                                                                     pPipelines);
    if (result == VK_SUCCESS) {
        for (uint32_t i = 0; i < createInfoCount; i++) {
            struct pipeline_info *pipeline;

            pipeline = get_graphics_pipeline_info(&pCreateInfos[i]);
            capture_pipeline(pipeline);
        }
    }

    return result;
}

static struct pipeline_info *
get_compute_pipeline_info(const VkComputePipelineCreateInfo *pCreateInfo)
{
    struct pipeline_info *pipeline;

    pipeline = (struct pipeline_info *)calloc(1, sizeof(*pipeline));

    pipeline->bindPoint = VK_PIPELINE_BIND_POINT_COMPUTE;
    pipeline->stageCount = 1;

    // shader stages and modules
    const VkShaderModuleCreateInfo *pShaderModuleInfo =
        get_shader_module_info(pCreateInfo->stage.module);
    
    pipeline->pShaderStagesInfo = (VkPipelineShaderStageCreateInfo *)
        calloc(pipeline->stageCount, sizeof(*pipeline->pShaderStagesInfo));
    pipeline->pShaderModulesInfo = (VkShaderModuleCreateInfo *)
        calloc(pipeline->stageCount, sizeof (*pipeline->pShaderModulesInfo));

    pipeline->pShaderStagesInfo[0] = pCreateInfo->stage;
    pipeline->pShaderModulesInfo[0] = *pShaderModuleInfo;

    // pipeline layout
    const VkPipelineLayoutCreateInfo *pPipelineLayoutInfo =
        get_pipeline_layout_info(pCreateInfo->layout);
    pipeline->pipelineLayoutInfo = *pPipelineLayoutInfo;

    pipeline->pSetLayoutsInfo = (VkDescriptorSetLayoutCreateInfo *)
        malloc(sizeof(*pipeline->pSetLayoutsInfo) * pipeline->pipelineLayoutInfo.setLayoutCount);

    for (uint32_t i = 0; i < pPipelineLayoutInfo->setLayoutCount; i++) {
        const VkDescriptorSetLayout *layout =
            &pPipelineLayoutInfo->pSetLayouts[i];
        const VkDescriptorSetLayoutCreateInfo *pSetLayoutInfo =
            get_descriptor_set_layout_info(*layout);

        pipeline->pSetLayoutsInfo[i] = *pSetLayoutInfo;
    }

    return pipeline;
}

VK_LAYER_EXPORT VkResult VKAPI_CALL
vkpipeline_db_CreateComputePipelines(
    VkDevice                                    device,
    VkPipelineCache                             pipelineCache,
    uint32_t                                    createInfoCount,
    const VkComputePipelineCreateInfo*          pCreateInfos,
    const VkAllocationCallbacks*                pAllocator,
    VkPipeline*                                 pPipelines)
{
    VkResult result;

    result = device_dispatch[GetKey(device)].CreateComputePipelines(device,
                                                                    pipelineCache,
                                                                    createInfoCount,
                                                                    pCreateInfos,
                                                                    pAllocator,
                                                                    pPipelines);
    if (result == VK_SUCCESS) {
        for (uint32_t i = 0; i < createInfoCount; i++) {
            struct pipeline_info *pipeline;

            pipeline = get_compute_pipeline_info(&pCreateInfos[i]);
            capture_pipeline(pipeline);
        }
    }

    return result;
}

static PFN_vkVoidFunction
vkpipeline_db_GetDeviceProcAddr(VkDevice device, const char *pName);
static PFN_vkVoidFunction
vkpipeline_db_GetInstanceProcAddr(VkInstance instance, const char *pName);

#define FUNC(name) (void *)vkpipeline_db_##name

// list of functions that are intercepted by this layer
static const struct {
    const char *name;
    void *ptr;
} funcs[] = {
    { "vkGetInstanceProcAddr", FUNC(GetInstanceProcAddr) },
    { "vkEnumerateInstanceLayerProperties", FUNC(EnumerateInstanceLayerProperties) },
    { "vkEnumerateInstanceExtensionProperties", FUNC(EnumerateInstanceExtensionProperties) },
    { "vkCreateInstance", FUNC(CreateInstance) },
    { "vkDestroyInstance", FUNC(DestroyInstance) },
    { "vkGetDeviceProcAddr", FUNC(GetDeviceProcAddr) },
    { "vkEnumerateDeviceLayerProperties", FUNC(EnumerateDeviceLayerProperties) },
    { "vkEnumerateDeviceExtensionProperties", FUNC(EnumerateDeviceExtensionProperties) },
    { "vkCreateDevice", FUNC(CreateDevice) },
    { "vkDestroyDevice", FUNC(DestroyDevice) },
    { "vkCreateSampler", FUNC(CreateSampler) },
    { "vkDestroySampler", FUNC(DestroySampler) },
    { "vkCreateDescriptorSetLayout", FUNC(CreateDescriptorSetLayout) },
    { "vkDestroyDescriptorSetLayout", FUNC(DestroyDescriptorSetLayout) },
    { "vkCreatePipelineLayout", FUNC(CreatePipelineLayout) },
    { "vkDestroyPipelineLayout", FUNC(DestroyPipelineLayout) },
    { "vkCreateShaderModule", FUNC(CreateShaderModule) },
    { "vkDestroyShaderModule", FUNC(DestroyShaderModule) },
    { "vkCreateRenderPass", FUNC(CreateRenderPass) },
    { "vkDestroyRenderPass", FUNC(DestroyRenderPass) },
    { "vkCreateGraphicsPipelines", FUNC(CreateGraphicsPipelines) },
    { "vkCreateComputePipelines", FUNC(CreateComputePipelines) },
};

#undef FUNC

static PFN_vkVoidFunction
vkpipeline_db_GetDeviceProcAddr(
    VkDevice                                    device,
    const char                                  *pName)
{
    for (uint32_t i = 0; i < sizeof(funcs) / sizeof(funcs[0]); i++) {
        if (!strcmp(pName, funcs[i].name))
            return (PFN_vkVoidFunction)funcs[i].ptr;
    }

    lock_guard_t l(global_lock);
    return device_dispatch[GetKey(device)].GetDeviceProcAddr(device, pName);
}

static PFN_vkVoidFunction
vkpipeline_db_GetInstanceProcAddr(
    VkInstance                                  instance,
    const char                                  *pName)
{
    for (uint32_t i = 0; i < sizeof(funcs) / sizeof(funcs[0]); i++) {
        if (!strcmp(pName, funcs[i].name))
            return (PFN_vkVoidFunction)funcs[i].ptr;
    }

    lock_guard_t l(global_lock);
    return instance_dispatch[GetKey(instance)].GetInstanceProcAddr(instance, pName);
}

VK_LAYER_EXPORT VkResult VKAPI_CALL
vkpipeline_db_NegotiateLoaderLayerInterfaceVersion(
    VkNegotiateLayerInterface *pVersionStruct)
{
    if (pVersionStruct->loaderLayerInterfaceVersion > 2)
        pVersionStruct->loaderLayerInterfaceVersion = 2;

    pVersionStruct->pfnGetInstanceProcAddr = vkpipeline_db_GetInstanceProcAddr;
    pVersionStruct->pfnGetDeviceProcAddr = vkpipeline_db_GetDeviceProcAddr;
    pVersionStruct->pfnGetPhysicalDeviceProcAddr = NULL;

    return VK_SUCCESS;
}