// 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/decompose_storage_access.h"

#include <utility>

#include "src/tint/lang/core/ir/builder.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"
#include "src/tint/lang/hlsl/ir/member_builtin_call.h"
#include "src/tint/lang/hlsl/type/byte_address_buffer.h"
#include "src/tint/lang/hlsl/type/matrix_layout.h"

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

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

/// PIMPL state for the transform.
struct State {
    /// The IR module.
    core::ir::Module& ir;

    /// The IR builder.
    core::ir::Builder b{ir};

    /// The type manager.
    core::type::Manager& ty{ir.Types()};

    using VarTypePair = std::pair<core::ir::Var*, const core::type::Type*>;
    /// Maps a struct to the load function
    Hashmap<VarTypePair, core::ir::Function*, 2> var_and_type_to_load_fn_{};
    /// Maps a struct to the store function
    Hashmap<VarTypePair, core::ir::Function*, 2> var_and_type_to_store_fn_{};

    /// Process the module.
    void Process() {
        Vector<core::ir::Var*, 4> var_worklist;
        for (auto* inst : *ir.root_block) {
            // Allow this to run before or after PromoteInitializers by handling non-var root_block
            // entries
            auto* var = inst->As<core::ir::Var>();
            if (!var) {
                continue;
            }

            // Var must be a pointer
            auto* var_ty = var->Result()->Type()->As<core::type::Pointer>();
            TINT_IR_ASSERT(ir, var_ty);

            // Only care about storage address space variables.
            if (var_ty->AddressSpace() != core::AddressSpace::kStorage) {
                continue;
            }

            var_worklist.Push(var);
        }

        for (auto* var : var_worklist) {
            auto* result = var->Result();

            ReplaceUses(result, var, {});

            // Swap the result type of the `var` to the new HLSL result type
            auto* var_ty = result->Type()->As<core::type::Pointer>();
            result->SetType(ty.Get<hlsl::type::ByteAddressBuffer>(var_ty->Access()));
        }
    }

    void BufferLength(core::ir::Var* var, core::ir::CoreBuiltinCall* call) {
        auto* buffer_ty = var->Result()->Type()->UnwrapPtr()->As<core::type::Buffer>();
        TINT_IR_ASSERT(
            ir,
            buffer_ty &&
                (buffer_ty->Count()
                     ->IsAnyOf<core::type::RuntimeArrayCount, core::type::ConstantArrayCount>()));

        if (call->Args().size() > 1) {
            // Length was directly encoded by previous passes
            call->Result()->ReplaceAllUsesWith(call->Args()[1]);
        } else if (auto* cnst = buffer_ty->Count()->As<core::type::ConstantArrayCount>()) {
            call->Result()->ReplaceAllUsesWith(b.Constant(u32(cnst->value)));
        } else {
            b.InsertBefore(call, [&] {
                // The `GetDimensions` call uses out parameters for all return values, there is no
                // return value. This ends up being the result value we care about.
                //
                // This creates a var with an access which means that when we emit the HLSL we'll
                // emit the correct `var` name.
                core::ir::Instruction* inst = b.Var(ty.ptr(function, ty.u32()));
                b.MemberCall<hlsl::ir::MemberBuiltinCall>(ty.void_(), BuiltinFn::kGetDimensions,
                                                          var, inst->Result());

                // Unlike arrayLength, bufferLength is just the size of the variable in bytes.
                inst = b.Load(inst);
                call->Result()->ReplaceAllUsesWith(inst->Result());
            });
        }
        call->Destroy();
    }

    struct OffsetData {
        uint32_t byte_offset = 0;
        Vector<core::ir::Value*, 4> byte_offset_expr{};

        uint32_t byte_struct_offset = 0;

        uint32_t byte_size = 0;
        core::ir::Value* byte_size_expr = nullptr;

        uint32_t byte_length = 0;
        core::ir::Value* byte_length_expr = nullptr;
    };

    void ArrayLength(core::ir::Var* var,
                     core::ir::CoreBuiltinCall* call,
                     const core::type::Type* type,
                     OffsetData offset) {
        auto* arr_ty = type->As<core::type::Array>();
        // If the `arrayLength` was called directly on the storage buffer then
        // it _must_ be a runtime array.
        TINT_IR_ASSERT(ir, arr_ty && arr_ty->Count()->As<core::type::RuntimeArrayCount>());

        b.InsertBefore(call, [&] {
            const bool has_size = HasSizeData(offset);
            const bool has_length = HasLengthData(offset);
            core::ir::Value* len = nullptr;
            if (has_size) {
                len = SizeToValue(offset);
            } else if (has_length) {
                len = LengthToValue(offset);
            } else {
                // The `GetDimensions` call uses out parameters for all return values, there is no
                // return value. This ends up being the result value we care about.
                //
                // This creates a var with an access which means that when we emit the HLSL we'll
                // emit the correct `var` name.
                core::ir::Instruction* inst = b.Var(ty.ptr(function, ty.u32()));
                b.MemberCall<hlsl::ir::MemberBuiltinCall>(ty.void_(), BuiltinFn::kGetDimensions,
                                                          var, inst->Result());

                len = b.Load(inst)->Result();
            }

            core::ir::Value* value = nullptr;
            TINT_IR_ASSERT(ir, offset.byte_offset_expr.Length() <= 1);
            if (has_size) {
                // BufferArrayView call preceded this length. Can't use accumulated offset because
                // in this case, but any struct member offset is still necessary.
                if (offset.byte_struct_offset != 0) {
                    value = b.Constant(u32(offset.byte_struct_offset));
                }
            } else {
                value = OffsetToValue(offset);
            }

            if (value) {
                auto* cnst = value->As<core::ir::Constant>();
                if (!cnst || cnst->Value()->ValueAs<uint32_t>() > 0) {
                    len = b.Subtract(len, value)->Result();
                }
            }

            auto* div = b.Divide(len, u32(arr_ty->ImplicitStride()));
            call->Result()->ReplaceAllUsesWith(div->Result());
        });
        call->Destroy();
    }

