// Copyright 2024 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/tint/lang/hlsl/writer/raise/shader_io.h"

#include <algorithm>
#include <limits>
#include <memory>
#include <utility>

#include "src/tint/lang/core/ir/builder.h"
#include "src/tint/lang/core/ir/module.h"
#include "src/tint/lang/core/ir/transform/shader_io.h"
#include "src/tint/lang/core/ir/validator.h"
#include "src/tint/lang/hlsl/builtin_fn.h"
#include "src/tint/lang/hlsl/ir/builtin_call.h"

using namespace tint::core::fluent_types;     // NOLINT
using namespace tint::core::number_suffixes;  // NOLINT

namespace tint::hlsl::writer::raise {

namespace {

/// PIMPL state for the parts of the shader IO transform specific to HLSL.
/// For HLSL, move all inputs to a struct passed as an entry point parameter, and wrap outputs in
/// a structure returned by the entry point.
struct StateImpl : core::ir::transform::ShaderIOBackendState {
    /// The config
    const ShaderIOConfig& config;

    /// The input parameter
    core::ir::FunctionParam* input_param = nullptr;

    Vector<uint32_t, 4> input_indices;
    Vector<uint32_t, 4> output_indices;

    /// The output struct type.
    core::type::Struct* output_struct = nullptr;

    /// The output values to return from the entry point.
    Vector<core::ir::Value*, 4> output_values;

    // Indices of inputs that require special handling
    std::optional<uint32_t> subgroup_id_index;
    std::optional<uint32_t> subgroup_invocation_id_index;
    std::optional<uint32_t> subgroup_size_index;
    std::optional<uint32_t> num_subgroups_index;
    std::optional<uint32_t> num_workgroups_index;
    std::optional<uint32_t> sample_mask_index;
    std::optional<uint32_t> sample_index_index;
    std::optional<uint32_t> local_invocation_index_index;
    std::optional<uint32_t> first_clip_distance_index;
    std::optional<uint32_t> second_clip_distance_index;
    std::optional<uint32_t> global_invocation_id_index;
    std::optional<uint32_t> global_invocation_index_index;
    std::optional<uint32_t> workgroup_index_index;
    std::optional<uint32_t> workgroup_id_index;
    Hashset<uint32_t, 4> truncated_indices;

    // If set, points to a var of type struct with fields for offsets to apply to vertex_index and
    // instance_index
    core::ir::Var* tint_first_index_offset = nullptr;

    // If set, holds a workgroup atomic var that is used to polyfill subgroup_id and num_subgroups.
    core::ir::Var* tint_subgroup_id_counter = nullptr;
    // If set, holds a function var that contains the subgroup_id value.
    core::ir::Var* tint_subgroup_id = nullptr;
    // If set, holds a function var that contains the num_workgroups value.
    core::ir::Value* tint_num_workgroups = nullptr;
    // If the entry point has a linear workgroup size, this will hold that linearized size.
    std::optional<uint32_t> linear_workgroup_size = std::nullopt;

    /// Constructor
    StateImpl(core::ir::Module& mod, core::ir::Function* f, const ShaderIOConfig& c)
        : ShaderIOBackendState(mod, f), config(c) {
        // Check if we have a linear workgroup size.
        if (auto wgsize = func->WorkgroupSizeAsConst()) {
            workgroup_size = wgsize;
            linear_workgroup_size = 1;
            for (uint32_t i = 0; i < 3; i++) {
                if (wgsize->at(i) > 1u) {
                    if (linear_workgroup_size > 1u) {
                        // We have already seen dimension with size > 1, so the size is not linear.
                        linear_workgroup_size = std::nullopt;
                        break;
                    }
                    linear_workgroup_size = wgsize->at(i);
                }
            }
        }
    }

    /// Destructor
    ~StateImpl() override {}

