// Copyright (c) 2023-2025 Arm Ltd.

//

// 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.

// Validates correctness of graph instructions.

#include <deque>

#include "source/opcode.h"

#include "source/spirv_target_env.h"

#include "source/val/validate.h"

#include "source/val/validation_state.h"

namespace spvtools {

namespace val {

namespace {

bool IsTensorArray(ValidationState_t& _, uint32_t id) {

  auto def = _.FindDef(id);

  if (!def || (def->opcode() != spv::Op::OpTypeArray &&

               def->opcode() != spv::Op::OpTypeRuntimeArray)) {

    return false;

  }

  auto tdef = _.FindDef(def->word(2));

  if (!tdef || tdef->opcode() != spv::Op::OpTypeTensorARM) {

    return false;

  }

  return true;

}

bool IsGraphInterfaceType(ValidationState_t& _, uint32_t id) {

  return _.IsTensorType(id) || IsTensorArray(_, id);

}

bool IsGraph(ValidationState_t& _, uint32_t id) {

  auto def = _.FindDef(id);

  if (!def || def->opcode() != spv::Op::OpGraphARM) {

    return false;

  }

  return true;

}

bool IsGraphType(ValidationState_t& _, uint32_t id) {

  auto def = _.FindDef(id);

  if (!def || def->opcode() != spv::Op::OpTypeGraphARM) {

    return false;

  }

  return true;

}

bool IsConstantInstruction(ValidationState_t& _, uint32_t id) {

  auto def = _.FindDef(id);

  return def && spvOpcodeIsConstant(def->opcode());

}

const uint32_t kGraphTypeIOStartWord = 3;

uint32_t GraphTypeInstNumIO(const Instruction* inst) {

  return static_cast<uint32_t>(inst->words().size()) - kGraphTypeIOStartWord;

}

uint32_t GraphTypeInstNumInputs(const Instruction* inst) {

  return inst->word(2);

}

uint32_t GraphTypeInstNumOutputs(const Instruction* inst) {

  return GraphTypeInstNumIO(inst) - GraphTypeInstNumInputs(inst);

}

uint32_t GraphTypeInstGetOutputAtIndex(const Instruction* inst,

                                       uint64_t index) {

  return inst->word(kGraphTypeIOStartWord + GraphTypeInstNumInputs(inst) +

                    static_cast<uint32_t>(index));

}

uint32_t GraphTypeInstGetInputAtIndex(const Instruction* inst, uint64_t index) {

  return inst->word(kGraphTypeIOStartWord + static_cast<uint32_t>(index));

}

spv_result_t ValidateGraphType(ValidationState_t& _, const Instruction* inst) {

  // Check there are at least NumInputs types

  uint32_t NumInputs = GraphTypeInstNumInputs(inst);

  size_t NumIOTypes = GraphTypeInstNumIO(inst);

  if (NumIOTypes < NumInputs) {

    return _.diag(SPV_ERROR_INVALID_DATA, inst)

           << NumIOTypes << " I/O types were provided but the graph has "

           << NumInputs << " inputs.";

  }

  // Check there is at least one output

  if (NumIOTypes == NumInputs) {

    return _.diag(SPV_ERROR_INVALID_DATA, inst)

           << "A graph type must have at least one output.";

  }

  // Check all I/O types are graph interface type

  for (unsigned i = kGraphTypeIOStartWord; i < inst->words().size(); i++) {

    auto tid = inst->word(i);

    if (!IsGraphInterfaceType(_, tid)) {

      return _.diag(SPV_ERROR_INVALID_DATA, inst)

             << "I/O type " << _.getIdName(tid)

             << " is not a Graph Interface Type.";

    }

  }

  return SPV_SUCCESS;

}

spv_result_t ValidateGraphConstant(ValidationState_t& _,

                                   const Instruction* inst) {

  // Check Result Type

  if (!_.IsTensorType(inst->type_id())) {

    return _.diag(SPV_ERROR_INVALID_DATA, inst)

           << spvOpcodeString(inst->opcode())

           << " must have a Result Type that is a tensor type.";

  }

  // Check the instruction is not preceded by another OpGraphConstantARM with

  // the same ID

  const uint32_t cst_id = inst->word(3);

  size_t inst_num = inst->LineNum() - 1;

  while (--inst_num) {

    auto prev_inst = &_.ordered_instructions()[inst_num];

    if (prev_inst->opcode() == spv::Op::OpGraphConstantARM) {

      const uint32_t prev_cst_id = prev_inst->word(3);

      if (prev_cst_id == cst_id) {

        return _.diag(SPV_ERROR_INVALID_DATA, inst)

               << "No two OpGraphConstantARM instructions may have the same "

                  "GraphConstantID";

      }

    }

  }

  return SPV_SUCCESS;

}

spv_result_t ValidateGraphEntryPoint(ValidationState_t& _,

                                     const Instruction* inst) {

  // Graph must be an OpGraphARM

  uint32_t graph = inst->GetOperandAs<uint32_t>(0);

  auto graph_inst = _.FindDef(graph);

  if (!IsGraph(_, graph)) {

    return _.diag(SPV_ERROR_INVALID_DATA, inst)

           << spvOpcodeString(inst->opcode())

           << " Graph must be a OpGraphARM but found "

           << spvOpcodeString(graph_inst->opcode()) << ".";

  }

  // Check number of Interface IDs matches number of I/Os of graph

  auto graph_type_inst = _.FindDef(graph_inst->type_id());

  size_t graph_type_num_io = GraphTypeInstNumIO(graph_type_inst);

  size_t graph_entry_point_num_interface_id = inst->operands().size() - 2;

  if (graph_type_inst->opcode() != spv::Op::OpTypeGraphARM) {

    // This is invalid but we want ValidateGraph to report a clear error

    // so stop validating the graph entry point instruction

    return SPV_SUCCESS;

  }

  if (graph_type_num_io != graph_entry_point_num_interface_id) {

    return _.diag(SPV_ERROR_INVALID_DATA, inst)

           << spvOpcodeString(inst->opcode()) << " Interface list contains "

           << graph_entry_point_num_interface_id << " IDs but Graph's type "

           << _.getIdName(graph_inst->type_id()) << " has " << graph_type_num_io

           << " inputs and outputs.";

  }

  // Check Interface IDs

  for (uint32_t i = 2; i < inst->operands().size(); i++) {

    uint32_t interface_id = inst->GetOperandAs<uint32_t>(i);

    auto interface_inst = _.FindDef(interface_id);

    // Check interface IDs come from OpVariable

    if ((interface_inst->opcode() != spv::Op::OpVariable) ||

        (interface_inst->GetOperandAs<spv::StorageClass>(2) !=

         spv::StorageClass::UniformConstant)) {

      return _.diag(SPV_ERROR_INVALID_DATA, interface_inst)

             << spvOpcodeString(inst->opcode()) << " Interface ID "

             << _.getIdName(interface_id)

             << " must come from OpVariable with UniformConstant Storage "

                "Class.";

    }

    // Check type of interface variable matches type of the corresponding graph

    // I/O

    uint32_t corresponding_graph_io_type =

        graph_type_inst->GetOperandAs<uint32_t>(i);

    uint32_t interface_ptr_type = interface_inst->type_id();

    auto interface_ptr_inst = _.FindDef(interface_ptr_type);

    auto interface_pointee_type = interface_ptr_inst->GetOperandAs<uint32_t>(2);

    if (interface_pointee_type != corresponding_graph_io_type) {

      return _.diag(SPV_ERROR_INVALID_DATA, inst)

             << spvOpcodeString(inst->opcode()) << " Interface ID type "

             << _.getIdName(interface_pointee_type)

             << " must match the type of the corresponding graph I/O "

             << _.getIdName(corresponding_graph_io_type);

    }

  }

  return SPV_SUCCESS;

}

spv_result_t ValidateGraph(ValidationState_t& _, const Instruction* inst) {

  // Result Type must be an OpTypeGraphARM

  if (!IsGraphType(_, inst->type_id())) {

    return _.diag(SPV_ERROR_INVALID_DATA, inst)

           << spvOpcodeString(inst->opcode())

           << " Result Type must be an OpTypeGraphARM.";

  }

  return SPV_SUCCESS;

}

spv_result_t ValidateGraphInput(ValidationState_t& _, const Instruction* inst) {

  const uint32_t input_index_id = inst->GetOperandAs<uint32_t>(2);

  // Check type of InputIndex

  auto input_index_inst = _.FindDef(input_index_id);

  if (!input_index_inst ||

      !_.IsIntScalarType(input_index_inst->type_id(), 32)) {

    return _.diag(SPV_ERROR_INVALID_DATA, inst)

           << spvOpcodeString(inst->opcode())

           << " InputIndex must be a 32-bit integer.";

  }

  bool has_element_index = inst->operands().size() > 3;

  // Check type of ElementIndex

  uint32_t element_index_id = 0;

  if (has_element_index) {

    element_index_id = inst->GetOperandAs<uint32_t>(3);

    auto element_index_inst = _.FindDef(element_index_id);

    if (!element_index_inst ||

        !_.IsIntScalarType(element_index_inst->type_id(), 32)) {

      return _.diag(SPV_ERROR_INVALID_DATA, inst)

             << spvOpcodeString(inst->opcode())

             << " ElementIndex must be a 32-bit integer.";

    }

  }

  if (spvIsVulkanEnv(_.context()->target_env)) {

    if (!IsConstantInstruction(_, input_index_id)) {

      return _.diag(SPV_ERROR_INVALID_DATA, inst)

             << _.VkErrorID(9931) << "OpGraphInputARM InputIndex must be the "

             << "<id> of a constant instruction.";

    }

    if (has_element_index && !IsConstantInstruction(_, element_index_id)) {

      return _.diag(SPV_ERROR_INVALID_DATA, inst)

             << _.VkErrorID(9931) << "OpGraphInputARM ElementIndex must be "

             << "the <id> of a constant instruction.";

    }

  }

  // Find graph definition

  size_t inst_num = inst->LineNum() - 1;

  auto graph_inst = &_.ordered_instructions()[inst_num];

  while (--inst_num) {

    graph_inst = &_.ordered_instructions()[inst_num];

    if (graph_inst->opcode() == spv::Op::OpGraphARM) {

      break;

    }

  }

  // Can the InputIndex be evaluated?

  // If not, there's nothing more we can validate here.

  uint64_t input_index;

  if (!_.EvalConstantValUint64(inst->GetOperandAs<uint32_t>(2), &input_index)) {

    return SPV_SUCCESS;

  }

  auto const graph_type_inst = _.FindDef(graph_inst->type_id());

  size_t graph_type_num_inputs = graph_type_inst->GetOperandAs<uint32_t>(1);

  // Check InputIndex is in range

  if (input_index >= graph_type_num_inputs) {

    std::string disassembly = _.Disassemble(*inst);

    return _.diag(SPV_ERROR_INVALID_DATA, nullptr)

           << "Type " << _.getIdName(graph_type_inst->id()) << " for graph "

           << _.getIdName(graph_inst->id()) << " has " << graph_type_num_inputs

           << " inputs but found an OpGraphInputARM instruction with an "

              "InputIndex that is "

           << input_index << ": " << disassembly;

  }

  uint32_t graph_type_input_type =

      GraphTypeInstGetInputAtIndex(graph_type_inst, input_index);

  if (has_element_index) {

    // Check ElementIndex is allowed

    if (!IsTensorArray(_, graph_type_input_type)) {

      return _.diag(SPV_ERROR_INVALID_DATA, inst)

             << "OpGraphInputARM ElementIndex not allowed when the graph input "

                "selected by "

             << "InputIndex is not an OpTypeArray or OpTypeRuntimeArray";

    }

    // Check ElementIndex is in range if it can be evaluated and the input is a

    // fixed-sized array whose Length can be evaluated

    uint64_t element_index;

    if (_.IsArrayType(graph_type_input_type) &&

        _.EvalConstantValUint64(inst->GetOperandAs<uint32_t>(3),

                                &element_index)) {

      uint64_t array_length;

      auto graph_type_input_type_inst = _.FindDef(graph_type_input_type);

      if (_.EvalConstantValUint64(

              graph_type_input_type_inst->GetOperandAs<uint32_t>(2),

              &array_length)) {

        if (element_index >= array_length) {

          return _.diag(SPV_ERROR_INVALID_DATA, inst)

                 << "OpGraphInputARM ElementIndex out of range. The type of "

                    "the graph input being accessed "

                 << _.getIdName(graph_type_input_type) << " is an array of "

                 << array_length << " elements but " << "ElementIndex is "

                 << element_index;

        }

      }

    }

  }

  // Check result type matches with graph type

  if (has_element_index) {

    uint32_t expected_type = _.GetComponentType(graph_type_input_type);

    if (inst->type_id() != expected_type) {

      return _.diag(SPV_ERROR_INVALID_DATA, inst)

             << "Result Type " << _.getIdName(inst->type_id())

             << " of graph input instruction " << _.getIdName(inst->id())

             << " does not match the component type "

             << _.getIdName(expected_type) << " of input " << input_index

             << " in the graph type.";

    }

  } else {

    if (inst->type_id() != graph_type_input_type) {

      return _.diag(SPV_ERROR_INVALID_DATA, inst)

             << "Result Type " << _.getIdName(inst->type_id())

             << " of graph input instruction " << _.getIdName(inst->id())

             << " does not match the type "

             << _.getIdName(graph_type_input_type) << " of input "

             << input_index << " in the graph type.";

    }

  }

  return SPV_SUCCESS;

}

spv_result_t ValidateGraphSetOutput(ValidationState_t& _,

                                    const Instruction* inst) {

  const uint32_t output_index_id = inst->GetOperandAs<uint32_t>(1);

  // Check type of OutputIndex

  auto output_index_inst = _.FindDef(output_index_id);

  if (!output_index_inst ||

      !_.IsIntScalarType(output_index_inst->type_id(), 32)) {

    return _.diag(SPV_ERROR_INVALID_DATA, inst)

           << spvOpcodeString(inst->opcode())

           << " OutputIndex must be a 32-bit integer.";

  }

  bool has_element_index = inst->operands().size() > 2;

  // Check type of ElementIndex

  uint32_t element_index_id = 0;

  if (has_element_index) {

    element_index_id = inst->GetOperandAs<uint32_t>(2);

    auto element_index_inst = _.FindDef(element_index_id);

    if (!element_index_inst ||

        !_.IsIntScalarType(element_index_inst->type_id(), 32)) {

      return _.diag(SPV_ERROR_INVALID_DATA, inst)

             << spvOpcodeString(inst->opcode())

             << " ElementIndex must be a 32-bit integer.";

    }

  }

  if (spvIsVulkanEnv(_.context()->target_env)) {

    if (!IsConstantInstruction(_, output_index_id)) {

      return _.diag(SPV_ERROR_INVALID_DATA, inst)

             << _.VkErrorID(9932) << "OpGraphSetOutputARM OutputIndex must "

             << "be the <id> of a constant instruction.";

    }

    if (has_element_index && !IsConstantInstruction(_, element_index_id)) {

      return _.diag(SPV_ERROR_INVALID_DATA, inst)

             << _.VkErrorID(9932) << "OpGraphSetOutputARM ElementIndex must "

             << "be the <id> of a constant instruction.";

    }

  }

  // Find graph definition

  size_t inst_num = inst->LineNum() - 1;

  auto graph_inst = &_.ordered_instructions()[inst_num];

  while (--inst_num) {

    graph_inst = &_.ordered_instructions()[inst_num];

    if (graph_inst->opcode() == spv::Op::OpGraphARM) {

      break;

    }

  }

  // Can the OutputIndex be evaluated?

  // If not, there's nothing more we can validate here.

  uint64_t output_index;

  if (!_.EvalConstantValUint64(inst->GetOperandAs<uint32_t>(1),

                               &output_index)) {

    return SPV_SUCCESS;

  }

  // Check that the OutputIndex is valid with respect to the graph type

  auto graph_type_inst = _.FindDef(graph_inst->type_id());

  size_t graph_type_num_outputs = GraphTypeInstNumOutputs(graph_type_inst);

  if (output_index >= graph_type_num_outputs) {

    return _.diag(SPV_ERROR_INVALID_DATA, inst)

           << spvOpcodeString(inst->opcode()) << " setting OutputIndex "

           << output_index << " but graph only has " << graph_type_num_outputs

           << " outputs.";

  }

  uint32_t graph_type_output_type =

      GraphTypeInstGetOutputAtIndex(graph_type_inst, output_index);

  if (has_element_index) {

    // Check ElementIndex is allowed

    if (!IsTensorArray(_, graph_type_output_type)) {

      return _.diag(SPV_ERROR_INVALID_DATA, inst)

             << "OpGraphSetOutputARM ElementIndex not allowed when the graph "

                "output selected by "

             << "OutputIndex is not an OpTypeArray or OpTypeRuntimeArray";

    }

    // Check ElementIndex is in range if it can be evaluated and the output is a

    // fixed-sized array whose Length can be evaluated

    uint64_t element_index;

    if (_.IsArrayType(graph_type_output_type) &&

        _.EvalConstantValUint64(inst->GetOperandAs<uint32_t>(2),

                                &element_index)) {

      uint64_t array_length;

      auto graph_type_output_type_inst = _.FindDef(graph_type_output_type);

      if (_.EvalConstantValUint64(

              graph_type_output_type_inst->GetOperandAs<uint32_t>(2),

              &array_length)) {

        if (element_index >= array_length) {

          return _.diag(SPV_ERROR_INVALID_DATA, inst)

                 << "OpGraphSetOutputARM ElementIndex out of range. The type "

                    "of the graph output being accessed "

                 << _.getIdName(graph_type_output_type) << " is an array of "

                 << array_length << " elements but " << "ElementIndex is "

                 << element_index;

        }

      }

    }

  }

  // Check Value's type matches with graph type

  uint32_t value = inst->GetOperandAs<uint32_t>(0);

  uint32_t value_type = _.FindDef(value)->type_id();

  if (has_element_index) {

    uint32_t expected_type = _.GetComponentType(graph_type_output_type);

    if (value_type != expected_type) {

      return _.diag(SPV_ERROR_INVALID_DATA, inst)

             << "The type " << _.getIdName(value_type)

             << " of Value provided to the graph output instruction "

             << _.getIdName(value) << " does not match the component type "

             << _.getIdName(expected_type) << " of output " << output_index

             << " in the graph type.";

    }

  } else {

    if (value_type != graph_type_output_type) {

      return _.diag(SPV_ERROR_INVALID_DATA, inst)

             << "The type " << _.getIdName(value_type)

             << " of Value provided to the graph output instruction "

             << _.getIdName(value) << " does not match the type "

             << _.getIdName(graph_type_output_type) << " of output "

             << output_index << " in the graph type.";

    }

  }

  return SPV_SUCCESS;

}

bool InputOutputInstructionsHaveDuplicateIndices(

    ValidationState_t& _, std::deque<const Instruction*>& inout_insts,

    const Instruction** first_dup) {

  std::set<std::pair<uint64_t, uint64_t>> inout_element_indices;

  for (auto const inst : inout_insts) {

    const bool is_input = inst->opcode() == spv::Op::OpGraphInputARM;

    bool has_element_index = inst->operands().size() > (is_input ? 3 : 2);

    uint64_t inout_index;

    if (!_.EvalConstantValUint64(inst->GetOperandAs<uint32_t>(is_input ? 2 : 1),

                                 &inout_index)) {

      continue;

    }

    uint64_t element_index = -1;  // -1 means no ElementIndex

    if (has_element_index) {

      if (!_.EvalConstantValUint64(

              inst->GetOperandAs<uint32_t>(is_input ? 3 : 2), &element_index)) {

        continue;

      }

    }

    auto inout_element_pair = std::make_pair(inout_index, element_index);

    auto inout_noelement_pair = std::make_pair(inout_index, -1);

    if (inout_element_indices.count(inout_element_pair) ||

        inout_element_indices.count(inout_noelement_pair)) {

      *first_dup = inst;

      return true;

    }

    inout_element_indices.insert(inout_element_pair);

  }

  return false;

}

spv_result_t ValidateGraphEnd(ValidationState_t& _, const Instruction* inst) {

  size_t end_inst_num = inst->LineNum() - 1;

  // Gather OpGraphInputARM and OpGraphSetOutputARM instructions

  std::deque<const Instruction*> graph_inputs, graph_outputs;

  size_t in_inst_num = end_inst_num;

  auto graph_inst = &_.ordered_instructions()[in_inst_num];

  while (--in_inst_num) {

    graph_inst = &_.ordered_instructions()[in_inst_num];

    if (graph_inst->opcode() == spv::Op::OpGraphInputARM) {

      graph_inputs.push_front(graph_inst);

      continue;

    }

    if (graph_inst->opcode() == spv::Op::OpGraphSetOutputARM) {

      graph_outputs.push_front(graph_inst);

      continue;

    }

    if (graph_inst->opcode() == spv::Op::OpGraphARM) {

      break;

    }

  }

  const Instruction* first_dup;

  // Check that there are no duplicate InputIndex and ElementIndex values

  if (InputOutputInstructionsHaveDuplicateIndices(_, graph_inputs,

                                                  &first_dup)) {

    return _.diag(SPV_ERROR_INVALID_DATA, first_dup)

           << "Two OpGraphInputARM instructions with the same InputIndex "

              "must not be part of the same "

           << "graph definition unless ElementIndex is present in both with "

              "different values.";

  }

  // Check that there are no duplicate OutputIndex and ElementIndex values

  if (InputOutputInstructionsHaveDuplicateIndices(_, graph_outputs,

                                                  &first_dup)) {

    return _.diag(SPV_ERROR_INVALID_DATA, first_dup)

           << "Two OpGraphSetOutputARM instructions with the same "

              "OutputIndex must not be part of the same "

           << "graph definition unless ElementIndex is present in both with "

              "different values.";

  }

  return SPV_SUCCESS;

}

}  // namespace

// Validates correctness of graph instructions.

spv_result_t GraphPass(ValidationState_t& _, const Instruction* inst) {

  switch (inst->opcode()) {

    case spv::Op::OpTypeGraphARM:

      return ValidateGraphType(_, inst);

    case spv::Op::OpGraphConstantARM:

      return ValidateGraphConstant(_, inst);

    case spv::Op::OpGraphEntryPointARM:

      return ValidateGraphEntryPoint(_, inst);

    case spv::Op::OpGraphARM:

      return ValidateGraph(_, inst);

    case spv::Op::OpGraphInputARM:

      return ValidateGraphInput(_, inst);

    case spv::Op::OpGraphSetOutputARM:

      return ValidateGraphSetOutput(_, inst);

    case spv::Op::OpGraphEndARM:

      return ValidateGraphEnd(_, inst);

    default:

      break;

  }

  return SPV_SUCCESS;

}

}  // namespace val

}  // namespace spvtools