    void BufferView(core::ir::Var* var, core::ir::CoreBuiltinCall* call) {
        OffsetData offset;
        b.InsertBefore(call, [&] {
            // Offset is in bytes
            UpdateOffsetData(call->Args()[1], 1, &offset);
            if (call->Func() == core::BuiltinFn::kBufferArrayView) {
                UpdateSizeData(call->Args()[2], &offset);
                if (call->Args().size() > 3) {
                    UpdateLengthData(call->Args()[3], &offset);
                }
            } else if (call->Args().size() > 2) {
                UpdateLengthData(call->Args()[2], &offset);
            }
        });

        ReplaceUses(call->Result(), var, offset);
        call->Destroy();
    }

    void Interlocked(core::ir::Var* var,
                     core::ir::CoreBuiltinCall* call,
                     const OffsetData& offset,
                     BuiltinFn fn) {
        auto args = call->Args();
        auto* type = args[1]->Type();

        b.InsertBefore(call, [&] {
            auto* original_value = b.Var(ty.ptr(function, type));
            original_value->SetInitializer(b.Zero(type));

            b.MemberCall<hlsl::ir::MemberBuiltinCall>(ty.void_(), fn, var, OffsetToValue(offset),
                                                      args[1], original_value);
            b.LoadWithResult(call->DetachResult(), original_value);
        });
        call->Destroy();
    }

    void InterlockedStore(core::ir::Var* var,
                          core::ir::CoreBuiltinCall* call,
                          const OffsetData& offset,
                          BuiltinFn fn) {
        auto args = call->Args();
        auto* type = args[1]->Type();
        b.InsertBefore(call, [&] {
            TINT_ASSERT(type->Is<core::type::U64>());

            b.MemberCall<hlsl::ir::MemberBuiltinCall>(ty.void_(), fn, var, OffsetToValue(offset),
                                                      args[1]);
        });
        call->Destroy();
    }

    void AtomicAnd(core::ir::Var* var, core::ir::CoreBuiltinCall* call, const OffsetData& offset) {
        Interlocked(var, call, offset, BuiltinFn::kInterlockedAnd);
    }

    void AtomicOr(core::ir::Var* var, core::ir::CoreBuiltinCall* call, const OffsetData& offset) {
        Interlocked(var, call, offset, BuiltinFn::kInterlockedOr);
    }

    void AtomicXor(core::ir::Var* var, core::ir::CoreBuiltinCall* call, const OffsetData& offset) {
        Interlocked(var, call, offset, BuiltinFn::kInterlockedXor);
    }

    void AtomicMin(core::ir::Var* var, core::ir::CoreBuiltinCall* call, const OffsetData& offset) {
        Interlocked(var, call, offset, BuiltinFn::kInterlockedMin);
    }

    void AtomicMax(core::ir::Var* var, core::ir::CoreBuiltinCall* call, const OffsetData& offset) {
        Interlocked(var, call, offset, BuiltinFn::kInterlockedMax);
    }

    void AtomicStoreMin(core::ir::Var* var,
                        core::ir::CoreBuiltinCall* call,
                        const OffsetData& offset) {
        InterlockedStore(var, call, offset, BuiltinFn::kInterlockedMin64);
    }

    void AtomicStoreMax(core::ir::Var* var,
                        core::ir::CoreBuiltinCall* call,
                        const OffsetData& offset) {
        InterlockedStore(var, call, offset, BuiltinFn::kInterlockedMax64);
    }

    void AtomicAdd(core::ir::Var* var, core::ir::CoreBuiltinCall* call, const OffsetData& offset) {
        Interlocked(var, call, offset, BuiltinFn::kInterlockedAdd);
    }

    void AtomicExchange(core::ir::Var* var,
                        core::ir::CoreBuiltinCall* call,
                        const OffsetData& offset) {
        Interlocked(var, call, offset, BuiltinFn::kInterlockedExchange);
    }