    /// FXC is sensitive to field order in structures, this is used by StructMemberComparator to
    /// ensure that FXC is happy with the order of emitted fields.
    uint32_t BuiltinOrder(core::BuiltinValue builtin) {
        switch (builtin) {
            case core::BuiltinValue::kPosition:
                return 1;
            case core::BuiltinValue::kVertexIndex:
                return 2;
            case core::BuiltinValue::kInstanceIndex:
                return 3;
            case core::BuiltinValue::kFrontFacing:
                return 4;
            case core::BuiltinValue::kFragDepth:
                return 5;
            case core::BuiltinValue::kLocalInvocationId:
                return 6;
            case core::BuiltinValue::kLocalInvocationIndex:
                return 7;
            case core::BuiltinValue::kGlobalInvocationId:
                return 8;
            case core::BuiltinValue::kWorkgroupId:
                return 9;
            case core::BuiltinValue::kNumWorkgroups:
                return 10;
            case core::BuiltinValue::kSampleIndex:
                return 11;
            case core::BuiltinValue::kSampleMask:
                return 12;
            case core::BuiltinValue::kPointSize:
                return 13;
            case core::BuiltinValue::kClipDistances:
                return 14;
            case core::BuiltinValue::kPrimitiveIndex:
                return 15;
            case core::BuiltinValue::kBarycentricCoord:
                return 16;
            case core::BuiltinValue::kSubgroupId:
            case core::BuiltinValue::kSubgroupInvocationId:
            case core::BuiltinValue::kSubgroupSize:
            case core::BuiltinValue::kNumSubgroups:
            case core::BuiltinValue::kGlobalInvocationIndex:
            case core::BuiltinValue::kWorkgroupIndex:
                // These are sorted, but don't actually end up as members. Value doesn't really
                // matter, so just make it larger than the rest.
                return std::numeric_limits<uint32_t>::max();
            default:
                break;
        }
        TINT_IR_UNREACHABLE(ir) << "Unhandled builtin value: " << ToString(builtin);
    }

    struct MemberInfo {
        core::type::Manager::StructMemberDesc member;
        uint32_t idx;
    };

    /// Comparison function used to reorder struct members such that all members with
    /// color attributes appear first (ordered by color slot), then location attributes (ordered by
    /// location slot), then blend_src attributes (ordered by blend_src slot), followed by those
    /// with builtin attributes (ordered by BuiltinOrder).
    /// @param x a struct member
    /// @param y another struct member
    /// @returns true if a comes before b
    bool StructMemberComparator(const MemberInfo& x, const MemberInfo& y) {
        if (x.member.attributes.color.has_value() && y.member.attributes.color.has_value() &&
            x.member.attributes.color != y.member.attributes.color) {
            // Both have color attributes: smallest goes first.
            return x.member.attributes.color < y.member.attributes.color;
        } else if (x.member.attributes.color.has_value() != y.member.attributes.color.has_value()) {
            // The member with the color goes first
            return x.member.attributes.color.has_value();
        }

        if (x.member.attributes.location.has_value() && y.member.attributes.location.has_value() &&
            x.member.attributes.location != y.member.attributes.location) {
            // Both have location attributes: smallest goes first.
            return x.member.attributes.location < y.member.attributes.location;
        } else if (x.member.attributes.location.has_value() !=
                   y.member.attributes.location.has_value()) {
            // The member with the location goes first
            return x.member.attributes.location.has_value();
        }

        if (x.member.attributes.blend_src.has_value() &&
            y.member.attributes.blend_src.has_value() &&
            x.member.attributes.blend_src != y.member.attributes.blend_src) {
            // Both have blend_src attributes: smallest goes first.
            return x.member.attributes.blend_src < y.member.attributes.blend_src;
        } else if (x.member.attributes.blend_src.has_value() !=
                   y.member.attributes.blend_src.has_value()) {
            // The member with the blend_src goes first
            return x.member.attributes.blend_src.has_value();
        }

        auto x_blt = x.member.attributes.builtin;
        auto y_blt = y.member.attributes.builtin;
        if (x_blt.has_value() && y_blt.has_value()) {
            // Both are builtins: order matters for FXC.
            auto order_a = BuiltinOrder(*x_blt);
            auto order_b = BuiltinOrder(*y_blt);
            if (order_a != order_b) {
                return order_a < order_b;
            }
        } else if (x_blt.has_value() != y_blt.has_value()) {
            // The member with the builtin goes first
            return x_blt.has_value();
        }

        // Control flow reaches here if x is the same as y.
        // Sort algorithms sometimes do that.
        return false;
    }

