// Copyright (c) 2016 Google Inc.

// Modifications Copyright (C) 2024 Advanced Micro Devices, Inc. All rights

// reserved.

//

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

#include <algorithm>

#include <cassert>

#include <iterator>

#include <sstream>

#include <string>

#include <unordered_map>

#include <unordered_set>

#include "source/binary.h"

#include "source/latest_version_spirv_header.h"

#include "source/parsed_operand.h"

#include "source/table2.h"

#include "source/to_string.h"

#include "spirv-tools/libspirv.h"

namespace spvtools {

NameMapper GetTrivialNameMapper() {

  return [](uint32_t i) { return spvtools::to_string(i); };

}

FriendlyNameMapper::FriendlyNameMapper(const spv_const_context context,

                                       const uint32_t* code,

                                       const size_t wordCount, uint32_t options)

    : grammar_(AssemblyGrammar(context)) {

  spv_diagnostic diag = nullptr;

  // We don't care if the parse fails.

  spvBinaryParseWithOptions(context, this, code, wordCount, nullptr,

                            ParseInstructionForwarder, &diag, options);

  spvDiagnosticDestroy(diag);

}

std::string FriendlyNameMapper::NameForId(uint32_t id) {

  auto iter = name_for_id_.find(id);

  if (iter == name_for_id_.end()) {

    // It must have been an invalid module, so just return a trivial mapping.

    // We don't care about uniqueness.

    return to_string(id);

  } else {

    return iter->second;

  }

}

std::string FriendlyNameMapper::Sanitize(const std::string& suggested_name) {

  if (suggested_name.empty()) return "_";

  // Otherwise, replace invalid characters by '_'.

  std::string result;

  std::string valid =

      "abcdefghijklmnopqrstuvwxyz"

      "ABCDEFGHIJKLMNOPQRSTUVWXYZ"

      "_0123456789";

  std::transform(suggested_name.begin(), suggested_name.end(),

                 std::back_inserter(result), [&valid](const char c) {

                   return (std::string::npos == valid.find(c)) ? '_' : c;

                 });

  return result;

}

void FriendlyNameMapper::SaveName(uint32_t id,

                                  const std::string& suggested_name) {

  if (name_for_id_.find(id) != name_for_id_.end()) return;

  std::string base_name;

  // Limit the size of the name to limit memory usage.

  const uint32_t kMaxSize = 256;

  if (suggested_name.size() > kMaxSize) {

    base_name = to_string(id);

  } else {

    base_name = Sanitize(suggested_name);

  }

  std::string name = base_name;

  auto inserted = used_names_.insert(name);

  if (!inserted.second) {

    const std::string base_name_prefix = base_name + "_";

    for (uint32_t index = 0; !inserted.second; ++index) {

      name = base_name_prefix + to_string(index);

      inserted = used_names_.insert(name);

    }

  }

  name_for_id_[id] = name;

}

void FriendlyNameMapper::SaveBuiltInName(uint32_t target_id,

                                         uint32_t built_in) {

#define GLCASE(name)                  \

  case spv::BuiltIn::name:            \

    SaveName(target_id, "gl_" #name); \

    return;

#define GLCASE2(name, suggested)           \

  case spv::BuiltIn::name:                 \

    SaveName(target_id, "gl_" #suggested); \

    return;

#define CASE(name)              \

  case spv::BuiltIn::name:      \

    SaveName(target_id, #name); \

    return;

  switch (spv::BuiltIn(built_in)) {

    GLCASE(Position)

    GLCASE(PointSize)

    GLCASE(ClipDistance)

    GLCASE(CullDistance)

    GLCASE2(VertexId, VertexID)

    GLCASE2(InstanceId, InstanceID)

    GLCASE2(PrimitiveId, PrimitiveID)

    GLCASE2(InvocationId, InvocationID)

    GLCASE(Layer)

    GLCASE(ViewportIndex)

    GLCASE(TessLevelOuter)

    GLCASE(TessLevelInner)

    GLCASE(TessCoord)

    GLCASE(PatchVertices)

    GLCASE(FragCoord)

    GLCASE(PointCoord)

    GLCASE(FrontFacing)

    GLCASE2(SampleId, SampleID)

    GLCASE(SamplePosition)

    GLCASE(SampleMask)

    GLCASE(FragDepth)

    GLCASE(HelperInvocation)

    GLCASE2(NumWorkgroups, NumWorkGroups)

    GLCASE2(WorkgroupSize, WorkGroupSize)

    GLCASE2(WorkgroupId, WorkGroupID)

    GLCASE2(LocalInvocationId, LocalInvocationID)

    GLCASE2(GlobalInvocationId, GlobalInvocationID)

    GLCASE(LocalInvocationIndex)

    CASE(WorkDim)

    CASE(GlobalSize)

    CASE(EnqueuedWorkgroupSize)

    CASE(GlobalOffset)

    CASE(GlobalLinearId)

    CASE(SubgroupSize)

    CASE(SubgroupMaxSize)

    CASE(NumSubgroups)

    CASE(NumEnqueuedSubgroups)

    CASE(SubgroupId)

    CASE(SubgroupLocalInvocationId)

    GLCASE(VertexIndex)

    GLCASE(InstanceIndex)

    GLCASE(BaseInstance)

    CASE(SubgroupEqMaskKHR)

    CASE(SubgroupGeMaskKHR)

    CASE(SubgroupGtMaskKHR)

    CASE(SubgroupLeMaskKHR)

    CASE(SubgroupLtMaskKHR)

    default:

      break;

  }

#undef GLCASE

#undef GLCASE2

#undef CASE

}

spv_result_t FriendlyNameMapper::ParseInstruction(

    const spv_parsed_instruction_t& inst) {

  const auto result_id = inst.result_id;

  if (inst.num_operands == 0) return SPV_SUCCESS;

  switch (spv::Op(inst.opcode)) {

    case spv::Op::OpName:

      SaveName(inst.words[1], spvDecodeLiteralStringOperand(inst, 1));

      break;

    case spv::Op::OpDecorate:

      // Decorations come after OpName.  So OpName will take precedence over

      // decorations.

      //

      // In theory, we should also handle OpGroupDecorate.  But that's unlikely

      // to occur.

      if (spv::Decoration(inst.words[2]) == spv::Decoration::BuiltIn) {

        assert(inst.num_words > 3);

        SaveBuiltInName(inst.words[1], inst.words[3]);

      }

      break;

    case spv::Op::OpTypeVoid:

      SaveName(result_id, "void");

      break;

    case spv::Op::OpTypeBool:

      SaveName(result_id, "bool");

      break;

    case spv::Op::OpTypeInt: {

      std::string signedness;

      std::string root;

      const auto bit_width = inst.words[2];

      switch (bit_width) {

        case 8:

          root = "char";

          break;

        case 16:

          root = "short";

          break;

        case 32:

          root = "int";

          break;

        case 64:

          root = "long";

          break;

        default:

          root = to_string(bit_width);

          signedness = "i";

          break;

      }

      if (0 == inst.words[3]) signedness = "u";

      SaveName(result_id, signedness + root);

    } break;

    case spv::Op::OpTypeFloat: {

      const auto bit_width = inst.words[2];

      if (inst.num_words > 3) {

        if (spv::FPEncoding(inst.words[3]) == spv::FPEncoding::BFloat16KHR) {

          SaveName(result_id, "bfloat16");

          break;

        }

        if (spv::FPEncoding(inst.words[3]) == spv::FPEncoding::Float8E4M3EXT) {

          SaveName(result_id, "fp8e4m3");

          break;

        }

        if (spv::FPEncoding(inst.words[3]) == spv::FPEncoding::Float8E5M2EXT) {

          SaveName(result_id, "fp8e5m2");

          break;

        }

        if (spv::FPEncoding(inst.words[3]) == spv::FPEncoding::Float6E2M3EXT) {

          SaveName(result_id, "fp6e2m3");

          break;

        }

        if (spv::FPEncoding(inst.words[3]) == spv::FPEncoding::Float6E3M2EXT) {

          SaveName(result_id, "fp6e3m2");

          break;

        }

        if (spv::FPEncoding(inst.words[3]) == spv::FPEncoding::Float4E2M1EXT) {

          SaveName(result_id, "fp4e2m1");

          break;

        }

        if (spv::FPEncoding(inst.words[3]) ==

            spv::FPEncoding::Float8UnsignedE8M0EXT) {

          SaveName(result_id, "fp8e8m0");

          break;

        }

        if (spv::FPEncoding(inst.words[3]) == spv::FPEncoding::MXInt8EXT) {

          SaveName(result_id, "mxint8");

          break;

        }

      }

      switch (bit_width) {

        case 16:

          SaveName(result_id, "half");

          break;

        case 32:

          SaveName(result_id, "float");

          break;

        case 64:

          SaveName(result_id, "double");

          break;

        default:

          SaveName(result_id, std::string("fp") + to_string(bit_width));

          break;

      }

    } break;

    case spv::Op::OpTypeVector:

      SaveName(result_id, std::string("v") + to_string(inst.words[3]) +

                              NameForId(inst.words[2]));

      break;

    case spv::Op::OpTypeMatrix:

      SaveName(result_id, std::string("mat") + to_string(inst.words[3]) +

                              NameForId(inst.words[2]));

      break;

    case spv::Op::OpTypeArray:

      SaveName(result_id, std::string("_arr_") + NameForId(inst.words[2]) +

                              "_" + NameForId(inst.words[3]));

      break;

    case spv::Op::OpTypeRuntimeArray:

      SaveName(result_id,

               std::string("_runtimearr_") + NameForId(inst.words[2]));

      break;

    case spv::Op::OpTypeNodePayloadArrayAMDX:

      SaveName(result_id,

               std::string("_payloadarr_") + NameForId(inst.words[2]));

      break;

    case spv::Op::OpTypePointer:

      SaveName(result_id, std::string("_ptr_") +

                              NameForEnumOperand(SPV_OPERAND_TYPE_STORAGE_CLASS,

                                                 inst.words[2]) +

                              "_" + NameForId(inst.words[3]));

      break;

    case spv::Op::OpTypeUntypedPointerKHR:

      SaveName(result_id, std::string("_ptr_") +

                              NameForEnumOperand(SPV_OPERAND_TYPE_STORAGE_CLASS,

                                                 inst.words[2]));

      break;

    case spv::Op::OpTypePipe:

      SaveName(result_id,

               std::string("Pipe") +

                   NameForEnumOperand(SPV_OPERAND_TYPE_ACCESS_QUALIFIER,

                                      inst.words[2]));

      break;

    case spv::Op::OpTypeEvent:

      SaveName(result_id, "Event");

      break;

    case spv::Op::OpTypeDeviceEvent:

      SaveName(result_id, "DeviceEvent");

      break;

    case spv::Op::OpTypeReserveId:

      SaveName(result_id, "ReserveId");

      break;

    case spv::Op::OpTypeQueue:

      SaveName(result_id, "Queue");

      break;

    case spv::Op::OpTypeOpaque:

      SaveName(result_id, std::string("Opaque_") +

                              Sanitize(spvDecodeLiteralStringOperand(inst, 1)));

      break;

    case spv::Op::OpTypePipeStorage:

      SaveName(result_id, "PipeStorage");

      break;

    case spv::Op::OpTypeNamedBarrier:

      SaveName(result_id, "NamedBarrier");

      break;

    case spv::Op::OpTypeStruct:

      // Structs are mapped rather simplisitically. Just indicate that they

      // are a struct and then give the raw Id number.

      SaveName(result_id, std::string("_struct_") + to_string(result_id));

      break;

    case spv::Op::OpConstantTrue:

      SaveName(result_id, "true");

      break;

    case spv::Op::OpConstantFalse:

      SaveName(result_id, "false");

      break;

    case spv::Op::OpConstant: {

      std::ostringstream value;

      EmitNumericLiteral(&value, inst, inst.operands[2]);

      auto value_str = value.str();

      // Use 'n' to signify negative. Other invalid characters will be mapped

      // to underscore.

      for (auto& c : value_str)

        if (c == '-') c = 'n';

      SaveName(result_id, NameForId(inst.type_id) + "_" + value_str);

    } break;

    default:

      // If this instruction otherwise defines an Id, then save a mapping for

      // it.  This is needed to ensure uniqueness in there is an OpName with

      // string something like "1" that might collide with this result_id.

      // We should only do this if a name hasn't already been registered by some

      // previous forward reference.

      if (result_id && name_for_id_.find(result_id) == name_for_id_.end())

        SaveName(result_id, to_string(result_id));

      break;

  }

  return SPV_SUCCESS;

}

std::string FriendlyNameMapper::NameForEnumOperand(spv_operand_type_t type,

                                                   uint32_t word) {

  const spvtools::OperandDesc* desc = nullptr;

  if (SPV_SUCCESS == spvtools::LookupOperand(type, word, &desc)) {

    return desc->name().data();

  } else {

    // Invalid input.  Just give something.

    return std::string("StorageClass") + to_string(word);

  }

}

}  // namespace spvtools