    // An atomic sub is a negated atomic add
    void AtomicSub(core::ir::Var* var, core::ir::CoreBuiltinCall* call, const OffsetData& offset) {
        auto args = call->Args();
        auto* type = args[1]->Type();

        b.InsertBefore(call, [&] {
            auto* original_value = b.Var(ty.ptr(function, type));
            original_value->SetInitializer(b.Zero(type));

            auto* val = b.Subtract(b.Zero(type), args[1]);
            b.MemberCall<hlsl::ir::MemberBuiltinCall>(ty.void_(), BuiltinFn::kInterlockedAdd, var,
                                                      OffsetToValue(offset), val, original_value);
            b.LoadWithResult(call->DetachResult(), original_value);
        });
        call->Destroy();
    }

    void AtomicCompareExchangeWeak(core::ir::Var* var,
                                   core::ir::CoreBuiltinCall* call,
                                   const OffsetData& offset) {
        auto args = call->Args();
        auto* type = args[1]->Type();
        b.InsertBefore(call, [&] {
            auto* original_value = b.Var(ty.ptr(function, type));
            original_value->SetInitializer(b.Zero(type));

            auto* cmp = args[1];
            b.MemberCall<hlsl::ir::MemberBuiltinCall>(
                ty.void_(), BuiltinFn::kInterlockedCompareExchange, var, OffsetToValue(offset), cmp,
                args[2], original_value);

            auto* o = b.Load(original_value);
            b.ConstructWithResult(call->DetachResult(), o, b.Equal(o, cmp));
        });
        call->Destroy();
    }

    // An atomic load is an Or with 0
    void AtomicLoad(core::ir::Var* var, core::ir::CoreBuiltinCall* call, const OffsetData& offset) {
        auto* type = call->Result()->Type();
        b.InsertBefore(call, [&] {
            auto* original_value = b.Var(ty.ptr(function, type));
            original_value->SetInitializer(b.Zero(type));

            b.MemberCall<hlsl::ir::MemberBuiltinCall>(ty.void_(), BuiltinFn::kInterlockedOr, var,
                                                      OffsetToValue(offset), b.Zero(type),
                                                      original_value);
            b.LoadWithResult(call->DetachResult(), original_value);
        });
        call->Destroy();
    }

    void AtomicStore(core::ir::Var* var,
                     core::ir::CoreBuiltinCall* call,
                     const OffsetData& offset) {
        auto args = call->Args();
        auto* type = args[1]->Type();

        b.InsertBefore(call, [&] {
            auto* original_value = b.Var(ty.ptr(function, type));
            original_value->SetInitializer(b.Zero(type));

            b.MemberCall<hlsl::ir::MemberBuiltinCall>(ty.void_(), BuiltinFn::kInterlockedExchange,
                                                      var, OffsetToValue(offset), args[1],
                                                      original_value);
        });
        call->Destroy();
    }

    core::ir::Constant* ColMajorToMatrixLayout(core::ir::Value* col_major) {
        auto* const_col_major = col_major->As<core::ir::Constant>();
        TINT_IR_ASSERT(ir, const_col_major);
        return b.Constant(ir.constant_values.Get<core::constant::Scalar<u32>>(
            ty.Get<type::MatrixLayout>(),
            u32(const_col_major->Value()->ValueAs<bool>() ? type::MatrixLayoutEnum::kColMajor
                                                          : type::MatrixLayoutEnum::kRowMajor)));
    }

    void SubgroupMatrixLoad(core::ir::Var* var,
                            core::ir::CoreBuiltinCall* call,
                            OffsetData offset) {
        auto args = call->Args();
        auto* call_offset = args[1];
        auto* col_major = args[2];
        auto* stride = args[3];

        auto* sm = call->Result()->Type()->As<core::type::SubgroupMatrix>();
        TINT_IR_ASSERT(ir, sm);

        b.InsertBefore(call, [&] {
            UpdateOffsetData(call_offset, sm->Type()->Size(), &offset);

            auto* layout = ColMajorToMatrixLayout(col_major);
            // TODO(crbug.com/490062439): This will need to be updated if we change stride.
            uint32_t bytes_per_element = sm->Type()->Size();
            if (auto* cnst = stride->As<core::ir::Constant>()) {
                stride = b.Constant(u32(cnst->Value()->ValueAs<uint32_t>() * bytes_per_element));
            } else {
                stride = b.Multiply(stride, u32(bytes_per_element))->Result();
            }
            auto* load = b.CallExplicit<hlsl::ir::BuiltinCall>(
                sm, BuiltinFn::kLoad, Vector{sm}, var, OffsetToValue(offset), stride, layout);
            call->Result()->ReplaceAllUsesWith(load->Result());
        });
        call->Destroy();
    }

