// Copyright 2018 The Dawn & Tint Authors
//
// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions are met:
//
// 1. Redistributions of source code must retain the above copyright notice, this
//    list of conditions and the following disclaimer.
//
// 2. Redistributions in binary form must reproduce the above copyright notice,
//    this list of conditions and the following disclaimer in the documentation
//    and/or other materials provided with the distribution.
//
// 3. Neither the name of the copyright holder nor the names of its
//    contributors may be used to endorse or promote products derived from
//    this software without specific prior written permission.
//
// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
// AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
// DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
// FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
// DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
// SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
// CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
// OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

#include "src/dawn/native/vulkan/BindGroupLayoutVk.h"

#include <utility>

#include "absl/container/flat_hash_map.h"
#include "src/dawn/common/MatchVariant.h"
#include "src/dawn/common/Range.h"
#include "src/dawn/common/ityp_vector.h"
#include "src/dawn/native/CacheKey.h"
#include "src/dawn/native/vulkan/DescriptorSetAllocator.h"
#include "src/dawn/native/vulkan/DeviceVk.h"
#include "src/dawn/native/vulkan/FencedDeleter.h"
#include "src/dawn/native/vulkan/PhysicalDeviceVk.h"
#include "src/dawn/native/vulkan/SamplerVk.h"
#include "src/dawn/native/vulkan/UtilsVulkan.h"
#include "src/dawn/native/vulkan/VulkanError.h"