    /// @copydoc ShaderIO::BackendState::FinalizeInputs
    Result<Vector<core::ir::FunctionParam*, 4>> FinalizeInputs() override {
        if (config.add_input_position_member) {
            RequireBuiltinInput(core::BuiltinValue::kPosition, ty.vec4f(), "pos");
        }

        // Check if we need to add certain builtins for polyfills:
        // * local_invocation_index for global_invocation_index or subgroup_id (with linear
        // workgroup size)
        // * workgroup_id for workgroup_index or global_invocation_index
        // * num_workgroups for workgroup_index or global_invocation_index
        const bool has_subgroup_id = HasBuiltinInput(core::BuiltinValue::kSubgroupId);
        const bool has_global_invocation_index =
            HasBuiltinInput(core::BuiltinValue::kGlobalInvocationIndex);
        const bool has_workgroup_index = HasBuiltinInput(core::BuiltinValue::kWorkgroupIndex);
        const bool needs_local_invocation_index = has_subgroup_id && linear_workgroup_size;
        if (needs_local_invocation_index) {
            RequireBuiltinInput(core::BuiltinValue::kLocalInvocationIndex, ty.u32(),
                                "local_invocation_index");
        }
        const bool needs_workgroup_id = has_workgroup_index;
        if (needs_workgroup_id) {
            RequireBuiltinInput(core::BuiltinValue::kWorkgroupId, ty.vec3u(), "workgroup_id");
        }
        const bool needs_num_workgroups = has_workgroup_index || has_global_invocation_index;
        if (needs_num_workgroups) {
            RequireBuiltinInput(core::BuiltinValue::kNumWorkgroups, ty.vec3u(), "num_workgroups");
        }
        const bool needs_global_invocation_id = has_global_invocation_index;
        if (needs_global_invocation_id) {
            RequireBuiltinInput(core::BuiltinValue::kGlobalInvocationId, ty.vec3u(),
                                "global_invocation_id");
        }
        if (config.polyfill_sample_mask &&
            func->Stage() == core::ir::Function::PipelineStage::kFragment &&
            HasBuiltinInput(core::BuiltinValue::kSampleMask)) {
            RequireBuiltinInput(core::BuiltinValue::kSampleIndex, ty.u32(), "sample_index");
        }

        Vector<MemberInfo, 4> input_data;
        bool has_vertex_or_instance_index = false;
        for (uint32_t i = 0; i < inputs.Length(); ++i) {
            // Save the index of certain builtins for GetIndex. Although struct members will not be
            // added for these inputs, we still add entries to input_data so that other inputs with
            // struct members can index input_indices properly in GetIndex.
            if (auto builtin = inputs[i].attributes.builtin) {
                if (*builtin == core::BuiltinValue::kSubgroupId) {
                    subgroup_id_index = i;
                } else if (*builtin == core::BuiltinValue::kSubgroupInvocationId) {
                    subgroup_invocation_id_index = i;
                } else if (*builtin == core::BuiltinValue::kSubgroupSize) {
                    subgroup_size_index = i;
                } else if (*builtin == core::BuiltinValue::kNumSubgroups) {
                    num_subgroups_index = i;
                } else if (*builtin == core::BuiltinValue::kNumWorkgroups) {
                    num_workgroups_index = i;
                } else if (*builtin == core::BuiltinValue::kSampleMask) {
                    sample_mask_index = i;
                } else if (*builtin == core::BuiltinValue::kSampleIndex) {
                    sample_index_index = i;
                } else if (*builtin == core::BuiltinValue::kLocalInvocationIndex) {
                    local_invocation_index_index = i;
                } else if (*builtin == core::BuiltinValue::kVertexIndex) {
                    has_vertex_or_instance_index = true;
                } else if (*builtin == core::BuiltinValue::kInstanceIndex) {
                    has_vertex_or_instance_index = true;
                } else if (*builtin == core::BuiltinValue::kGlobalInvocationIndex) {
                    global_invocation_index_index = i;
                } else if (*builtin == core::BuiltinValue::kWorkgroupIndex) {
                    workgroup_index_index = i;
                } else if (*builtin == core::BuiltinValue::kWorkgroupId) {
                    workgroup_id_index = i;
                } else if (*builtin == core::BuiltinValue::kGlobalInvocationId) {
                    global_invocation_id_index = i;
                }
            }

            input_data.Push(MemberInfo{inputs[i], i});
        }

        input_indices.Resize(input_data.Length());

        if (config.first_index_offset_binding.has_value() && has_vertex_or_instance_index) {
            // Create a FirstIndexOffset uniform buffer. GetInput will use this to offset the
            // vertex/instance index.
            TINT_IR_ASSERT(ir, func->IsVertex());
            tint::Vector<tint::core::type::Manager::StructMemberDesc, 2> members;
            auto* str = ty.Struct(ir.symbols.New("tint_first_index_offset_struct"),
                                  {
                                      {ir.symbols.New("vertex_index"), ty.u32(), {}},
                                      {ir.symbols.New("instance_index"), ty.u32(), {}},
                                  });
            tint_first_index_offset = b.Var("tint_first_index_offset", uniform, str);
            tint_first_index_offset->SetBindingPoint(config.first_index_offset_binding->group,
                                                     config.first_index_offset_binding->binding);
            ir.root_block->Append(tint_first_index_offset);
        }

        // Sort the struct members to satisfy HLSL interfacing matching rules.
        // We use stable_sort so that two members with the same attributes maintain their relative
        // ordering (e.g. kClipDistance).
        std::stable_sort(input_data.begin(), input_data.end(),
                         [&](auto& x, auto& y) { return StructMemberComparator(x, y); });

        Vector<core::type::Manager::StructMemberDesc, 4> input_struct_members;
        for (auto& input : input_data) {
            // Don't add members for certain builtins
            if (input.idx == subgroup_id_index ||              //
                input.idx == subgroup_invocation_id_index ||   //
                input.idx == subgroup_size_index ||            //
                input.idx == num_subgroups_index ||            //
                input.idx == num_workgroups_index ||           //
                input.idx == global_invocation_index_index ||  //
                input.idx == workgroup_index_index) {
                // Invalid value, should not be indexed
                input_indices[input.idx] = std::numeric_limits<uint32_t>::max();
                continue;
            }
            input_indices[input.idx] = static_cast<uint32_t>(input_struct_members.Length());
            input_struct_members.Push(input.member);
        }

        if (!input_struct_members.IsEmpty()) {
            auto* input_struct = ty.Struct(ir.symbols.New(ir.NameOf(func).Name() + "_inputs"),
                                           std::move(input_struct_members));
            switch (func->Stage()) {
                case core::ir::Function::PipelineStage::kFragment:
                    input_struct->AddUsage(core::type::PipelineStageUsage::kFragmentInput);
                    break;
                case core::ir::Function::PipelineStage::kVertex:
                    input_struct->AddUsage(core::type::PipelineStageUsage::kVertexInput);
                    break;
                case core::ir::Function::PipelineStage::kCompute:
                    input_struct->AddUsage(core::type::PipelineStageUsage::kComputeInput);
                    break;
                case core::ir::Function::PipelineStage::kUndefined:
                    TINT_IR_UNREACHABLE(ir);
            }
            input_param = b.FunctionParam("inputs", input_struct);
            return Vector<core::ir::FunctionParam*, 4>{input_param};
        }

        return Vector<core::ir::FunctionParam*, 4>{};
    }