    void SubgroupMatrixStore(core::ir::Var* var,
                             core::ir::CoreBuiltinCall* call,
                             OffsetData offset) {
        auto args = call->Args();
        auto* call_offset = args[1];
        auto* value = args[2];
        auto* col_major = args[3];
        auto* stride = args[4];

        auto* sm = value->Type()->As<core::type::SubgroupMatrix>();
        TINT_IR_ASSERT(ir, sm);

        b.InsertBefore(call, [&] {
            UpdateOffsetData(call_offset, sm->Type()->Size(), &offset);

            auto* layout = ColMajorToMatrixLayout(col_major);

            // TODO(crbug.com/490062439): This will need to be updated if we change stride.
            uint32_t bytes_per_element = sm->Type()->Size();
            if (auto* cnst = stride->As<core::ir::Constant>()) {
                stride = b.Constant(u32(cnst->Value()->ValueAs<uint32_t>() * bytes_per_element));
            } else {
                stride = b.Multiply(stride, u32(bytes_per_element))->Result();
            }
            b.MemberCall<hlsl::ir::MemberBuiltinCall>(ty.void_(), BuiltinFn::kStore, value, var,
                                                      OffsetToValue(offset), stride, layout);
        });
        call->Destroy();
    }

    // Note, must be called inside a builder insert block (Append, InsertBefore, etc)
    void UpdateOffsetData(core::ir::Value* v, uint32_t elm_size, OffsetData* offset) {
        tint::Switch(
            v,  //
            [&](core::ir::Constant* idx_value) {
                offset->byte_offset += idx_value->Value()->ValueAs<uint32_t>() * elm_size;
            },
            [&](core::ir::Value* val) {
                auto* idx = b.InsertConvertIfNeeded(ty.u32(), val);
                offset->byte_offset_expr.Push(b.Multiply(idx, u32(elm_size))->Result());
            },
            TINT_ICE_ON_NO_MATCH);
    }

    // Note, must be called inside a builder insert block (Append, InsertBefore, etc)
    void UpdateSizeData(core::ir::Value* v, OffsetData* offset) {
        tint::Switch(
            v,  //
            [&](core::ir::Constant* idx_value) {
                offset->byte_size += idx_value->Value()->ValueAs<uint32_t>();
            },
            [&](core::ir::Value* val) {
                TINT_IR_ASSERT(ir, offset->byte_size_expr == nullptr);
                offset->byte_size_expr = b.InsertConvertIfNeeded(ty.u32(), val);
            },
            TINT_ICE_ON_NO_MATCH);
    }

    // Note, must be called inside a builder insert block (Append, InsertBefore, etc)
    void UpdateLengthData(core::ir::Value* v, OffsetData* offset) {
        tint::Switch(
            v,  //
            [&](core::ir::Constant* idx_value) {
                offset->byte_length += idx_value->Value()->ValueAs<uint32_t>();
            },
            [&](core::ir::Value* val) {
                TINT_IR_ASSERT(ir, offset->byte_length_expr == nullptr);
                offset->byte_length_expr = b.InsertConvertIfNeeded(ty.u32(), val);
            },
            TINT_ICE_ON_NO_MATCH);
    }

    // Note, must be called inside a builder insert block (Append, InsertBefore, etc)
    core::ir::Value* OffsetToValue(const OffsetData& offset) {
        core::ir::Value* val = b.Value(u32(offset.byte_offset));
        for (core::ir::Value* expr : offset.byte_offset_expr) {
            val = b.Add(val, expr)->Result();
        }
        return val;
    }

    // Note, must be called inside a builder insert block (Append, InsertBefore, etc)
    core::ir::Value* SizeToValue(const OffsetData& offset) {
        if (offset.byte_size_expr != nullptr) {
            TINT_IR_ASSERT(ir, offset.byte_size == 0);
            return offset.byte_size_expr;
        }
        return b.Constant(u32(offset.byte_size));
    }

    bool HasSizeData(const OffsetData& offset) {
        return offset.byte_size != 0 || offset.byte_size_expr != nullptr;
    }

    // Note, must be called inside a builder insert block (Append, InsertBefore, etc)
    core::ir::Value* LengthToValue(const OffsetData& offset) {
        if (offset.byte_length_expr != nullptr) {
            TINT_IR_ASSERT(ir, offset.byte_length == 0);
            return offset.byte_length_expr;
        }
        return b.Constant(u32(offset.byte_length));
    }

    bool HasLengthData(const OffsetData& offset) {
        return offset.byte_length != 0 || offset.byte_length_expr != nullptr;
    }

    // Creates the appropriate store instructions for the given result type.
    void MakeStore(core::ir::Instruction* inst,
                   core::ir::Var* var,
                   core::ir::Value* from,
                   core::ir::Value* offset) {
        auto* store_ty = from->Type();
        if (store_ty->IsNumericScalarOrVector()) {
            MakeScalarOrVectorStore(var, from, offset);
            return;
        }

        tint::Switch(
            from->Type(),  //
            [&](const core::type::Struct* s) {
                auto* fn = GetStoreFunctionFor(inst, var, s);
                b.Call(fn, offset, from);
            },
            [&](const core::type::Matrix* m) {
                auto* fn = GetStoreFunctionFor(inst, var, m);
                b.Call(fn, offset, from);
            },
            [&](const core::type::Array* a) {
                auto* fn = GetStoreFunctionFor(inst, var, a);
                b.Call(fn, offset, from);
            },
            TINT_ICE_ON_NO_MATCH);
    }

