// Copyright (c) 2022 The Khronos Group Inc.
//
// 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 ray tracing instructions from SPV_KHR_ray_tracing

#include "source/opcode.h"
#include "source/val/instruction.h"
#include "source/val/validate.h"
#include "source/val/validation_state.h"

namespace spvtools {
namespace val {

spv_result_t ValidateTraceRay(ValidationState_t& _, const Instruction* inst) {
  _.function(inst->function()->id())
      ->RegisterExecutionModelLimitation(
          [](spv::ExecutionModel model, std::string* message) {
            if (model != spv::ExecutionModel::RayGenerationKHR &&
                model != spv::ExecutionModel::ClosestHitKHR &&
                model != spv::ExecutionModel::MissKHR) {
              if (message) {
                *message =
                    "OpTraceRayKHR requires RayGenerationKHR, "
                    "ClosestHitKHR and MissKHR execution models";
              }
              return false;
            }
            return true;
          });

  if (_.GetIdOpcode(_.GetOperandTypeId(inst, 0)) !=
      spv::Op::OpTypeAccelerationStructureKHR) {
    return _.diag(SPV_ERROR_INVALID_DATA, inst)
           << "Expected Acceleration Structure to be of type "
              "OpTypeAccelerationStructureKHR";
  }

  const uint32_t ray_flags = _.GetOperandTypeId(inst, 1);
  if (!_.IsIntScalarType(ray_flags, 32)) {
    return _.diag(SPV_ERROR_INVALID_DATA, inst)
           << "Ray Flags must be a 32-bit int scalar";
  }

  const uint32_t cull_mask = _.GetOperandTypeId(inst, 2);
  if (!_.IsIntScalarType(cull_mask, 32)) {
    return _.diag(SPV_ERROR_INVALID_DATA, inst)
           << "Cull Mask must be a 32-bit int scalar";
  }

  const uint32_t sbt_offset = _.GetOperandTypeId(inst, 3);
  if (!_.IsIntScalarType(sbt_offset, 32)) {
    return _.diag(SPV_ERROR_INVALID_DATA, inst)
           << "SBT Offset must be a 32-bit int scalar";
  }

  const uint32_t sbt_stride = _.GetOperandTypeId(inst, 4);
  if (!_.IsIntScalarType(sbt_stride, 32)) {
    return _.diag(SPV_ERROR_INVALID_DATA, inst)
           << "SBT Stride must be a 32-bit int scalar";
  }

  const uint32_t miss_index = _.GetOperandTypeId(inst, 5);
  if (!_.IsIntScalarType(miss_index, 32)) {
    return _.diag(SPV_ERROR_INVALID_DATA, inst)
           << "Miss Index must be a 32-bit int scalar";
  }

  const uint32_t ray_origin = _.GetOperandTypeId(inst, 6);
  if (!_.IsFloatVectorType(ray_origin) || _.GetDimension(ray_origin) != 3 ||
      _.GetBitWidth(ray_origin) != 32) {
    return _.diag(SPV_ERROR_INVALID_DATA, inst)
           << "Ray Origin must be a 32-bit float 3-component vector";
  }

  const uint32_t ray_tmin = _.GetOperandTypeId(inst, 7);
  if (!_.IsFloatScalarType(ray_tmin, 32)) {
    return _.diag(SPV_ERROR_INVALID_DATA, inst)
           << "Ray TMin must be a 32-bit float scalar";
  }

  const uint32_t ray_direction = _.GetOperandTypeId(inst, 8);
  if (!_.IsFloatVectorType(ray_direction) ||
      _.GetDimension(ray_direction) != 3 ||
      _.GetBitWidth(ray_direction) != 32) {
    return _.diag(SPV_ERROR_INVALID_DATA, inst)
           << "Ray Direction must be a 32-bit float 3-component vector";
  }

  const uint32_t ray_tmax = _.GetOperandTypeId(inst, 9);
  if (!_.IsFloatScalarType(ray_tmax, 32)) {
    return _.diag(SPV_ERROR_INVALID_DATA, inst)
           << "Ray TMax must be a 32-bit float scalar";
  }

  const Instruction* payload = _.FindDef(inst->GetOperandAs<uint32_t>(10));
  if (payload->opcode() != spv::Op::OpVariable) {
    return _.diag(SPV_ERROR_INVALID_DATA, inst)
           << "Payload must be the result of a OpVariable";
  } else if (payload->GetOperandAs<spv::StorageClass>(2) !=
                 spv::StorageClass::RayPayloadKHR &&
             payload->GetOperandAs<spv::StorageClass>(2) !=
                 spv::StorageClass::IncomingRayPayloadKHR) {
    return _.diag(SPV_ERROR_INVALID_DATA, inst)
           << "Payload must have storage class RayPayloadKHR or "
              "IncomingRayPayloadKHR";
  }

  return SPV_SUCCESS;
}
spv_result_t ValidateReportIntersection(ValidationState_t& _,
                                        const Instruction* inst) {
  _.function(inst->function()->id())
      ->RegisterExecutionModelLimitation(
          [](spv::ExecutionModel model, std::string* message) {
            if (model != spv::ExecutionModel::IntersectionKHR) {
              if (message) {
                *message =
                    "OpReportIntersectionKHR requires IntersectionKHR "
                    "execution model";
              }
              return false;
            }
            return true;
          });

  const uint32_t result_type = inst->type_id();
  if (!_.IsBoolScalarType(result_type)) {
    return _.diag(SPV_ERROR_INVALID_DATA, inst)
           << "expected Result Type to be bool scalar type";
  }

  const uint32_t hit = _.GetOperandTypeId(inst, 2);
  if (!_.IsFloatScalarType(hit, 32)) {
    return _.diag(SPV_ERROR_INVALID_DATA, inst)
           << "Hit must be a 32-bit int scalar";
  }

  const uint32_t hit_kind = _.GetOperandTypeId(inst, 3);
  if (!_.IsUnsignedIntScalarType(hit_kind) || _.GetBitWidth(hit_kind) != 32) {
    return _.diag(SPV_ERROR_INVALID_DATA, inst)
           << "Hit Kind must be a 32-bit unsigned int scalar";
  }
  return SPV_SUCCESS;
}

spv_result_t ValidateExecuteCallable(ValidationState_t& _,
                                     const Instruction* inst) {
  _.function(inst->function()->id())
      ->RegisterExecutionModelLimitation(
          [](spv::ExecutionModel model, std::string* message) {
            if (model != spv::ExecutionModel::RayGenerationKHR &&
                model != spv::ExecutionModel::ClosestHitKHR &&
                model != spv::ExecutionModel::MissKHR &&
                model != spv::ExecutionModel::CallableKHR) {
              if (message) {
                *message =
                    "OpExecuteCallableKHR requires RayGenerationKHR, "
                    "ClosestHitKHR, MissKHR and CallableKHR execution models";
              }
              return false;
            }
            return true;
          });

  const uint32_t sbt_index = _.GetOperandTypeId(inst, 0);
  if (!_.IsUnsignedIntScalarType(sbt_index) || _.GetBitWidth(sbt_index) != 32) {
    return _.diag(SPV_ERROR_INVALID_DATA, inst)
           << "SBT Index must be a 32-bit unsigned int scalar";
  }

  const auto callable_data = _.FindDef(inst->GetOperandAs<uint32_t>(1));
  if (callable_data->opcode() != spv::Op::OpVariable) {
    return _.diag(SPV_ERROR_INVALID_DATA, inst)
           << "Callable Data must be the result of a OpVariable";
  } else if (callable_data->GetOperandAs<spv::StorageClass>(2) !=
                 spv::StorageClass::CallableDataKHR &&
             callable_data->GetOperandAs<spv::StorageClass>(2) !=
                 spv::StorageClass::IncomingCallableDataKHR) {
    return _.diag(SPV_ERROR_INVALID_DATA, inst)
           << "Callable Data must have storage class CallableDataKHR or "
              "IncomingCallableDataKHR";
  }
  return SPV_SUCCESS;
}

spv_result_t RayTracingPass(ValidationState_t& _, const Instruction* inst) {
  const spv::Op opcode = inst->opcode();

  switch (opcode) {
    case spv::Op::OpTraceRayKHR:
      return ValidateTraceRay(_, inst);
    case spv::Op::OpReportIntersectionKHR:
      return ValidateReportIntersection(_, inst);
    case spv::Op::OpExecuteCallableKHR:
      return ValidateExecuteCallable(_, inst);
    default:
      break;
  }

  return SPV_SUCCESS;
}
}  // namespace val
}  // namespace spvtools