namespace dawn::native::vulkan {

namespace {

// Helper function (and result structure) that precomputes all the information related to static
// bindings that might be for Vulkan BindGroupLayout.
struct VulkanStaticBindings {
    ityp::vector<BindingIndex, VkDescriptorSetLayoutBinding> bindings;
    ityp::vector<BindingIndex, VkDescriptorUpdateTemplateEntry> updateTemplateEntries;
    absl::flat_hash_map<VkDescriptorType, uint32_t> descriptorCountPerType;
    TextureToStaticSamplerMap textureToStaticSampler;
};
ResultOrError<VulkanStaticBindings> ComputeVulkanStaticBindings(
    Device* device,
    const BindGroupLayoutInternalBase* layout,
    const BindGroupLayout::StaticSamplerSpecializationMap& staticSamplerSpecializations) {
    VulkanStaticBindings res;

    // Build a map of texture indices to sampler indices. This maps the texture to
    // the sampler which will be used in VK in the combined image sampler entry.
    for (BindingIndex bindingIndex : layout->GetStaticSamplerIndices()) {
        auto samplerBindingInfo =
            std::get<StaticSamplerBindingInfo>(layout->GetBindingInfo(bindingIndex).bindingLayout);
        // This is a static sampler combined with textures dynamically in the shader.
        if (samplerBindingInfo.use == StaticSamplerUse::Freestanding) {
            continue;
        }

        res.textureToStaticSampler[samplerBindingInfo.sampledTextureIndex] = bindingIndex;
    }

    // Compute the bindings that will be chained in the DescriptorSetLayout create info. We add
    // one entry per binding set. This could be optimized by computing continuous ranges of
    // bindings of the same type.
    res.bindings.reserve(layout->GetBindingCount());
    res.updateTemplateEntries.reserve(layout->GetBindingCount());

    for (BindingIndex bindingIndex : Range(layout->GetBindingCount())) {
        const BindingInfo& bindingInfo = layout->GetBindingInfo(bindingIndex);

        // Skip over bindings that cannot be seen by any shaders as they could cause us to create
        // bindgroups with more bindings than the VkDevice's limits. However keep dynamic buffers
        // as the amount of dynamic offsets need to stay the same as WebGPU's so we can passthrough
        // the dynamic offsets.
        if (bindingInfo.visibility == wgpu::ShaderStage::None &&
            bindingIndex >= layout->GetDynamicBufferCount()) {
            continue;
        }

        // This texture will be bound into the VkDescriptorSet at the index for the sampler itself.
        if (res.textureToStaticSampler.contains(bindingIndex)) {
            continue;
        }

        // Vulkan descriptor set layouts have one entry for binding_array. Only handle their first
        // element as subsequent ones will be part of the already added
        // VkDescriptorSetLayoutBinding.
        if (bindingInfo.indexInArray != BindingIndex(0u)) {
            continue;
        }

        VkDescriptorSetLayoutBinding vkBinding{
            .binding = uint32_t{bindingIndex},
            .descriptorType = VulkanDescriptorType(bindingInfo),
            .descriptorCount = uint32_t{bindingInfo.arraySize},
            .stageFlags = VulkanShaderStages(bindingInfo.visibility),
            .pImmutableSamplers = nullptr,
        };
        size_t descriptorCount = vkBinding.descriptorCount;

        // Static samplers are set at VkDescriptorSetLayout creation time.
        if (std::holds_alternative<StaticSamplerBindingInfo>(bindingInfo.bindingLayout)) {
            auto samplerLayout = std::get<StaticSamplerBindingInfo>(bindingInfo.bindingLayout);
            auto sampler = ToBackend(samplerLayout.sampler);

            // Override with the specialization's sampler if there's one. This is used to replace
            // samplers with the correct YCbCr sampler when JITing pipelines.
            if (auto it = staticSamplerSpecializations.find(bindingIndex);
                it != staticSamplerSpecializations.end()) {
                DAWN_CHECK(samplerLayout.use == StaticSamplerUse::InternalForExternalTexture);
                DAWN_TRY_ASSIGN(sampler, Sampler::Create(device, it->second));
            }

            vkBinding.pImmutableSamplers = &sampler->GetHandle().GetHandle();

            if (sampler->IsYCbCr()) {
                // A YCbCr sampler can take up multiple Vk descriptor slots.  There is a
                // recommended Vulkan API to query how many slots a YCbCr sampler should take, but
                // it is not clear how to actually pass the Android external format to that API.
                // However, the spec for that API says the following:
                // "combinedImageSamplerDescriptorCount is a number between 1 and the number of
                // planes in the format. A descriptor set layout binding with immutable Y′CBCR
                // conversion samplers will have a maximum combinedImageSamplerDescriptorCount
                // which is the maximum across all formats supported by its samplers of the
                // combinedImageSamplerDescriptorCount for each format." Hence, we simply hardcode
                // the maximum number of planes that an external format can have here. The number
                // of overall YCbCr descriptors will be relatively small and these pools are not an
                // overall bottleneck on memory usage.
                DAWN_CHECK(bindingInfo.arraySize == BindingIndex(1u));
                descriptorCount = 3;
            }

            if (samplerLayout.use != StaticSamplerUse::Freestanding) {
                // Write the info of the texture at the binding index for the static sampler.
                res.updateTemplateEntries.push_back({
                    .dstBinding = vkBinding.binding,
                    .dstArrayElement = 0,
                    .descriptorCount = vkBinding.descriptorCount,
                    .descriptorType = vkBinding.descriptorType,
                    .offset =
                        uint32_t{samplerLayout.sampledTextureIndex} * sizeof(DescriptorUpdateData),
                    .stride = sizeof(DescriptorUpdateData),
                });
            }
        } else {
            res.updateTemplateEntries.push_back({
                .dstBinding = vkBinding.binding,
                .dstArrayElement = 0,
                .descriptorCount = vkBinding.descriptorCount,
                .descriptorType = vkBinding.descriptorType,
                .offset = uint32_t{bindingIndex} * sizeof(DescriptorUpdateData),
                .stride = sizeof(DescriptorUpdateData),
            });
        }

        res.bindings.emplace_back(vkBinding);

        // absl:flat_hash_map::operator[] will return 0 if the key doesn't exist.
        res.descriptorCountPerType[vkBinding.descriptorType] += descriptorCount;
    }

    return std::move(res);
}

}  // anonymous namespace

VkDescriptorType VulkanDescriptorType(const BindingInfo& bindingInfo) {
    return MatchVariant(
        bindingInfo.bindingLayout,
        [](const BufferBindingInfo& layout) -> VkDescriptorType {
            switch (layout.type) {
                case wgpu::BufferBindingType::Uniform:
                    if (layout.hasDynamicOffset) {
                        return VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER_DYNAMIC;
                    }
                    return VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER;
                case wgpu::BufferBindingType::Storage:
                case kInternalStorageBufferBinding:
                case wgpu::BufferBindingType::ReadOnlyStorage:
                case kInternalReadOnlyStorageBufferBinding:
                    if (layout.hasDynamicOffset) {
                        return VK_DESCRIPTOR_TYPE_STORAGE_BUFFER_DYNAMIC;
                    }
                    return VK_DESCRIPTOR_TYPE_STORAGE_BUFFER;
                case wgpu::BufferBindingType::BindingNotUsed:
                case wgpu::BufferBindingType::Undefined:
                    DAWN_UNREACHABLE();
                    return VK_DESCRIPTOR_TYPE_STORAGE_BUFFER;
            }
            DAWN_UNREACHABLE();
        },
        [](const SamplerBindingInfo&) { return VK_DESCRIPTOR_TYPE_SAMPLER; },
        [](const StaticSamplerBindingInfo& layout) {
            // Make this entry into a combined image sampler iff the client
            // specified a single texture binding to be paired with it.
            return (layout.use == StaticSamplerUse::Freestanding)
                       ? VK_DESCRIPTOR_TYPE_SAMPLER
                       : VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER;
        },
        [](const TextureBindingInfo&) { return VK_DESCRIPTOR_TYPE_SAMPLED_IMAGE; },
        [](const StorageTextureBindingInfo&) { return VK_DESCRIPTOR_TYPE_STORAGE_IMAGE; },
        [](const TexelBufferBindingInfo& layout) -> VkDescriptorType {
            switch (layout.access) {
                case wgpu::TexelBufferAccess::ReadOnly:
                    // TODO(crbug.com/382544164): Investigate whether read-only texel buffers
                    // should use VK_DESCRIPTOR_TYPE_UNIFORM_TEXEL_BUFFER for broader format
                    // support, or stay on VK_DESCRIPTOR_TYPE_STORAGE_TEXEL_BUFFER for bindless
                    // compatibility (uniform texel buffers have limited bindless support on
                    // Vulkan and would require a separate descriptor array in the resource
                    // table).
                    [[fallthrough]];
                case wgpu::TexelBufferAccess::ReadWrite:
                    return VK_DESCRIPTOR_TYPE_STORAGE_TEXEL_BUFFER;
                case wgpu::TexelBufferAccess::Undefined:
                    DAWN_UNREACHABLE();
            }
            DAWN_UNREACHABLE();
        },

        [](const InputAttachmentBindingInfo&) { return VK_DESCRIPTOR_TYPE_INPUT_ATTACHMENT; },
        [](const ExternalTextureBindingInfo&) -> VkDescriptorType { DAWN_UNREACHABLE(); });
}

// static
ResultOrError<Ref<BindGroupLayout>> BindGroupLayout::Create(
    Device* device,
    const UnpackedPtr<BindGroupLayoutDescriptor>& descriptor) {
    Ref<BindGroupLayout> bgl = AcquireRef(new BindGroupLayout(device, descriptor));
    DAWN_TRY(bgl->Initialize());
    return bgl;
}

BindGroupLayout::BindGroupLayout(DeviceBase* device,
                                 const UnpackedPtr<BindGroupLayoutDescriptor>& descriptor)
    : BindGroupLayoutInternalBase(device, descriptor),
      mBindGroupAllocator(MakeFrontendBindGroupAllocator<BindGroup>(4096)) {}

BindGroupLayout::~BindGroupLayout() = default;

MaybeError BindGroupLayout::Initialize() {
    Device* device = ToBackend(GetDevice());

    VulkanStaticBindings bindings;
    DAWN_TRY_ASSIGN(bindings, ComputeVulkanStaticBindings(device, this, {}));
    mTextureToStaticSampler = std::move(bindings.textureToStaticSampler);

    VkDescriptorSetLayoutCreateInfo createInfo{
        .sType = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_LAYOUT_CREATE_INFO,
        .pNext = nullptr,
        .flags = 0,
        .bindingCount = uint32_t{bindings.bindings.size()},
        .pBindings = bindings.bindings.data(),
    };

    // Record cache key information now since the createInfo is not stored.
    StreamIn(&mCacheKey, createInfo);

    SpecializationResult r;
    DAWN_TRY_ASSIGN(r, GetOrCreateSpecialization({}));

    mHandle = r.layout;
    mDescriptorUpdateTemplate = r.descriptorUpdateTemplate;
    mDescriptorSetAllocator = r.allocator.Get();

    SetLabelImpl();

    return {};
}

ResultOrError<BindGroupLayout::SpecializationResult> BindGroupLayout::GetOrCreateSpecialization(
    const Specialization& specialization) {
    if (auto specialized = mSpecializations.ConstUse(
            [&](auto specializations) -> std::optional<SpecializationResult> {
                if (auto it = specializations->find(specialization); it != specializations->end()) {
                    return it->second;
                }
                return std::nullopt;
            });
        specialized) {
        return *specialized;
    }

    Device* device = ToBackend(GetDevice());
    VulkanStaticBindings bindings;
    DAWN_TRY_ASSIGN(bindings,
                    ComputeVulkanStaticBindings(device, this, specialization.staticSamplers));

    VkDescriptorSetLayoutCreateInfo createInfo{
        .sType = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_LAYOUT_CREATE_INFO,
        .pNext = nullptr,
        .flags = 0,
        .bindingCount = uint32_t{bindings.bindings.size()},
        .pBindings = bindings.bindings.data(),
    };

    VkDescriptorSetLayout specialized;
    DAWN_TRY(CheckVkSuccess(device->fn.CreateDescriptorSetLayout(device->GetVkDevice(), &createInfo,
                                                                 nullptr, &*specialized),
                            "CreateDescriptorSetLayout"));

    // VUID-VkDescriptorUpdateTemplateCreateInfo-descriptorUpdateEntryCount-arraylength requires
    // descriptorUpdateEntryCount to be greater than 0.
    VkDescriptorUpdateTemplate updateTemplate = VK_NULL_HANDLE;
    if (device->IsToggleEnabled(Toggle::VulkanUseDescriptorUpdateTemplate) &&
        !bindings.updateTemplateEntries.empty()) {
        VkDescriptorUpdateTemplateCreateInfo templateCreateInfo{
            .sType = VK_STRUCTURE_TYPE_DESCRIPTOR_UPDATE_TEMPLATE_CREATE_INFO,
            .pNext = nullptr,
            .flags = 0,
            .descriptorUpdateEntryCount = uint32_t{bindings.updateTemplateEntries.size()},
            .pDescriptorUpdateEntries = bindings.updateTemplateEntries.data(),
            .templateType = VK_DESCRIPTOR_UPDATE_TEMPLATE_TYPE_DESCRIPTOR_SET,
            .descriptorSetLayout = specialized,
            // pipelineBindPoint, pipelineLayout, and set are only used when templateType is
            // VK_DESCRIPTOR_UPDATE_TEMPLATE_TYPE_PUSH_DESCRIPTORS. For
            // VK_DESCRIPTOR_UPDATE_TEMPLATE_TYPE_DESCRIPTOR_SET, the Vulkan spec ignores these
            // three fields because the template only updates a standalone VkDescriptorSet, which
            // can subsequently be bound to both graphics and compute pipelines.
            .pipelineBindPoint = VK_PIPELINE_BIND_POINT_GRAPHICS,
            .pipelineLayout = VkPipelineLayout(VK_NULL_HANDLE),
            .set = 0,
        };
        DAWN_TRY_WITH_CLEANUP(
            CheckVkSuccess(
                device->fn.CreateDescriptorUpdateTemplate(
                    device->GetVkDevice(), &templateCreateInfo, nullptr, &*updateTemplate),
                "CreateDescriptorUpdateTemplate"),
            {
                device->fn.DestroyDescriptorSetLayout(device->GetVkDevice(), specialized, nullptr);
            });
    }

    Ref<DescriptorSetAllocator> allocator =
        DescriptorSetAllocator::Create(device, std::move(bindings.descriptorCountPerType));

    return mSpecializations.Use([&](auto specializations) -> ResultOrError<SpecializationResult> {
        SpecializationResult result{.layout = specialized,
                                    .descriptorUpdateTemplate = updateTemplate,
                                    .allocator = std::move(allocator)};
        auto [it, inserted] = specializations->insert({specialization, result});
        if (!inserted) {
            device->fn.DestroyDescriptorSetLayout(device->GetVkDevice(), specialized, nullptr);
            device->fn.DestroyDescriptorUpdateTemplate(device->GetVkDevice(), updateTemplate,
                                                       nullptr);
            return it->second;
        }
        return result;
    });
}

ResultOrError<VkDescriptorSetLayout> BindGroupLayout::GetOrCreateSpecializedHandle(
    const Specialization& specialization) {
    SpecializationResult specialized;
    DAWN_TRY_ASSIGN(specialized, GetOrCreateSpecialization(specialization));
    return specialized.layout;
}

void BindGroupLayout::DestroyImpl(DestroyReason reason) {
    BindGroupLayoutInternalBase::DestroyImpl(reason);

    Device* device = ToBackend(GetDevice());

    // mHandle and mDescriptorUpdateTemplate are destroyed in the loop below.
    mHandle = VK_NULL_HANDLE;
    mDescriptorUpdateTemplate = VK_NULL_HANDLE;

    mDescriptorSetAllocator = nullptr;

    // DescriptorSetLayouts and DescriptorUpdateTemplates aren't used by execution on the GPU and
    // can be deleted at any time, so we can destroy them immediately instead of using the
    // FencedDeleter.
    mSpecializations.Use([&](auto specializations) {
        for (auto& [_, specialized] : *specializations) {
            device->fn.DestroyDescriptorSetLayout(device->GetVkDevice(), specialized.layout,
                                                  nullptr);
            device->fn.DestroyDescriptorUpdateTemplate(
                device->GetVkDevice(), specialized.descriptorUpdateTemplate, nullptr);
        }
        specializations->clear();
    });
}

VkDescriptorSetLayout BindGroupLayout::GetHandle() const {
    return mHandle;
}

VkDescriptorUpdateTemplate BindGroupLayout::GetDescriptorUpdateTemplate() const {
    return mDescriptorUpdateTemplate;
}

ResultOrError<Ref<BindGroup>> BindGroupLayout::AllocateBindGroup(
    const UnpackedPtr<BindGroupDescriptor>& descriptor) {
    DescriptorSetAllocation descriptorSetAllocation;
    DAWN_TRY_ASSIGN(descriptorSetAllocation, mDescriptorSetAllocator->Allocate(GetHandle()));

    return AcquireRef(
        mBindGroupAllocator->Allocate(ToBackend(GetDevice()), descriptor, descriptorSetAllocation));
}

void BindGroupLayout::DeallocateBindGroup(BindGroup* bindGroup) {
    mBindGroupAllocator->Deallocate(bindGroup);
}

void BindGroupLayout::DeallocateDescriptorSet(DescriptorSetAllocation* descriptorSetAllocation) {
    mDescriptorSetAllocator->Deallocate(descriptorSetAllocation);
}

void BindGroupLayout::ReduceMemoryUsage() {
    mBindGroupAllocator->DeleteEmptySlabs();
}

ResultOrError<std::unique_ptr<OwnedDescriptorSet>> BindGroupLayout::GetSpecializedSetFor(
    const BindGroup* bg,
    const Specialization& specialization) {
    DAWN_ASSERT(bg->GetLayout() == this);

    SpecializationResult specialized;
    DAWN_TRY_ASSIGN(specialized, GetOrCreateSpecialization(specialization));

    DescriptorSetAllocation dsAllocation;
    DAWN_TRY_ASSIGN(dsAllocation, specialized.allocator->Allocate(specialized.layout));

    bg->WriteDescriptorSet(dsAllocation.set, specialized.descriptorUpdateTemplate);
    return std::make_unique<OwnedDescriptorSet>(std::move(specialized.allocator), dsAllocation);
}

const TextureToStaticSamplerMap& BindGroupLayout::GetTextureToStaticSamplerMap() const {
    return mTextureToStaticSampler;
}

void BindGroupLayout::SetLabelImpl() {
    SetDebugName(ToBackend(GetDevice()), mHandle, "Dawn_BindGroupLayout", GetLabel());
    SetDebugName(ToBackend(GetDevice()), mDescriptorUpdateTemplate, "Dawn_DescriptorUpdateTemplate",
                 GetLabel());
}

// OwnedDescriptorSet

OwnedDescriptorSet::OwnedDescriptorSet(Ref<DescriptorSetAllocator> allocator,
                                       DescriptorSetAllocation allocation)
    : mAllocation(allocation), mAllocator(std::move(allocator)) {}

OwnedDescriptorSet::~OwnedDescriptorSet() {
    mAllocator->Deallocate(&mAllocation);
    mAllocator = nullptr;
}

VkDescriptorSet OwnedDescriptorSet::GetHandle() const {
    return mAllocation.set;
}

}  // namespace dawn::native::vulkan