    // Creates a `v.Store{2,3,4} offset, value` call based on the provided type. The stored value is
    // bitcast to a `u32` (or `u32` vector as needed).
    //
    // This only works for `u32`, `i32`, `f32` and the vector sizes of those types.
    void MakeScalarOrVectorStore(core::ir::Var* var,
                                 core::ir::Value* from,
                                 core::ir::Value* offset) {
        bool is_f16 = from->Type()->DeepestElement()->Is<core::type::F16>();
        bool is_u16 = from->Type()->DeepestElement()->Is<core::type::U16>();

        const core::type::Type* cast_ty = ty.MatchWidth(ty.u32(), from->Type());
        auto fn =
            is_f16 ? BuiltinFn::kStoreF16 : (is_u16 ? BuiltinFn::kStoreU16 : BuiltinFn::kStore);
        if (auto* vec = from->Type()->As<core::type::Vector>()) {
            switch (vec->Width()) {
                case 2:
                    fn = is_f16 ? BuiltinFn::kStore2F16
                                : (is_u16 ? BuiltinFn::kStore2U16 : BuiltinFn::kStore2);
                    break;
                case 3:
                    fn = is_f16 ? BuiltinFn::kStore3F16
                                : (is_u16 ? BuiltinFn::kStore3U16 : BuiltinFn::kStore3);
                    break;
                case 4:
                    fn = is_f16 ? BuiltinFn::kStore4F16
                                : (is_u16 ? BuiltinFn::kStore4U16 : BuiltinFn::kStore4);
                    break;
                default:
                    TINT_IR_UNREACHABLE(ir);
            }
        }

        core::ir::Value* cast = nullptr;
        // The `f16` and `u16` types are not cast in a store as the store itself ends up templated.
        if (is_f16 || is_u16) {
            cast = from;
        } else {
            cast = b.Bitcast(cast_ty, from)->Result();
        }
        b.MemberCall<hlsl::ir::MemberBuiltinCall>(ty.void_(), fn, var, offset, cast);
    }

    // Creates the appropriate load instructions for the given result type.
    core::ir::Call* MakeLoad(core::ir::Instruction* inst,
                             core::ir::Var* var,
                             const core::type::Type* result_ty,
                             core::ir::Value* offset) {
        if (result_ty->IsNumericScalarOrVector()) {
            return MakeScalarOrVectorLoad(var, result_ty, offset);
        }

        return tint::Switch(
            result_ty,  //
            [&](const core::type::Struct* s) {
                auto* fn = GetLoadFunctionFor(inst, var, s);
                return b.Call(fn, offset);
            },
            [&](const core::type::Matrix* m) {
                auto* fn = GetLoadFunctionFor(inst, var, m);
                return b.Call(fn, offset);
            },  //
            [&](const core::type::Array* a) {
                auto* fn = GetLoadFunctionFor(inst, var, a);
                return b.Call(fn, offset);
            },  //
            TINT_ICE_ON_NO_MATCH);
    }

    // Creates a `v.Load{2,3,4} offset` call based on the provided type. The load returns a
    // `u32` or vector of `u32` and then a `bitcast` is done to get back to the desired type.
    //
    // This only works for `u32`, `i32`, `f16`, `f32` and the vector sizes of those types.
    //
    // The `f16` type is special in that `f16` uses a templated load in HLSL `Load<float16_t>`
    // and returns the correct type, so there is no bitcast.
    core::ir::Call* MakeScalarOrVectorLoad(core::ir::Var* var,
                                           const core::type::Type* result_ty,
                                           core::ir::Value* offset) {
        bool is_f16 = result_ty->DeepestElement()->Is<core::type::F16>();
        bool is_u16 = result_ty->DeepestElement()->Is<core::type::U16>();

        const core::type::Type* load_ty = nullptr;
        // `f16` and `u16` loads return their native type instead of a `u32`
        if (is_f16) {
            load_ty = ty.MatchWidth(ty.f16(), result_ty);
        } else if (is_u16) {
            load_ty = ty.MatchWidth(ty.u16(), result_ty);
        } else {
            load_ty = ty.MatchWidth(ty.u32(), result_ty);
        }

        auto fn = is_f16 ? BuiltinFn::kLoadF16 : (is_u16 ? BuiltinFn::kLoadU16 : BuiltinFn::kLoad);
        if (auto* v = result_ty->As<core::type::Vector>()) {
            switch (v->Width()) {
                case 2:
                    fn = is_f16 ? BuiltinFn::kLoad2F16
                                : (is_u16 ? BuiltinFn::kLoad2U16 : BuiltinFn::kLoad2);
                    break;
                case 3:
                    fn = is_f16 ? BuiltinFn::kLoad3F16
                                : (is_u16 ? BuiltinFn::kLoad3U16 : BuiltinFn::kLoad3);
                    break;
                case 4:
                    fn = is_f16 ? BuiltinFn::kLoad4F16
                                : (is_u16 ? BuiltinFn::kLoad4U16 : BuiltinFn::kLoad4);
                    break;
                default:
                    TINT_IR_UNREACHABLE(ir);
            }
        }

        auto* builtin = b.MemberCall<hlsl::ir::MemberBuiltinCall>(load_ty, fn, var, offset);
        core::ir::Call* res = nullptr;

        // Do not bitcast `f16` or `u16` conversions as they use templated Load instructions
        if (is_f16 || is_u16) {
            res = builtin;
        } else {
            res = b.Bitcast(result_ty, builtin->Result());
        }
        return res;
    }