    /// @copydoc ShaderIO::BackendState::FinalizeOutputs
    Result<const core::type::Type*> FinalizeOutputs() override {
        if (outputs.IsEmpty()) {
            return ty.void_();
        }

        // If a clip_distances output is found, replace it with either one or two new outputs,
        // depending on the array size. The new outputs maintain the ClipDistance attribute, which
        // is translated to SV_ClipDistanceN in the printer.
        for (uint32_t i = 0; i < outputs.Length(); ++i) {
            if (outputs[i].attributes.builtin == core::BuiltinValue::kClipDistances) {
                auto* const type = outputs[i].type;
                auto const name = outputs[i].name;
                auto const attributes = outputs[i].attributes;
                // Compute new member element counts
                auto* arr = type->As<core::type::Array>();
                uint32_t arr_count = *arr->ConstantCount();
                TINT_IR_ASSERT(ir, arr_count >= 1 && arr_count <= 8);
                uint32_t count0, count1;
                if (arr_count >= 4) {
                    count0 = 4;
                    count1 = arr_count - 4;
                } else {
                    count0 = arr_count;
                    count1 = 0;
                }
                // Replace current output for the first one
                auto* ty0 = ty.MatchWidth(ty.f32(), count0);
                auto name0 = ir.symbols.New(name.Name() + std::to_string(0));
                outputs[i] = {name0, ty0, attributes};
                first_clip_distance_index = i;
                // And add a new output for the second, if any
                if (count1 > 0) {
                    auto* ty1 = ty.MatchWidth(ty.f32(), count1);
                    auto name1 = ir.symbols.New(name.Name() + std::to_string(1));
                    second_clip_distance_index = AddOutput(name1, ty1, attributes);
                }
                break;
            }
        }

        Vector<MemberInfo, 4> output_data;
        for (uint32_t i = 0; i < outputs.Length(); ++i) {
            output_data.Push(MemberInfo{outputs[i], i});
        }

        // Sort the struct members to satisfy HLSL interfacing matching rules.
        // We use stable_sort so that two members with the same attributes maintain their relative
        // ordering (e.g. kClipDistance).
        std::stable_sort(output_data.begin(), output_data.end(),
                         [&](auto& x, auto& y) { return StructMemberComparator(x, y); });

        output_indices.Resize(outputs.Length());
        output_values.Resize(outputs.Length(), nullptr);

        Vector<core::type::Manager::StructMemberDesc, 4> output_struct_members;
        for (size_t i = 0; i < output_data.Length(); ++i) {
            output_indices[output_data[i].idx] = static_cast<uint32_t>(i);

            // If we need to truncate this member, don't add it to the output struct
            if (config.truncate_interstage_variables) {
                if (auto loc = output_data[i].member.attributes.location) {
                    if (!config.interstage_locations.test(*loc)) {
                        truncated_indices.Add(output_data[i].idx);
                        continue;
                    }
                }
            }

            output_struct_members.Push(output_data[i].member);
        }
        if (output_struct_members.IsEmpty()) {
            // All members were truncated
            return ty.void_();
        }

        output_struct =
            ty.Struct(ir.symbols.New(ir.NameOf(func).Name() + "_outputs"), output_struct_members);
        switch (func->Stage()) {
            case core::ir::Function::PipelineStage::kFragment:
                output_struct->AddUsage(core::type::PipelineStageUsage::kFragmentOutput);
                break;
            case core::ir::Function::PipelineStage::kVertex:
                output_struct->AddUsage(core::type::PipelineStageUsage::kVertexOutput);
                break;
            case core::ir::Function::PipelineStage::kCompute:
                output_struct->AddUsage(core::type::PipelineStageUsage::kComputeOutput);
                break;
            case core::ir::Function::PipelineStage::kUndefined:
                TINT_IR_UNREACHABLE(ir);
        }
        return output_struct;
    }

