// Copyright (c) 2026 Google LLC
//
// 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 "source/opt/legalize_multidim_array_pass.h"

#include "source/opt/constants.h"
#include "source/opt/desc_sroa_util.h"
#include "source/opt/ir_builder.h"
#include "source/opt/ir_context.h"
#include "source/opt/type_manager.h"

namespace spvtools {
namespace opt {

Pass::Status LegalizeMultidimArrayPass::Process() {
  std::vector<Instruction*> vars_to_legalize;

  for (auto& var : context()->types_values()) {
    if (var.opcode() != spv::Op::OpVariable) continue;
    if (!IsMultidimArrayOfResources(&var)) continue;
    if (!CanLegalize(&var)) {
      context()->EmitErrorMessage("Unable to legalize multidimensional array: ",
                                  &var);
      return Status::Failure;
    }
    vars_to_legalize.push_back(&var);
  }

  if (vars_to_legalize.empty()) return Status::SuccessWithoutChange;

  for (auto* var : vars_to_legalize) {
    uint32_t old_ptr_type_id = var->type_id();
    uint32_t new_ptr_type_id = FlattenArrayType(var);
    if (new_ptr_type_id == 0) return Status::Failure;
    if (!RewriteAccessChains(var, old_ptr_type_id)) return Status::Failure;
  }

  return Status::SuccessWithChange;
}

bool LegalizeMultidimArrayPass::IsMultidimArrayOfResources(Instruction* var) {
  if (!descsroautil::IsDescriptorArray(context(), var)) return false;

  uint32_t type_id = var->type_id();
  Instruction* type_inst = context()->get_def_use_mgr()->GetDef(type_id);
  uint32_t pointee_type_id = type_inst->GetSingleWordInOperand(1);
  std::vector<uint32_t> dims;
  uint32_t element_type_id = 0;
  GetArrayDimensions(pointee_type_id, &dims, &element_type_id);

  return dims.size() > 1;
}

void LegalizeMultidimArrayPass::GetArrayDimensions(uint32_t type_id,
                                                   std::vector<uint32_t>* dims,
                                                   uint32_t* element_type_id) {
  assert(dims != nullptr && "dims cannot be null.");
  dims->clear();

  Instruction* type_inst = context()->get_def_use_mgr()->GetDef(type_id);
  while (type_inst->opcode() == spv::Op::OpTypeArray) {
    uint32_t length_id = type_inst->GetSingleWordInOperand(1);
    Instruction* length_inst = context()->get_def_use_mgr()->GetDef(length_id);
    // Assume OpConstant. According to the spec the length could also be an
    // OpSpecConstantOp. However, DXC will not generate that type of code. The
    // code to handle spec constants will be much more complicated.
    assert(length_inst->opcode() == spv::Op::OpConstant);
    uint32_t length = length_inst->GetSingleWordInOperand(0);
    dims->push_back(length);
    type_id = type_inst->GetSingleWordInOperand(0);
    type_inst = context()->get_def_use_mgr()->GetDef(type_id);
  }
  *element_type_id = type_id;
}

uint32_t LegalizeMultidimArrayPass::FlattenArrayType(Instruction* var) {
  analysis::TypeManager* type_mgr = context()->get_type_mgr();
  analysis::ConstantManager* constant_mgr = context()->get_constant_mgr();

  uint32_t ptr_type_id = var->type_id();
  Instruction* ptr_type_inst =
      context()->get_def_use_mgr()->GetDef(ptr_type_id);
  uint32_t pointee_type_id = ptr_type_inst->GetSingleWordInOperand(1);

  std::vector<uint32_t> dims;
  uint32_t element_type_id = 0;
  GetArrayDimensions(pointee_type_id, &dims, &element_type_id);

  uint32_t total_elements = 1;
  for (uint32_t dim : dims) {
    total_elements *= dim;
  }

  const analysis::Constant* total_elements_const =
      constant_mgr->GetIntConst(total_elements, 32, false);

  Instruction* total_elements_inst =
      constant_mgr->GetDefiningInstruction(total_elements_const);
  uint32_t total_elements_id = total_elements_inst->result_id();

  // Create new OpTypeArray.
  analysis::Type* element_type = type_mgr->GetType(element_type_id);
  analysis::Array::LengthInfo length_info = {
      total_elements_id,
      {analysis::Array::LengthInfo::kConstant, total_elements}};
  analysis::Array new_array_type(element_type, length_info);
  uint32_t new_array_type_id = type_mgr->GetTypeInstruction(&new_array_type);

  // Create new OpTypePointer.
  spv::StorageClass sc =
      static_cast<spv::StorageClass>(ptr_type_inst->GetSingleWordInOperand(0));
  analysis::Pointer new_ptr_type(type_mgr->GetType(new_array_type_id), sc);
  uint32_t new_ptr_type_id = type_mgr->GetTypeInstruction(&new_ptr_type);

  var->SetResultType(new_ptr_type_id);
  context()->UpdateDefUse(var);

  // Move the var after the new pointer type to avoid a def-before-use.
  var->InsertAfter(get_def_use_mgr()->GetDef(new_ptr_type_id));

  return new_ptr_type_id;
}

bool LegalizeMultidimArrayPass::RewriteAccessChains(Instruction* var,
                                                    uint32_t old_ptr_type_id) {
  uint32_t var_id = var->result_id();
  std::vector<Instruction*> users;
  // Use a worklist to handle transitive uses (e.g. through OpCopyObject)
  std::vector<Instruction*> worklist;

  context()->get_def_use_mgr()->ForEachUser(
      var_id, [&worklist](Instruction* user) { worklist.push_back(user); });

  Instruction* old_ptr_type_inst =
      context()->get_def_use_mgr()->GetDef(old_ptr_type_id);
  uint32_t old_pointee_type_id = old_ptr_type_inst->GetSingleWordInOperand(1);
  std::vector<uint32_t> dims;
  uint32_t element_type_id = 0;
  GetArrayDimensions(old_pointee_type_id, &dims, &element_type_id);
  assert(dims.size() != 0 &&
         "This variable should have been rejected earlier.");

  // Calculate strides once
  std::vector<uint32_t> strides(dims.size());
  strides[dims.size() - 1] = 1;
  for (int i = static_cast<int>(dims.size()) - 2; i >= 0; --i) {
    strides[i] = strides[i + 1] * dims[i + 1];
  }

  // Pre-calculate uint type id
  uint32_t uint_type_id = context()->get_type_mgr()->GetUIntTypeId();
  if (uint_type_id == 0) return false;

  while (!worklist.empty()) {
    Instruction* user = worklist.back();
    worklist.pop_back();

    if (user->opcode() == spv::Op::OpAccessChain ||
        user->opcode() == spv::Op::OpInBoundsAccessChain) {
      uint32_t num_indices = user->NumInOperands() - 1;
      assert(num_indices >= dims.size());

      InstructionBuilder builder(context(), user, IRContext::kAnalysisDefUse);

      uint32_t linearized_idx_id = 0;
      for (uint32_t i = 0; i < dims.size(); ++i) {
        uint32_t idx_id = user->GetSingleWordInOperand(i + 1);

        uint32_t term_id = idx_id;
        if (strides[i] != 1) {
          const analysis::Constant* stride_const =
              context()->get_constant_mgr()->GetConstant(
                  context()->get_type_mgr()->GetType(uint_type_id),
                  {strides[i]});
          Instruction* stride_inst =
              context()->get_constant_mgr()->GetDefiningInstruction(
                  stride_const);

          Instruction* mul_inst = builder.AddBinaryOp(
              uint_type_id, spv::Op::OpIMul, idx_id, stride_inst->result_id());
          if (mul_inst == nullptr) return false;
          term_id = mul_inst->result_id();
        }

        if (linearized_idx_id == 0) {
          linearized_idx_id = term_id;
        } else {
          Instruction* add_inst = builder.AddBinaryOp(
              uint_type_id, spv::Op::OpIAdd, linearized_idx_id, term_id);
          if (add_inst == nullptr) return false;
          linearized_idx_id = add_inst->result_id();
        }
      }

      // Create new AccessChain.
      Instruction::OperandList new_operands;
      new_operands.push_back(user->GetInOperand(0));
      new_operands.push_back({SPV_OPERAND_TYPE_ID, {linearized_idx_id}});
      for (uint32_t i = static_cast<uint32_t>(dims.size()); i < num_indices;
           ++i) {
        new_operands.push_back(user->GetInOperand(i + 1));
      }
      user->SetInOperands(std::move(new_operands));
      context()->UpdateDefUse(user);
    } else if (user->opcode() == spv::Op::OpCopyObject) {
      // The type of the variable has changed so the result type of the
      // OpCopyObject will change as well.

      uint32_t operand_id = user->GetSingleWordInOperand(0);
      Instruction* operand_inst =
          context()->get_def_use_mgr()->GetDef(operand_id);
      user->SetResultType(operand_inst->type_id());
      context()->UpdateDefUse(user);

      // Add users of this copy to worklist
      context()->get_def_use_mgr()->ForEachUser(
          user->result_id(),
          [&worklist](Instruction* u) { worklist.push_back(u); });
    }
  }
  return true;
}

bool LegalizeMultidimArrayPass::CheckUse(Instruction* inst,
                                         uint32_t max_depth) {
  if (inst->opcode() == spv::Op::OpAccessChain ||
      inst->opcode() == spv::Op::OpInBoundsAccessChain) {
    uint32_t num_indices = inst->NumInOperands() - 1;
    return num_indices >= max_depth;
  } else if (inst->opcode() == spv::Op::OpCopyObject) {
    bool ok = true;
    return !context()->get_def_use_mgr()->WhileEachUser(
        inst->result_id(),
        [&](Instruction* u) { return !CheckUse(u, max_depth); });
    return ok;
  } else if (inst->IsDecoration() || inst->opcode() == spv::Op::OpName ||
             inst->opcode() == spv::Op::OpMemberName) {
    // Metadata is fine.
    return true;
  }

  // Direct use of array or partial array without AccessChain is not allowed.
  return false;
}

bool LegalizeMultidimArrayPass::CanLegalize(Instruction* var) {
  bool ok = true;
  uint32_t ptr_type_id = var->type_id();
  Instruction* ptr_type_inst =
      context()->get_def_use_mgr()->GetDef(ptr_type_id);
  uint32_t pointee_type_id = ptr_type_inst->GetSingleWordInOperand(1);
  std::vector<uint32_t> dims;
  uint32_t element_type_id = 0;
  GetArrayDimensions(pointee_type_id, &dims, &element_type_id);

  context()->get_def_use_mgr()->ForEachUser(
      var->result_id(), [&](Instruction* u) {
        if (!CheckUse(u, static_cast<uint32_t>(dims.size()))) ok = false;
      });
  return ok;
}

}  // namespace opt
}  // namespace spvtools