    // Creates a load function for the given `var` and `struct` combination. Essentially creates
    // a function similar to:
    //
    // fn custom_load_S(offset: u32) {
    //   let a = <load S member 0>(offset + member 0 offset);
    //   let b = <load S member 1>(offset + member 1 offset);
    //   ...
    //   let z = <load S member last>(offset + member last offset);
    //   return S(a, b, ..., z);
    // }
    core::ir::Function* GetLoadFunctionFor(core::ir::Instruction* inst,
                                           core::ir::Var* var,
                                           const core::type::Struct* s) {
        return var_and_type_to_load_fn_.GetOrAdd(VarTypePair{var, s}, [&] {
            auto* p = b.FunctionParam("offset", ty.u32());
            auto* fn = b.Function(s);
            fn->SetParams({p});

            b.Append(fn->Block(), [&] {
                Vector<core::ir::Value*, 4> values;
                for (const auto* mem : s->Members()) {
                    values.Push(
                        MakeLoad(inst, var, mem->Type(), b.Add(p, u32(mem->Offset()))->Result())
                            ->Result());
                }

                b.Return(fn, b.Construct(s, values));
            });

            return fn;
        });
    }

    core::ir::Function* GetStoreFunctionFor(core::ir::Instruction* inst,
                                            core::ir::Var* var,
                                            const core::type::Struct* s) {
        return var_and_type_to_store_fn_.GetOrAdd(VarTypePair{var, s}, [&] {
            auto* p = b.FunctionParam("offset", ty.u32());
            auto* obj = b.FunctionParam("obj", s);
            auto* fn = b.Function(ty.void_());
            fn->SetParams({p, obj});

            b.Append(fn->Block(), [&] {
                for (const auto* mem : s->Members()) {
                    auto* from = b.Access(mem->Type(), obj, u32(mem->Index()));
                    MakeStore(inst, var, from->Result(), b.Add(p, u32(mem->Offset()))->Result());
                }

                b.Return(fn);
            });

            return fn;
        });
    }

    // Creates a load function for the given `var` and `matrix` combination. Essentially creates
    // a function similar to:
    //
    // fn custom_load_M(offset: u32) {
    //   let a = <load M column 1>(offset + (1 * ColumnStride));
    //   let b = <load M column 2>(offset + (2 * ColumnStride));
    //   ...
    //   let z = <load M column last>(offset + (last * ColumnStride));
    //   return M(a, b, ... z);
    // }
    core::ir::Function* GetLoadFunctionFor(core::ir::Instruction* inst,
                                           core::ir::Var* var,
                                           const core::type::Matrix* mat) {
        return var_and_type_to_load_fn_.GetOrAdd(VarTypePair{var, mat}, [&] {
            auto* p = b.FunctionParam("offset", ty.u32());
            auto* fn = b.Function(mat);
            fn->SetParams({p});

            b.Append(fn->Block(), [&] {
                Vector<core::ir::Value*, 4> values;
                for (size_t i = 0; i < mat->Columns(); ++i) {
                    auto* add = b.Add(p, u32(i * mat->ColumnStride()));
                    auto* load = MakeLoad(inst, var, mat->ColumnType(), add->Result());
                    values.Push(load->Result());
                }

                b.Return(fn, b.Construct(mat, values));
            });

            return fn;
        });
    }

    core::ir::Function* GetStoreFunctionFor(core::ir::Instruction* inst,
                                            core::ir::Var* var,
                                            const core::type::Matrix* mat) {
        return var_and_type_to_store_fn_.GetOrAdd(VarTypePair{var, mat}, [&] {
            auto* p = b.FunctionParam("offset", ty.u32());
            auto* obj = b.FunctionParam("obj", mat);
            auto* fn = b.Function(ty.void_());
            fn->SetParams({p, obj});

            b.Append(fn->Block(), [&] {
                Vector<core::ir::Value*, 4> values;
                for (size_t i = 0; i < mat->Columns(); ++i) {
                    auto* from = b.Access(mat->ColumnType(), obj, u32(i));
                    MakeStore(inst, var, from->Result(),
                              b.Add(p, u32(i * mat->ColumnStride()))->Result());
                }

                b.Return(fn);
            });

            return fn;
        });
    }