    /// Handles kNumWorkgroups builtin by emitting a UBO to hold the num_workgroups value,
    /// along with the load of the value. Returns the loaded value.
    core::ir::Value* GetInputForNumWorkgroups(core::ir::Builder& builder) {
        if (tint_num_workgroups) {
            return tint_num_workgroups;
        }

        // Create uniform var that will receive the number of workgroups
        core::ir::Var* num_wg_var = nullptr;
        builder.Append(ir.root_block, [&] {
            num_wg_var = builder.Var("tint_num_workgroups", ty.ptr(uniform, ty.vec3u()));
        });
        if (config.num_workgroups_binding.has_value()) {
            // If config.num_workgroups_binding holds a value, use it.
            auto bp = *config.num_workgroups_binding;
            num_wg_var->SetBindingPoint(bp.group, bp.binding);
        } else {
            // Otherwise, use the binding 0 of the largest used group plus 1, or group 0 if no
            // resources are bound.
            uint32_t group = 0;
            for (auto* inst : *ir.root_block) {
                if (auto* var = inst->As<core::ir::Var>()) {
                    if (const auto& bp = var->BindingPoint()) {
                        if (bp->group >= group) {
                            group = bp->group + 1;
                        }
                    }
                }
            }
            num_wg_var->SetBindingPoint(group, 0);
        }
        auto* load = builder.Load(num_wg_var);
        tint_num_workgroups = load->Result();
        return tint_num_workgroups;
    }

