/* Copyright (c) 2015-2026 The Khronos Group Inc.
 * Copyright (c) 2015-2026 Valve Corporation
 * Copyright (c) 2015-2026 LunarG, Inc.
 * Copyright (C) 2015-2026 Google Inc.
 * Modifications Copyright (C) 2020 Advanced Micro Devices, Inc. All rights reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include "state_tracker/pipeline_layout_state.h"

#include "error_message/error_strings.h"
#include "state_tracker/state_tracker.h"
#include "state_tracker/descriptor_sets.h"

// NOTE: The following canonical id dictionaries are global and thread-safe objects:
//      pipeline_layout_set_layouts_dict
//      pipeline_layout_compat_dict
//      push_constant_ranges_dict
//
// So far we did not have any issues with them. However, we had issues with DescriptorSetLayoutId
// canonical ids (can survice between test sessions and then incorrectly accessed by another device).
// That issue was solved by making DescriptorSetLayoutId dictionary local to DeviceState
// (DeviceState::descriptor_set_layout_canonical_ids_). If there are any issues with these global
// maps consider making them per-device or per-instance state.

// Dictionary of canonical form of the pipeline set layout of descriptor set layouts
static PipelineLayoutSetLayoutsDict pipeline_layout_set_layouts_dict;

// Dictionary of canonical form of the "compatible for set" records
static PipelineLayoutCompatDict pipeline_layout_compat_dict;

static PushConstantRangesDict push_constant_ranges_dict;

size_t PipelineLayoutCompatDef::hash() const {
    hash_util::HashCombiner hc;
    // The set number is integral to the CompatDef's distinctiveness
    hc << set << push_constant_ranges.get() << is_independent_sets;
    const auto& descriptor_set_layouts = *set_layouts_id.get();
    for (uint32_t i = 0; i <= set; i++) {
        hc << descriptor_set_layouts[i].get();
    }
    return hc.Value();
}

bool PipelineLayoutCompatDef::operator==(const PipelineLayoutCompatDef& other) const {
    if ((set != other.set) || (push_constant_ranges != other.push_constant_ranges) ||
        (is_independent_sets != other.is_independent_sets)) {
        return false;
    }

    if (set_layouts_id == other.set_layouts_id) {
        // if it's the same set_layouts_id, then *any* subset will match
        return true;
    }

    // They aren't exactly the same PipelineLayoutSetLayouts, so we need to check if the required subsets match
    const auto& descriptor_set_layouts = *set_layouts_id.get();
    assert(set < descriptor_set_layouts.size());
    const auto& other_ds_layouts = *other.set_layouts_id.get();
    assert(set < other_ds_layouts.size());
    for (uint32_t i = 0; i <= set; i++) {
        if (descriptor_set_layouts[i] != other_ds_layouts[i]) {
            if (other.is_independent_sets) {
                // This is hit because the vvl::ShaderObject can have NULL, independent layouts.
                // This never happens with pipelines because we get |merged_graphics_layout| that occurs during the final linking of
                // the pipeline which is creating a unified version for us.
                //
                // I used to have a check for |other.from_shader_object| but would get random failures
                //  (https://gitlab.khronos.org/Tracker/vk-gl-cts/-/issues/6512)
                // because the last DSL might have been created first with vkCreatePipelineLayout
                continue;
            }
            return false;
        }
    }
    return true;
}

// When comparing from draw time the |other| value is
//  - when using pipelines, the `vkCmdBindPipeline` layout
//  - when using shaderObject, the "fake" wrapper around VkShaderCreateInfoEXT::pSetLayouts
std::string PipelineLayoutCompatDef::DescribeDifference(const PipelineLayoutCompatDef& other) const {
    std::ostringstream ss;
    if (set != other.set) {
        assert(!other.from_shader_object);
        ss << "The set " << set << " is different from the non-compatible VkPipelineLayout (" << other.set << ")\n";
    } else if (push_constant_ranges != other.push_constant_ranges) {
        ss << "The VkPipelineLayout bound with last call to vkCmdBindDescriptorSets has following push constant ranges:\n";
        if (push_constant_ranges->empty()) {
            ss << "Empty\n";
        } else {
            for (const auto [pcr_i, pcr] : vvl::enumerate(push_constant_ranges->data(), push_constant_ranges->size())) {
                ss << "VkPushConstantRange[" << pcr_i << "]: " << string_VkPushConstantRange(pcr) << '\n';
            }
        }
        if (other.from_shader_object) {
            ss << "But the VkShaderCreateInfoEXT::pPushConstantRanges was created with the following push constant ranges:\n";
        } else {
            ss << "But the VkPipelineLayout layout of last vkCmdBindPipeline has following push constant ranges:\n";
        }
        if (other.push_constant_ranges->empty()) {
            ss << "Empty\n";
        } else {
            for (const auto [pcr_i, pcr] : vvl::enumerate(other.push_constant_ranges->data(), other.push_constant_ranges->size())) {
                ss << "VkPushConstantRange[" << pcr_i << "]: " << string_VkPushConstantRange(pcr) << '\n';
            }
        }
    } else if (is_independent_sets != other.is_independent_sets) {
        ss << "The VkPipelineLayout used to bind set " << set << " was created " << ((is_independent_sets) ? "with" : "without")
           << " VK_PIPELINE_LAYOUT_CREATE_INDEPENDENT_SETS_BIT_EXT when the ";

        if (other.from_shader_object) {
            ss << "VkShaderCreateInfoEXT::flags was";
            if (is_independent_sets) {
                ss << " not ";
            }
            ss << "created with VK_SHADER_CREATE_INDEPENDENT_SETS_BIT_KHR";
        } else {
            ss << "VkPipelineLayout of last bound pipeline was";
            if (is_independent_sets) {
                ss << " not";
            }
        }

    } else {
        const auto& descriptor_set_layouts = *set_layouts_id.get();
        const auto& other_ds_layouts = *other.set_layouts_id.get();
        for (uint32_t i = 0; i <= set; i++) {
            if (descriptor_set_layouts[i] != other_ds_layouts[i]) {
                if (!descriptor_set_layouts[i] || !other_ds_layouts[i]) {
                    ss << "Set layout " << i << " contains a null set which is considered non-compatible\n";
                    break;
                }
                return descriptor_set_layouts[i]->DescribeDifference(i, *other_ds_layouts[i]);
            }
        }
    }
    return ss.str();
}

static PipelineLayoutCompatId GetCanonicalId(const uint32_t set_index, const PushConstantRangesId& pcr_id,
                                             const PipelineLayoutSetLayoutsId& set_layouts_id, bool is_independent_sets,
                                             bool from_shader_object) {
    return pipeline_layout_compat_dict.LookUp(
        PipelineLayoutCompatDef(set_index, pcr_id, set_layouts_id, is_independent_sets, from_shader_object));
}

// For repeatable sorting, not very useful for "memory in range" search
struct PushConstantRangeCompare {
    bool operator()(const VkPushConstantRange* lhs, const VkPushConstantRange* rhs) const {
        if (lhs->offset == rhs->offset) {
            if (lhs->size == rhs->size) {
                // The comparison is arbitrary, but avoids false aliasing by comparing all fields.
                return lhs->stageFlags < rhs->stageFlags;
            }
            // If the offsets are the same then sorting by the end of range is useful for validation
            return lhs->size < rhs->size;
        }
        return lhs->offset < rhs->offset;
    }
};

PushConstantRangesId GetCanonicalId(uint32_t pushConstantRangeCount, const VkPushConstantRange* pPushConstantRanges) {
    if (!pPushConstantRanges) {
        // Hand back the empty entry (creating as needed)...
        return push_constant_ranges_dict.LookUp(PushConstantRanges());
    }

    // Sort the input ranges to ensure equivalent ranges map to the same id
    std::set<const VkPushConstantRange*, PushConstantRangeCompare> sorted;
    for (uint32_t i = 0; i < pushConstantRangeCount; i++) {
        sorted.insert(pPushConstantRanges + i);
    }

    PushConstantRanges ranges;
    ranges.reserve(sorted.size());
    for (const auto* range : sorted) {
        ranges.emplace_back(*range);
    }
    return push_constant_ranges_dict.LookUp(std::move(ranges));
}

static PushConstantRangesId GetPushConstantRangesFromLayouts(const vvl::span<const vvl::PipelineLayout* const>& layouts) {
    PushConstantRangesId ret{};
    for (const auto* layout : layouts) {
        if (layout && layout->push_constant_ranges_layout) {
            ret = layout->push_constant_ranges_layout;

            if (ret->size() > 0) {
                return ret;
            }
        }
    }
    return ret;
}

std::vector<PipelineLayoutCompatId> GetCompatForSet(const vvl::DescriptorSetLayoutList& set_layouts,
                                                    const PushConstantRangesId& push_constant_ranges, bool is_independent_sets,
                                                    bool from_shader_object) {
    PipelineLayoutSetLayoutsDef set_layout_ids(set_layouts.list.size());
    for (size_t i = 0; i < set_layouts.list.size(); i++) {
        if (set_layouts.list[i]) {
            set_layout_ids[i] = set_layouts.list[i]->GetLayoutId();
        }
    }
    auto set_layouts_id = pipeline_layout_set_layouts_dict.LookUp(set_layout_ids);

    std::vector<PipelineLayoutCompatId> set_compat_ids;
    set_compat_ids.reserve(set_layouts.list.size());

    for (uint32_t i = 0; i < set_layouts.list.size(); i++) {
        set_compat_ids.emplace_back(
            GetCanonicalId(i, push_constant_ranges, set_layouts_id, is_independent_sets, from_shader_object));
    }
    return set_compat_ids;
}

// This is called when merging the flags from the pipeline layouts in libraries
VkPipelineLayoutCreateFlags GetCreateFlags(const vvl::span<const vvl::PipelineLayout* const>& layouts) {
    // from https://github.com/KhronosGroup/Vulkan-ValidationLayers/issues/9870 and
    // https://gitlab.khronos.org/vulkan/vulkan/-/issues/4264
    // We do not actually combine the flags, instead we only take the flags from the final linked pipeline layout
    return layouts[0] ? layouts[0]->CreateFlags() : 0;
}

namespace vvl {

static DescriptorSetLayoutList GetSetLayouts(DeviceState& dev_data, const VkPipelineLayoutCreateInfo* pCreateInfo) {
    DescriptorSetLayoutList set_layouts(pCreateInfo->setLayoutCount);

    for (uint32_t i = 0; i < pCreateInfo->setLayoutCount; ++i) {
        set_layouts.list[i] = dev_data.Get<vvl::DescriptorSetLayout>(pCreateInfo->pSetLayouts[i]);
    }
    return set_layouts;
}

static DescriptorSetLayoutList GetSetLayouts(const vvl::span<const PipelineLayout* const>& layouts) {
    DescriptorSetLayoutList set_layouts;
    size_t num_layouts = 0;
    for (const auto& layout : layouts) {
        if (layout && (layout->set_layouts.list.size() > num_layouts)) {
            num_layouts = layout->set_layouts.list.size();
        }
    }

    set_layouts.list.reserve(num_layouts);
    for (size_t i = 0; i < num_layouts; ++i) {
        const PipelineLayout* used_layout = nullptr;
        for (const auto* layout : layouts) {
            if (layout) {
                if (layout->set_layouts.list.size() > i) {
                    // This _could_ be the layout we're looking for
                    used_layout = layout;

                    if (layout->set_layouts.list[i]) {
                        // This is the layout we're looking for. Any subsequent ones that match must be identical to this one.
                        break;
                    }
                }
            }
        }
        if (used_layout) {
            set_layouts.list.emplace_back(used_layout->set_layouts.list[i]);
        }
    }
    return set_layouts;
}

static bool HasDescriptorBuffer(const DescriptorSetLayoutList& set_layouts) {
    // It is valid to have null DSL (using GPL) so need to find the first valid
    for (const auto& set_layout : set_layouts.list) {
        // VU 08008 forces all layouts to have this flag, so only need to check first flag
        if (set_layout) {
            return (set_layout->GetCreateFlags() & VK_DESCRIPTOR_SET_LAYOUT_CREATE_DESCRIPTOR_BUFFER_BIT_EXT) != 0;
        }
    }
    return false;
}

static bool HasImmutableSamplers(const DescriptorSetLayoutList& set_layouts) {
    for (const auto& set_layout : set_layouts.list) {
        if (set_layout && set_layout->HasImmutableSamplers()) {
            return true;
        }
    }
    return false;
}

PipelineLayout::PipelineLayout(DeviceState& dev_data, VkPipelineLayout handle, const VkPipelineLayoutCreateInfo* pCreateInfo)
    : StateObject(handle, kVulkanObjectTypePipelineLayout),
      set_layouts(GetSetLayouts(dev_data, pCreateInfo)),
      push_constant_ranges_layout(GetCanonicalId(pCreateInfo->pushConstantRangeCount, pCreateInfo->pPushConstantRanges)),
      create_flags(pCreateInfo->flags),
      is_independent_set((create_flags & VK_PIPELINE_LAYOUT_CREATE_INDEPENDENT_SETS_BIT_EXT) != 0),
      has_descriptor_buffer(HasDescriptorBuffer(set_layouts)),
      has_immutable_samplers(HasImmutableSamplers(set_layouts)),
      set_compat_ids(GetCompatForSet(set_layouts, push_constant_ranges_layout, is_independent_set, false)) {}

PipelineLayout::PipelineLayout(const vvl::span<const PipelineLayout* const>& layouts)
    : StateObject(static_cast<VkPipelineLayout>(VK_NULL_HANDLE), kVulkanObjectTypePipelineLayout),
      set_layouts(GetSetLayouts(layouts)),
      push_constant_ranges_layout(GetPushConstantRangesFromLayouts(layouts)),  // TODO is this correct?
      create_flags(GetCreateFlags(layouts)),
      is_independent_set((create_flags & VK_PIPELINE_LAYOUT_CREATE_INDEPENDENT_SETS_BIT_EXT) != 0),
      has_descriptor_buffer(HasDescriptorBuffer(set_layouts)),
      has_immutable_samplers(HasImmutableSamplers(set_layouts)),
      set_compat_ids(GetCompatForSet(set_layouts, push_constant_ranges_layout, is_independent_set, false)) {}

}  // namespace vvl