    // Creates a load function for the given `var` and `array` combination. Essentially creates
    // a function similar to:
    //
    // fn custom_load_A(offset: u32) {
    //   A a = A();
    //   u32 i = 0;
    //   loop {
    //     if (i >= A length) {
    //       break;
    //     }
    //     a[i] = <load array type>(offset + (i * A->ImplicitStride()));
    //     i = i + 1;
    //   }
    //   return a;
    // }
    core::ir::Function* GetLoadFunctionFor(core::ir::Instruction* inst,
                                           core::ir::Var* var,
                                           const core::type::Array* arr) {
        return var_and_type_to_load_fn_.GetOrAdd(VarTypePair{var, arr}, [&] {
            auto* p = b.FunctionParam("offset", ty.u32());
            auto* fn = b.Function(arr);
            fn->SetParams({p});

            b.Append(fn->Block(), [&] {
                auto* result_arr = b.Var<function>("a", b.Zero(arr));

                auto* count = arr->Count()->As<core::type::ConstantArrayCount>();
                TINT_IR_ASSERT(ir, count);

                b.LoopRange(0_u, u32(count->value), 1_u, [&](core::ir::Value* idx) {
                    auto* access = b.Access(ty.ptr<function>(arr->ElemType()), result_arr, idx);
                    auto* stride = b.Multiply(idx, u32(arr->ImplicitStride()));
                    auto* byte_offset = b.Add(p, stride);
                    b.Store(access, MakeLoad(inst, var, arr->ElemType(), byte_offset->Result()));
                });

                b.Return(fn, b.Load(result_arr));
            });

            return fn;
        });
    }

    core::ir::Function* GetStoreFunctionFor(core::ir::Instruction* inst,
                                            core::ir::Var* var,
                                            const core::type::Array* arr) {
        return var_and_type_to_store_fn_.GetOrAdd(VarTypePair{var, arr}, [&] {
            auto* p = b.FunctionParam("offset", ty.u32());
            auto* obj = b.FunctionParam("obj", arr);
            auto* fn = b.Function(ty.void_());
            fn->SetParams({p, obj});

            b.Append(fn->Block(), [&] {
                auto* count = arr->Count()->As<core::type::ConstantArrayCount>();
                TINT_IR_ASSERT(ir, count);

                b.LoopRange(0_u, u32(count->value), 1_u, [&](core::ir::Value* idx) {
                    auto* from = b.Access(arr->ElemType(), obj, idx);
                    auto* stride = b.Multiply(idx, u32(arr->ImplicitStride()));
                    auto* byte_offset = b.Add(p, stride);
                    MakeStore(inst, var, from->Result(), byte_offset->Result());
                });

                b.Return(fn);
            });

            return fn;
        });
    }

    void Access(core::ir::Access* a,
                core::ir::Var* var,
                const core::type::Type* obj,
                OffsetData offset) {
        // Note, because we recurse through the `access` helper, the object passed in isn't
        // necessarily the originating `var` object, but maybe a partially resolved access chain
        // object.
        if (auto* view = obj->As<core::type::MemoryView>()) {
            obj = view->StoreType();
        }

        for (auto* idx_value : a->Indices()) {
            tint::Switch(
                obj,  //
                [&](const core::type::Vector* v) {
                    b.InsertBefore(
                        a, [&] { UpdateOffsetData(idx_value, v->Type()->Size(), &offset); });
                    obj = v->Type();
                },
                [&](const core::type::Matrix* m) {
                    b.InsertBefore(
                        a, [&] { UpdateOffsetData(idx_value, m->ColumnStride(), &offset); });
                    obj = m->ColumnType();
                },
                [&](const core::type::Array* ary) {
                    b.InsertBefore(
                        a, [&] { UpdateOffsetData(idx_value, ary->ImplicitStride(), &offset); });
                    obj = ary->ElemType();
                },
                [&](const core::type::Struct* s) {
                    auto* cnst = idx_value->As<core::ir::Constant>();

                    // A struct index must be a constant
                    TINT_IR_ASSERT(ir, cnst);

                    uint32_t idx = cnst->Value()->ValueAs<uint32_t>();
                    auto* mem = s->Members()[idx];
                    offset.byte_offset += mem->Offset();
                    obj = mem->Type();
                    if (!obj->HasFixedFootprint()) {
                        offset.byte_struct_offset = mem->Offset();
                    }
                },
                TINT_ICE_ON_NO_MATCH);
        }

        ReplaceUses(a->Result(), var, offset);
        a->Destroy();
    }