    /// Create the atomic counter used to polyfill subgroup_id and num_subgroups, and add the code
    /// to increment it from each subgroup.
    void MakeSubgroupIdCounter(core::ir::Builder& builder) {
        if (tint_subgroup_id_counter != nullptr) {
            return;
        }
        TINT_IR_ASSERT(ir, tint_subgroup_id == nullptr);

        // The polyfills for subgroup_id and num_subgroups look like this:
        //
        //   var<workgroup> tint_subgroup_id_counter: atomic<u32>;
        //
        //   @compute @workgroup_size(...)
        //   fn main(...) {
        //     // Initialize the counter to zero.
        //     atomicStore(&tint_subgroup_id_counter, 0u);
        //     workgroupBarrier();
        //
        //     // One invocation from each subgroup increments the counter to produce an ID.
        //     var tint_subgroup_id: u32;
        //     if (subgroup_invocation_id == 0) {
        //       tint_subgroup_id = atomicAdd(&tint_subgroup_id_counter, 1u);
        //     }
        //
        //     // Broadcast the resulting ID to the rest of the invocations in the subgroup.
        //     tint_subgroup_id = subgroupBroadcastFirst(tint_subgroup_id);
        //
        //     // Load the final counter value to get the number of subgroups.
        //     workgroupBarrier();
        //     let tint_num_subgroups = atomicLoad(&tint_subgroup_id_counter);
        //   }

        // Declare the atomic counter.
        builder.Append(ir.root_block, [&] {
            tint_subgroup_id_counter =
                builder.Var<workgroup, atomic<u32>>("tint_subgroup_id_counter");
        });

        // Initialize the counter to zero.
        builder.Call<void>(core::BuiltinFn::kAtomicStore, tint_subgroup_id_counter, 0_u);
        builder.Call<void>(core::BuiltinFn::kWorkgroupBarrier);

        // One invocation from each subgroup increments the counter to produce an ID.
        tint_subgroup_id = builder.Var<function, u32>("tint_subgroup_id");
        auto* first_invocation = builder.If(builder.Call<bool>(core::BuiltinFn::kSubgroupElect));
        builder.Append(first_invocation->True(), [&] {
            builder.Store(tint_subgroup_id, builder.Call<u32>(core::BuiltinFn::kAtomicAdd,
                                                              tint_subgroup_id_counter, 1_u));
            builder.ExitIf(first_invocation);
        });
    }

    core::ir::Value* PolyfillSubgroupId(core::ir::Builder& builder) {
        // If the workgroup size is linear, we can just divide the local invocation index by the
        // subgroup size.
        if (linear_workgroup_size) {
            TINT_IR_ASSERT(ir, local_invocation_index_index);
            auto* local_index = GetInput(builder, local_invocation_index_index.value());
            auto* subgroup_size = GetSubgroupSize(builder);
            return builder.Divide(local_index, subgroup_size)->Result();
        }

        // Otherwise, we use the atomic counter that is incremented by each subgroup, and then
        // broadcast the resulting subgroup ID to the rest of the invocations in the subgroup.
        MakeSubgroupIdCounter(builder);
        auto* id = builder.Load(tint_subgroup_id);
        return builder.Call<u32>(core::BuiltinFn::kSubgroupBroadcastFirst, id)->Result();
    }

    core::ir::Value* PolyfillNumSubgroups(core::ir::Builder& builder) {
        // If the workgroup size is linear, then we can just divide the workgroup size by the
        // subgroup size (rounding up).
        if (linear_workgroup_size) {
            auto* subgroup_size = GetSubgroupSize(builder);
            auto* add = builder.Add(u32(linear_workgroup_size.value() - 1), subgroup_size);
            return builder.Divide(add, subgroup_size)->Result();
        }

        // Otherwise, we use the atomic counter that is incremented by each subgroup, and then load
        // the final value from the atomic which will tell us how many subgroups there are.
        MakeSubgroupIdCounter(builder);
        TINT_IR_ASSERT(ir, tint_subgroup_id_counter);
        builder.Call<void>(core::BuiltinFn::kWorkgroupBarrier);
        return builder.Call<u32>(core::BuiltinFn::kAtomicLoad, tint_subgroup_id_counter)->Result();
    }

    core::ir::Value* GetSubgroupSize(core::ir::Builder& builder) {
        return builder.Call<hlsl::ir::BuiltinCall>(ty.u32(), hlsl::BuiltinFn::kWaveGetLaneCount)
            ->Result();
    }

    /// @copydoc ShaderIO::BackendState::GetInput
    core::ir::Value* GetInput(core::ir::Builder& builder, uint32_t idx) override {
        if (subgroup_id_index == idx) {
            return PolyfillSubgroupId(builder);
        }
        if (num_subgroups_index == idx) {
            return PolyfillNumSubgroups(builder);
        }
        if (global_invocation_index_index == idx) {
            return PolyfillGlobalInvocationIndex(builder, global_invocation_id_index.value(),
                                                 num_workgroups_index.value());
        }
        if (workgroup_index_index == idx) {
            return PolyfillWorkgroupIndex(builder, workgroup_id_index.value(),
                                          num_workgroups_index.value());
        }
        if (subgroup_invocation_id_index == idx) {
            return builder
                .Call<hlsl::ir::BuiltinCall>(ty.u32(), hlsl::BuiltinFn::kWaveGetLaneIndex)
                ->Result();
        }
        if (subgroup_size_index == idx) {
            return GetSubgroupSize(builder);
        }
        if (num_workgroups_index == idx) {
            if (config.num_workgroups_start_offset.has_value()) {
                auto* immediate_data = config.immediate_data_layout.var;
                auto num_workgroup_idx = u32(config.immediate_data_layout.IndexOf(
                    config.num_workgroups_start_offset.value()));
                auto* load = builder.Load(
                    builder.Access<ptr<immediate, vec3<u32>>>(immediate_data, num_workgroup_idx));
                return load->Result();
            }
            return GetInputForNumWorkgroups(builder);
        }

        auto index = input_indices[idx];
        auto* v = builder.Access(inputs[idx].type, input_param, u32(index))->Result();

        if (config.polyfill_sample_mask && sample_mask_index == idx) {
            // In Sample Shading mode only the sample for the fragment should be set in the sample
            // mask. The behavior of HLSL is to include the full raster sample mask. This is not in
            // line with WGSL spec and so we remove this additional information via a bitwise and.
            TINT_IR_ASSERT(ir, sample_index_index.has_value());
            auto* sample_index = GetInput(builder, sample_index_index.value());
            auto* mask = builder.ShiftLeft(1_u, sample_index);
            return builder.And(v, mask)->Result();
        }

        if (inputs[idx].attributes.builtin == core::BuiltinValue::kPosition) {
            // If this is an input position builtin we need to invert the 'w' component of the
            // vector.
            auto* w = builder.Access(ty.f32(), v, 3_u);
            auto* div = builder.Divide(1.0_f, w);
            auto* swizzle = builder.Swizzle(ty.vec3f(), v, {0, 1, 2});
            return builder.Construct(ty.vec4f(), swizzle, div)->Result();
        }

        if (config.first_index_offset_binding.has_value() &&
            inputs[idx].attributes.builtin == core::BuiltinValue::kVertexIndex) {
            // Apply vertex_index offset
            TINT_IR_ASSERT(ir, tint_first_index_offset);
            auto* vertex_index_offset =
                builder.Access(ty.ptr<uniform, u32>(), tint_first_index_offset, 0_u);
            return builder.Add(v, builder.Load(vertex_index_offset))->Result();
        }

        if (config.first_index_offset_binding.has_value() &&
            inputs[idx].attributes.builtin == core::BuiltinValue::kInstanceIndex) {
            // Apply instance_index offset
            TINT_IR_ASSERT(ir, tint_first_index_offset);
            auto* instance_index_offset =
                builder.Access(ty.ptr<uniform, u32>(), tint_first_index_offset, 1_u);
            return builder.Add(v, builder.Load(instance_index_offset))->Result();
        }

        if (config.first_index_offset.has_value() &&
            inputs[idx].attributes.builtin == core::BuiltinValue::kVertexIndex) {
            auto* immediate_data = config.immediate_data_layout.var;
            auto first_index_offset_idx =
                u32(config.immediate_data_layout.IndexOf(config.first_index_offset.value()));
            auto first_index_offset =
                builder.Access<ptr<immediate, u32>>(immediate_data, first_index_offset_idx);
            return builder.Add(v, builder.Load(first_index_offset))->Result();
        }

        if (config.first_instance_offset.has_value() &&
            inputs[idx].attributes.builtin == core::BuiltinValue::kInstanceIndex) {
            auto* immediate_data = config.immediate_data_layout.var;
            auto first_instance_offset_idx =
                u32(config.immediate_data_layout.IndexOf(config.first_instance_offset.value()));
            auto first_instance_offset =
                builder.Access<ptr<immediate, u32>>(immediate_data, first_instance_offset_idx);
            return builder.Add(v, builder.Load(first_instance_offset))->Result();
        }

        return v;
    }