    void ReplaceUses(core::ir::Value* value, core::ir::Var* var, OffsetData offset) {
        // Copy the usages into a vector so we can remove items from the hashset.
        auto usages = value->UsagesSorted();
        while (!usages.IsEmpty()) {
            auto usage = usages.Pop();
            tint::Switch(
                usage.instruction,
                [&](core::ir::Let* let) {
                    // The `let` is essentially an alias to the pointer value. So, add the `let`
                    // usages into the usage worklist, and replace the let with the pointer value.
                    for (auto& u : let->Result()->UsagesSorted()) {
                        usages.Push(u);
                    }
                    let->Result()->ReplaceAllUsesWith(value);
                    let->Destroy();
                },
                [&](core::ir::Access* a) {
                    // Treat an access chain as a continuation of the current offset.
                    Access(a, var, value->Type()->UnwrapPtr(), offset);
                },

                [&](core::ir::LoadVectorElement* lve) { LoadVectorElement(lve, var, offset); },
                [&](core::ir::Load* ld) { Load(ld, var, offset); },
                [&](core::ir::StoreVectorElement* sve) { StoreVectorElement(sve, var, offset); },
                [&](core::ir::Store* store) { Store(store, var, store->From(), offset); },

                [&](core::ir::CoreBuiltinCall* call) {
                    switch (call->Func()) {
                        case core::BuiltinFn::kArrayLength:
                            // If this pointer is being used in an `arrayLength` call then it _must_
                            // have resolved to a runtime array.
                            ArrayLength(var, call, value->Type()->UnwrapPtr(), offset);
                            break;
                        case core::BuiltinFn::kAtomicAnd:
                            AtomicAnd(var, call, offset);
                            break;
                        case core::BuiltinFn::kAtomicOr:
                            AtomicOr(var, call, offset);
                            break;
                        case core::BuiltinFn::kAtomicXor:
                            AtomicXor(var, call, offset);
                            break;
                        case core::BuiltinFn::kAtomicMin:
                            AtomicMin(var, call, offset);
                            break;
                        case core::BuiltinFn::kAtomicMax:
                            AtomicMax(var, call, offset);
                            break;
                        case core::BuiltinFn::kAtomicAdd:
                            AtomicAdd(var, call, offset);
                            break;
                        case core::BuiltinFn::kAtomicSub:
                            AtomicSub(var, call, offset);
                            break;
                        case core::BuiltinFn::kAtomicExchange:
                            AtomicExchange(var, call, offset);
                            break;
                        case core::BuiltinFn::kAtomicCompareExchangeWeak:
                            AtomicCompareExchangeWeak(var, call, offset);
                            break;
                        case core::BuiltinFn::kAtomicStoreMax:
                            AtomicStoreMax(var, call, offset);
                            break;
                        case core::BuiltinFn::kAtomicStoreMin:
                            AtomicStoreMin(var, call, offset);
                            break;
                        case core::BuiltinFn::kAtomicStore:
                            AtomicStore(var, call, offset);
                            break;
                        case core::BuiltinFn::kAtomicLoad:
                            AtomicLoad(var, call, offset);
                            break;
                        case core::BuiltinFn::kSubgroupMatrixStore:
                            SubgroupMatrixStore(var, call, offset);
                            break;
                        case core::BuiltinFn::kSubgroupMatrixLoad:
                            SubgroupMatrixLoad(var, call, offset);
                            break;
                        case core::BuiltinFn::kBufferLength:
                            BufferLength(var, call);
                            break;
                        case core::BuiltinFn::kBufferView:
                        case core::BuiltinFn::kBufferArrayView:
                            BufferView(var, call);
                            break;
                        default:
                            TINT_IR_UNREACHABLE(ir) << call->Func();
                    }
                },  //
                TINT_ICE_ON_NO_MATCH);
        }
    }

    void Store(core::ir::Instruction* inst,
               core::ir::Var* var,
               core::ir::Value* from,
               const OffsetData& offset) {
        b.InsertBefore(inst, [&] {
            auto* off = OffsetToValue(offset);
            MakeStore(inst, var, from, off);
        });
        inst->Destroy();
    }

    void Load(core::ir::Instruction* inst, core::ir::Var* var, const OffsetData& offset) {
        b.InsertBefore(inst, [&] {
            auto* off = OffsetToValue(offset);
            auto* call = MakeLoad(inst, var, inst->Result()->Type(), off);
            inst->Result()->ReplaceAllUsesWith(call->Result());
        });
        inst->Destroy();
    }

    // Converts to:
    //
    // %1:u32 = v.Load 0u
    // %b:f32 = bitcast %1
    void LoadVectorElement(core::ir::LoadVectorElement* lve,
                           core::ir::Var* var,
                           OffsetData offset) {
        b.InsertBefore(lve, [&] {
            UpdateOffsetData(lve->Index(), lve->Result()->Type()->Size(), &offset);

            auto* result =
                MakeScalarOrVectorLoad(var, lve->Result()->Type(), OffsetToValue(offset));
            lve->Result()->ReplaceAllUsesWith(result->Result());
        });

        lve->Destroy();
    }

    // Converts to:
    //
    // %1 = <sve->Value()>
    // %2:u32 = bitcast %1
    // %3:void = v.Store 0u, %2
    void StoreVectorElement(core::ir::StoreVectorElement* sve,
                            core::ir::Var* var,
                            OffsetData offset) {
        b.InsertBefore(sve, [&] {
            UpdateOffsetData(sve->Index(), sve->Value()->Type()->Size(), &offset);
            Store(sve, var, sve->Value(), offset);
        });
    }
};

}  // namespace

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

    State{ir}.Process();

    return Success;
}

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