    core::ir::Value* BuildClipDistanceInitValue(core::ir::Builder& builder,
                                                uint32_t output_index,
                                                core::ir::Value* src_array,
                                                uint32_t src_array_first_index) {
        // Copy from the array `src_array`
        auto* src_array_ty = src_array->Type()->As<core::type::Array>();
        TINT_IR_ASSERT(ir, src_array_ty);

        core::ir::Value* dst_value;
        if (auto* dst_vec_ty = outputs[output_index].type->As<core::type::Vector>()) {
            // Create a vector and copy array elements to it
            Vector<core::ir::Value*, 4> init;
            for (size_t i = 0; i < dst_vec_ty->Elements().count; ++i) {
                init.Push(builder.Access<f32>(src_array, u32(src_array_first_index + i))->Result());
            }
            dst_value = builder.Construct(dst_vec_ty, std::move(init))->Result();
        } else {
            TINT_IR_ASSERT(ir, outputs[output_index].type->As<core::type::Scalar>());
            dst_value = builder.Access<f32>(src_array, u32(src_array_first_index))->Result();
        }
        return dst_value;
    }

    /// @copydoc ShaderIO::BackendState::SetOutput
    void SetOutput(core::ir::Builder& builder, uint32_t idx, core::ir::Value* value) override {
        if (truncated_indices.Contains(idx)) {
            // Leave this output value as nullptr
            TINT_IR_ASSERT(ir, !output_values[output_indices[idx]]);
            return;
        }

        // If setting a ClipDistance output, build the initial value for one or both output values.
        if (idx == first_clip_distance_index) {
            auto index = output_indices[idx];
            output_values[index] = BuildClipDistanceInitValue(builder, idx, value, 0);

            if (second_clip_distance_index) {
                index = output_indices[*second_clip_distance_index];
                output_values[index] =
                    BuildClipDistanceInitValue(builder, *second_clip_distance_index, value, 4);
            }
            return;
        }

        auto index = output_indices[idx];
        output_values[index] = value;
    }

    /// @copydoc ShaderIO::BackendState::MakeReturnValue
    core::ir::Value* MakeReturnValue(core::ir::Builder& builder) override {
        if (!output_struct) {
            return nullptr;
        }

        if (truncated_indices.Count()) {
            // Remove all truncated values, which are nullptr in output_values
            output_values.EraseIf([](auto* v) { return v == nullptr; });
        }

        TINT_IR_ASSERT(ir, output_values.Length() == output_struct->Members().Length());
        return builder.Construct(output_struct, std::move(output_values))->Result();
    }
};
}  // namespace

Result<SuccessType> ShaderIO(core::ir::Module& ir, const ShaderIOConfig& config) {
    core::ir::AssertValid(ir,
                          core::ir::Capabilities{core::ir::Capability::kAllowDuplicateBindings,
                                                 core::ir::Capability::kAllow8BitIntegers,
                                                 core::ir::Capability::kAllow16BitIntegers},
                          "before hlsl.ShaderIO");

    TINT_CHECK_RESULT(core::ir::transform::RunShaderIOBase(
        ir, [&](core::ir::Module& mod, core::ir::Function* func) {
            return std::make_unique<StateImpl>(mod, func, config);
        }));

    return Success;
}

}  // namespace tint::hlsl::writer::raise
