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

#include <memory>
#include <utility>
#include <vector>

#include "effcee/effcee.h"
#include "gmock/gmock.h"
#include "gtest/gtest.h"
#include "source/opt/build_module.h"
#include "source/opt/instruction.h"
#include "spirv-tools/libspirv.hpp"

namespace spvtools {
namespace opt {
namespace analysis {
namespace {

bool Validate(const std::vector<uint32_t>& bin) {
  spv_target_env target_env = SPV_ENV_UNIVERSAL_1_2;
  spv_context spvContext = spvContextCreate(target_env);
  spv_diagnostic diagnostic = nullptr;
  spv_const_binary_t binary = {bin.data(), bin.size()};
  spv_result_t error = spvValidate(spvContext, &binary, &diagnostic);
  if (error != 0) spvDiagnosticPrint(diagnostic);
  spvDiagnosticDestroy(diagnostic);
  spvContextDestroy(spvContext);
  return error == 0;
}

void Match(const std::string& original, IRContext* context,
           bool do_validation = true) {
  std::vector<uint32_t> bin;
  context->module()->ToBinary(&bin, true);
  if (do_validation) {
    EXPECT_TRUE(Validate(bin));
  }
  std::string assembly;
  SpirvTools tools(SPV_ENV_UNIVERSAL_1_2);
  EXPECT_TRUE(
      tools.Disassemble(bin, &assembly, SpirvTools::kDefaultDisassembleOption))
      << "Disassembling failed for shader:\n"
      << assembly << std::endl;
  auto match_result = effcee::Match(assembly, original);
  EXPECT_EQ(effcee::Result::Status::Ok, match_result.status())
      << match_result.message() << "\nChecking result:\n"
      << assembly;
}

std::vector<std::unique_ptr<Type>> GenerateAllTypes() {
  // Types in this test case are only equal to themselves, nothing else.
  std::vector<std::unique_ptr<Type>> types;

  // Void, Bool
  types.emplace_back(new Void());
  auto* voidt = types.back().get();
  types.emplace_back(new Bool());
  auto* boolt = types.back().get();

  // Integer
  types.emplace_back(new Integer(32, true));
  auto* s32 = types.back().get();
  types.emplace_back(new Integer(32, false));
  types.emplace_back(new Integer(64, true));
  types.emplace_back(new Integer(64, false));
  auto* u64 = types.back().get();

  // Float
  types.emplace_back(new Float(32));
  auto* f32 = types.back().get();
  types.emplace_back(new Float(64));

  // Vector
  types.emplace_back(new Vector(s32, 2));
  types.emplace_back(new Vector(s32, 3));
  auto* v3s32 = types.back().get();
  types.emplace_back(new Vector(u64, 4));
  types.emplace_back(new Vector(f32, 3));
  auto* v3f32 = types.back().get();

  // Matrix
  types.emplace_back(new Matrix(v3s32, 3));
  types.emplace_back(new Matrix(v3s32, 4));
  types.emplace_back(new Matrix(v3f32, 4));

  // Images
  types.emplace_back(new Image(s32, spv::Dim::Dim2D, 0, 0, 0, 0,
                               spv::ImageFormat::Rg8,
                               spv::AccessQualifier::ReadOnly));
  auto* image1 = types.back().get();
  types.emplace_back(new Image(s32, spv::Dim::Dim2D, 0, 1, 0, 0,
                               spv::ImageFormat::Rg8,
                               spv::AccessQualifier::ReadOnly));
  types.emplace_back(new Image(s32, spv::Dim::Dim3D, 0, 1, 0, 0,
                               spv::ImageFormat::Rg8,
                               spv::AccessQualifier::ReadOnly));
  types.emplace_back(new Image(voidt, spv::Dim::Dim3D, 0, 1, 0, 1,
                               spv::ImageFormat::Rg8,
                               spv::AccessQualifier::ReadWrite));
  auto* image2 = types.back().get();

  // Sampler
  types.emplace_back(new Sampler());

  // Sampled Image
  types.emplace_back(new SampledImage(image1));
  types.emplace_back(new SampledImage(image2));

  // Array
  types.emplace_back(new Array(f32, Array::LengthInfo{100, {0, 100u}}));
  types.emplace_back(new Array(f32, Array::LengthInfo{42, {0, 42u}}));
  auto* a42f32 = types.back().get();
  types.emplace_back(new Array(u64, Array::LengthInfo{24, {0, 24u}}));

  // RuntimeArray
  types.emplace_back(new RuntimeArray(v3f32));
  types.emplace_back(new RuntimeArray(v3s32));
  auto* rav3s32 = types.back().get();

  // Struct
  types.emplace_back(new Struct(std::vector<const Type*>{s32}));
  types.emplace_back(new Struct(std::vector<const Type*>{s32, f32}));
  auto* sts32f32 = types.back().get();
  types.emplace_back(
      new Struct(std::vector<const Type*>{u64, a42f32, rav3s32}));

  // Opaque
  types.emplace_back(new Opaque(""));
  types.emplace_back(new Opaque("hello"));
  types.emplace_back(new Opaque("world"));

  // Pointer
  types.emplace_back(new Pointer(f32, spv::StorageClass::Input));
  types.emplace_back(new Pointer(sts32f32, spv::StorageClass::Function));
  types.emplace_back(new Pointer(a42f32, spv::StorageClass::Function));
  types.emplace_back(new Pointer(nullptr, spv::StorageClass::Uniform));

  // Function
  types.emplace_back(new Function(voidt, {}));
  types.emplace_back(new Function(voidt, {boolt}));
  types.emplace_back(new Function(voidt, {boolt, s32}));
  types.emplace_back(new Function(s32, {boolt, s32}));

  // Event, Device Event, Reserve Id, Queue,
  types.emplace_back(new Event());
  types.emplace_back(new DeviceEvent());
  types.emplace_back(new ReserveId());
  types.emplace_back(new Queue());

  // Pipe, Forward Pointer, PipeStorage, NamedBarrier, AccelerationStructureNV,
  // CooperativeMatrixNV
  types.emplace_back(new Pipe(spv::AccessQualifier::ReadWrite));
  types.emplace_back(new Pipe(spv::AccessQualifier::ReadOnly));
  types.emplace_back(new ForwardPointer(1, spv::StorageClass::Input));
  types.emplace_back(new ForwardPointer(2, spv::StorageClass::Input));
  types.emplace_back(new ForwardPointer(2, spv::StorageClass::Uniform));
  types.emplace_back(new PipeStorage());
  types.emplace_back(new NamedBarrier());
  types.emplace_back(new AccelerationStructureNV());
  types.emplace_back(new CooperativeMatrixNV(f32, 24, 24, 24));
  types.emplace_back(new CooperativeMatrixKHR(f32, 8, 8, 8, 1002));
  types.emplace_back(new RayQueryKHR());
  types.emplace_back(new HitObjectNV());
  types.emplace_back(new HitObjectEXT());
  types.emplace_back(new CooperativeVectorNV(f32, 16));

  // SPV_AMDX_shader_enqueue
  types.emplace_back(new NodePayloadArrayAMDX(sts32f32));

  // Tensors
  types.emplace_back(new TensorARM(f32));
  auto* tensor_f32 = types.back().get();
  types.emplace_back(new TensorARM(f32, 4));
  auto* tensor_f32_ranked = types.back().get();
  types.emplace_back(new TensorARM(f32, 4, 44));
  auto* tensor_f32_shaped = types.back().get();

  // Graph
  types.emplace_back(new GraphARM(0, {tensor_f32}));
  types.emplace_back(new GraphARM(1, {tensor_f32_ranked, tensor_f32_ranked}));
  types.emplace_back(new GraphARM(1, {tensor_f32_shaped, tensor_f32_shaped}));

  types.emplace_back(new TensorLayoutNV(1002, 1000));
  types.emplace_back(new TensorViewNV(1002, 1003, {1000, 1001}));

  return types;
}

TEST(TypeManager, GenerateAllTypesGeneratesAllTypes) {
  std::set<Type::Kind> generated_types;
  for (auto& type : GenerateAllTypes()) {
    generated_types.insert(type->kind());
  }

  std::vector<Type::Kind> all_types;
  for (uint32_t kind = 0; kind != Type::Kind::kLast; ++kind) {
    all_types.push_back(static_cast<Type::Kind>(kind));
  }

  EXPECT_THAT(generated_types, testing::UnorderedElementsAreArray(all_types));
}

TEST(TypeManager, TypeStrings) {
  const std::string text = R"(
    OpDecorate %spec_const_with_id SpecId 99
    OpTypeForwardPointer %p Uniform
    %void    = OpTypeVoid
    %bool    = OpTypeBool
    %u32     = OpTypeInt 32 0
    %id4     = OpConstant %u32 4
    %s32     = OpTypeInt 32 1
    %f64     = OpTypeFloat 64
    %v3u32   = OpTypeVector %u32 3
    %m3x3    = OpTypeMatrix %v3u32 3
    %img1    = OpTypeImage %s32 Cube 0 1 1 0 R32f ReadWrite
    %img2    = OpTypeImage %s32 Cube 0 1 1 0 R32f
    %sampler = OpTypeSampler
    %si1     = OpTypeSampledImage %img1
    %si2     = OpTypeSampledImage %img2
    %a5u32   = OpTypeArray %u32 %id4
    %af64    = OpTypeRuntimeArray %f64
    %st1     = OpTypeStruct %u32
    %st2     = OpTypeStruct %f64 %s32 %v3u32
    %opaque1 = OpTypeOpaque ""
    %opaque2 = OpTypeOpaque "opaque"
    %p       = OpTypePointer Uniform %st1
    %f       = OpTypeFunction %void %u32 %u32
    %event   = OpTypeEvent
    %de      = OpTypeDeviceEvent
    %ri      = OpTypeReserveId
    %queue   = OpTypeQueue
    %pipe    = OpTypePipe ReadOnly
    %ps      = OpTypePipeStorage
    %nb      = OpTypeNamedBarrier
    %rtacc   = OpTypeAccelerationStructureNV
    ; Set up other kinds of OpTypeArray
    %s64     = OpTypeInt 64 1
    ; ID 32
    %spec_const_without_id = OpSpecConstant %s32 44
    %spec_const_with_id = OpSpecConstant %s32 42 ;; This is ID 1
    %long_constant = OpConstant %s64 5000000000
    %spec_const_op = OpSpecConstantOp %s32 IAdd %id4 %id4
    ; ID 35
    %arr_spec_const_without_id = OpTypeArray %s32 %spec_const_without_id
    %arr_spec_const_with_id = OpTypeArray %s32 %spec_const_with_id
    %arr_long_constant = OpTypeArray %s32 %long_constant
    %arr_spec_const_op = OpTypeArray %s32 %spec_const_op
    %cm   = OpTypeCooperativeMatrixNV %f64 %id4 %id4 %id4
    %id2    = OpConstant %u32 2
    %cmkhr  = OpTypeCooperativeMatrixKHR %f64 %id4 %id4 %id4 %id2
    %untyped = OpTypeUntypedPointerKHR Uniform
    ; ID 43
    %ts_shape = OpConstantComposite %a5u32 %id4 %id4 %id4 %id4
    %ts  = OpTypeTensorARM %u32
    %tsr = OpTypeTensorARM %u32 %id4
    %tss = OpTypeTensorARM %u32 %id4 %ts_shape
    %g_noin = OpTypeGraphARM 0 %ts
    %g_onein = OpTypeGraphARM 1 %tsr %tsr
    %g_shaped = OpTypeGraphARM 1 %tss %tss
  )";

  std::vector<std::pair<uint32_t, std::string>> type_id_strs = {
      {3, "void"},
      {4, "bool"},
      {5, "uint32"},
      // Id 6 is used by the constant.
      {7, "sint32"},
      {8, "float64"},
      {9, "<uint32, 3>"},
      {10, "<<uint32, 3>, 3>"},
      {11, "image(sint32, 3, 0, 1, 1, 0, 3, 2)"},
      {12, "image(sint32, 3, 0, 1, 1, 0, 3, 0)"},
      {13, "sampler"},
      {14, "sampled_image(image(sint32, 3, 0, 1, 1, 0, 3, 2))"},
      {15, "sampled_image(image(sint32, 3, 0, 1, 1, 0, 3, 0))"},
      {16, "[uint32, id(6), words(0,4)]"},
      {17, "[float64]"},
      {18, "{uint32}"},
      {19, "{float64, sint32, <uint32, 3>}"},
      {20, "opaque('')"},
      {21, "opaque('opaque')"},
      {2, "{uint32} 2*"},  // Include storage class number
      {22, "(uint32, uint32) -> void"},
      {23, "event"},
      {24, "device_event"},
      {25, "reserve_id"},
      {26, "queue"},
      {27, "pipe(0)"},
      {28, "pipe_storage"},
      {29, "named_barrier"},
      {30, "accelerationStructureNV"},
      {31, "sint64"},
      {35, "[sint32, id(32), words(0,44)]"},
      {36, "[sint32, id(1), words(1,99,42)]"},
      {37, "[sint32, id(33), words(0,705032704,1)]"},
      {38, "[sint32, id(34), words(2,34)]"},
      {39, "<float64, 6, 6, 6>"},
      {41, "<float64, 6, 6, 6, 40>"},
      {42, "untyped_ptr 2*"},  // Include storage class number
      // Id 43 is OpConstantComposite %a5u32 %id4 %id4 %id4 %id4
      {44, "tensor<uint32, id(0), id(0)>"},
      {45, "tensor<uint32, id(6), id(0)>"},
      {46, "tensor<uint32, id(6), id(43)>"},
      {47, "graph<0,tensor<uint32, id(0), id(0)>>"},
      {48,
       "graph<1,tensor<uint32, id(6), id(0)>,tensor<uint32, id(6), id(0)>>"},
      {49,
       "graph<1,tensor<uint32, id(6), id(43)>,tensor<uint32, id(6), id(43)>>"},
  };

  std::unique_ptr<IRContext> context =
      BuildModule(SPV_ENV_UNIVERSAL_1_4, nullptr, text);
  ASSERT_NE(nullptr, context.get());  // It assembled
  TypeManager manager(nullptr, context.get());

  EXPECT_EQ(type_id_strs.size(), manager.NumTypes());

  for (const auto& p : type_id_strs) {
    ASSERT_NE(nullptr, manager.GetType(p.first));
    EXPECT_EQ(p.second, manager.GetType(p.first)->str())
        << " id is " << p.first;
    EXPECT_EQ(p.first, manager.GetId(manager.GetType(p.first)));
  }
}

TEST(TypeManager, StructWithFwdPtr) {
  const std::string text = R"(
               OpCapability Addresses
               OpCapability Kernel
          %1 = OpExtInstImport "OpenCL.std"
               OpMemoryModel Physical64 OpenCL
               OpEntryPoint Kernel %7 "test"
               OpSource OpenCL_C 102000
               OpDecorate %11 FuncParamAttr NoCapture
         %11 = OpDecorationGroup
               OpGroupDecorate %11 %8 %9
               OpTypeForwardPointer %100 CrossWorkgroup
       %void = OpTypeVoid
  %150 = OpTypeStruct %100
%100 = OpTypePointer CrossWorkgroup %150
          %6 = OpTypeFunction %void %100 %100
          %7 = OpFunction %void Pure %6
          %8 = OpFunctionParameter %100
          %9 = OpFunctionParameter %100
         %10 = OpLabel
               OpReturn
               OpFunctionEnd
  )";

  std::unique_ptr<IRContext> context =
      BuildModule(SPV_ENV_UNIVERSAL_1_1, nullptr, text,
                  SPV_TEXT_TO_BINARY_OPTION_PRESERVE_NUMERIC_IDS);
  TypeManager manager(nullptr, context.get());

  Type* p100 = manager.GetType(100);
  Type* s150 = manager.GetType(150);

  EXPECT_TRUE(p100->AsPointer());
  EXPECT_EQ(p100->AsPointer()->pointee_type(), s150);

  EXPECT_TRUE(s150->AsStruct());
  EXPECT_EQ(s150->AsStruct()->element_types()[0], p100);
}

TEST(TypeManager, CircularFwdPtr) {
  const std::string text = R"(
               OpCapability Addresses
               OpCapability Kernel
          %1 = OpExtInstImport "OpenCL.std"
               OpMemoryModel Physical64 OpenCL
               OpEntryPoint Kernel %7 "test"
               OpSource OpenCL_C 102000
               OpDecorate %11 FuncParamAttr NoCapture
         %11 = OpDecorationGroup
               OpGroupDecorate %11 %8 %9
               OpTypeForwardPointer %100 CrossWorkgroup
               OpTypeForwardPointer %200 CrossWorkgroup
       %void = OpTypeVoid
        %int = OpTypeInt 32 0
      %float = OpTypeFloat 32
  %150 = OpTypeStruct %200 %int
  %250 = OpTypeStruct %100 %float
%100 = OpTypePointer CrossWorkgroup %150
%200 = OpTypePointer CrossWorkgroup %250
          %6 = OpTypeFunction %void %100 %200
          %7 = OpFunction %void Pure %6
          %8 = OpFunctionParameter %100
          %9 = OpFunctionParameter %200
         %10 = OpLabel
               OpReturn
               OpFunctionEnd
  )";

  std::unique_ptr<IRContext> context =
      BuildModule(SPV_ENV_UNIVERSAL_1_1, nullptr, text,
                  SPV_TEXT_TO_BINARY_OPTION_PRESERVE_NUMERIC_IDS);
  TypeManager manager(nullptr, context.get());

  Type* p100 = manager.GetType(100);
  Type* s150 = manager.GetType(150);
  Type* p200 = manager.GetType(200);
  Type* s250 = manager.GetType(250);

  EXPECT_TRUE(p100->AsPointer());
  EXPECT_EQ(p100->AsPointer()->pointee_type(), s150);

  EXPECT_TRUE(p200->AsPointer());
  EXPECT_EQ(p200->AsPointer()->pointee_type(), s250);

  EXPECT_TRUE(s150->AsStruct());
  EXPECT_EQ(s150->AsStruct()->element_types()[0], p200);

  EXPECT_TRUE(s250->AsStruct());
  EXPECT_EQ(s250->AsStruct()->element_types()[0], p100);
}

TEST(TypeManager, IsomorphicStructWithFwdPtr) {
  const std::string text = R"(
               OpCapability Addresses
               OpCapability Kernel
          %1 = OpExtInstImport "OpenCL.std"
               OpMemoryModel Physical64 OpenCL
               OpEntryPoint Kernel %7 "test"
               OpSource OpenCL_C 102000
               OpDecorate %11 FuncParamAttr NoCapture
         %11 = OpDecorationGroup
               OpGroupDecorate %11 %8 %9
               OpTypeForwardPointer %100 CrossWorkgroup
               OpTypeForwardPointer %200 CrossWorkgroup
       %void = OpTypeVoid
  %_struct_1 = OpTypeStruct %100
  %_struct_2 = OpTypeStruct %200
%100 = OpTypePointer CrossWorkgroup %_struct_1
%200 = OpTypePointer CrossWorkgroup %_struct_2
          %6 = OpTypeFunction %void %100 %200
          %7 = OpFunction %void Pure %6
          %8 = OpFunctionParameter %100
          %9 = OpFunctionParameter %200
         %10 = OpLabel
               OpReturn
               OpFunctionEnd
  )";

  std::unique_ptr<IRContext> context =
      BuildModule(SPV_ENV_UNIVERSAL_1_1, nullptr, text,
                  SPV_TEXT_TO_BINARY_OPTION_PRESERVE_NUMERIC_IDS);
  TypeManager manager(nullptr, context.get());

  EXPECT_EQ(manager.GetType(100), manager.GetType(200));
}

TEST(TypeManager, IsomorphicCircularFwdPtr) {
  const std::string text = R"(
               OpCapability Addresses
               OpCapability Kernel
          %1 = OpExtInstImport "OpenCL.std"
               OpMemoryModel Physical64 OpenCL
               OpEntryPoint Kernel %7 "test"
               OpSource OpenCL_C 102000
               OpDecorate %11 FuncParamAttr NoCapture
         %11 = OpDecorationGroup
               OpGroupDecorate %11 %8 %9
               OpTypeForwardPointer %100 CrossWorkgroup
               OpTypeForwardPointer %200 CrossWorkgroup
               OpTypeForwardPointer %300 CrossWorkgroup
               OpTypeForwardPointer %400 CrossWorkgroup
       %void = OpTypeVoid
        %int = OpTypeInt 32 0
      %float = OpTypeFloat 32
  %150 = OpTypeStruct %200 %int
  %250 = OpTypeStruct %100 %float
  %350 = OpTypeStruct %400 %int
  %450 = OpTypeStruct %300 %float
%100 = OpTypePointer CrossWorkgroup %150
%200 = OpTypePointer CrossWorkgroup %250
%300 = OpTypePointer CrossWorkgroup %350
%400 = OpTypePointer CrossWorkgroup %450
          %6 = OpTypeFunction %void %100 %200
          %7 = OpFunction %void Pure %6
          %8 = OpFunctionParameter %100
          %9 = OpFunctionParameter %200
         %10 = OpLabel
               OpReturn
               OpFunctionEnd
  )";

  std::unique_ptr<IRContext> context =
      BuildModule(SPV_ENV_UNIVERSAL_1_1, nullptr, text,
                  SPV_TEXT_TO_BINARY_OPTION_PRESERVE_NUMERIC_IDS);
  TypeManager manager(nullptr, context.get());

  Type* p100 = manager.GetType(100);
  Type* p300 = manager.GetType(300);
  EXPECT_EQ(p100, p300);
  Type* p200 = manager.GetType(200);
  Type* p400 = manager.GetType(400);
  EXPECT_EQ(p200, p400);

  Type* p150 = manager.GetType(150);
  Type* p350 = manager.GetType(350);
  EXPECT_EQ(p150, p350);
  Type* p250 = manager.GetType(250);
  Type* p450 = manager.GetType(450);
  EXPECT_EQ(p250, p450);
}

TEST(TypeManager, PartialIsomorphicFwdPtr) {
  const std::string text = R"(
               OpCapability Addresses
               OpCapability Kernel
          %1 = OpExtInstImport "OpenCL.std"
               OpMemoryModel Physical64 OpenCL
               OpEntryPoint Kernel %7 "test"
               OpSource OpenCL_C 102000
               OpDecorate %11 FuncParamAttr NoCapture
         %11 = OpDecorationGroup
               OpGroupDecorate %11 %8 %9
               OpTypeForwardPointer %100 CrossWorkgroup
               OpTypeForwardPointer %200 CrossWorkgroup
       %void = OpTypeVoid
        %int = OpTypeInt 32 0
      %float = OpTypeFloat 32
  %150 = OpTypeStruct %200 %int
  %250 = OpTypeStruct %200 %int
%100 = OpTypePointer CrossWorkgroup %150
%200 = OpTypePointer CrossWorkgroup %250
          %6 = OpTypeFunction %void %100 %200
          %7 = OpFunction %void Pure %6
          %8 = OpFunctionParameter %100
          %9 = OpFunctionParameter %200
         %10 = OpLabel
               OpReturn
               OpFunctionEnd
  )";

  std::unique_ptr<IRContext> context =
      BuildModule(SPV_ENV_UNIVERSAL_1_1, nullptr, text,
                  SPV_TEXT_TO_BINARY_OPTION_PRESERVE_NUMERIC_IDS);
  TypeManager manager(nullptr, context.get());

  Type* p100 = manager.GetType(100);
  Type* p200 = manager.GetType(200);
  EXPECT_EQ(p100->AsPointer()->pointee_type(),
            p200->AsPointer()->pointee_type());
}

TEST(TypeManager, DecorationOnStruct) {
  const std::string text = R"(
    OpDecorate %struct1 Block
    OpDecorate %struct2 Block
    OpDecorate %struct3 Block
    OpDecorate %struct4 Block

    %u32 = OpTypeInt 32 0             ; id: 5
    %f32 = OpTypeFloat 32             ; id: 6
    %struct1 = OpTypeStruct %u32 %f32 ; base
    %struct2 = OpTypeStruct %f32 %u32 ; different member order
    %struct3 = OpTypeStruct %f32      ; different member list
    %struct4 = OpTypeStruct %u32 %f32 ; the same
    %struct7 = OpTypeStruct %f32      ; no decoration
  )";
  std::unique_ptr<IRContext> context =
      BuildModule(SPV_ENV_UNIVERSAL_1_1, nullptr, text);
  TypeManager manager(nullptr, context.get());

  ASSERT_EQ(7u, manager.NumTypes());
  // Make sure we get ids correct.
  ASSERT_EQ("uint32", manager.GetType(5)->str());
  ASSERT_EQ("float32", manager.GetType(6)->str());

  // Try all combinations of pairs. Expect to be the same type only when the
  // same id or (1, 4).
  for (const auto id1 : {1, 2, 3, 4, 7}) {
    for (const auto id2 : {1, 2, 3, 4, 7}) {
      if (id1 == id2 || (id1 == 1 && id2 == 4) || (id1 == 4 && id2 == 1)) {
        EXPECT_TRUE(manager.GetType(id1)->IsSame(manager.GetType(id2)))
            << "%struct" << id1 << " is expected to be the same as %struct"
            << id2;
      } else {
        EXPECT_FALSE(manager.GetType(id1)->IsSame(manager.GetType(id2)))
            << "%struct" << id1 << " is expected to be different with %struct"
            << id2;
      }
    }
  }
}

TEST(TypeManager, DecorationOnMember) {
  const std::string text = R"(
    OpMemberDecorate %struct1  0 Offset 0
    OpMemberDecorate %struct2  0 Offset 0
    OpMemberDecorate %struct3  0 Offset 0
    OpMemberDecorate %struct4  0 Offset 0
    OpMemberDecorate %struct5  1 Offset 0
    OpMemberDecorate %struct6  0 Offset 4

    OpDecorate %struct7 Block
    OpMemberDecorate %struct7  0 Offset 0

    %u32 = OpTypeInt 32 0              ; id: 8
    %f32 = OpTypeFloat 32              ; id: 9
    %struct1  = OpTypeStruct %u32 %f32 ; base
    %struct2  = OpTypeStruct %f32 %u32 ; different member order
    %struct3  = OpTypeStruct %f32      ; different member list
    %struct4  = OpTypeStruct %u32 %f32 ; the same
    %struct5  = OpTypeStruct %u32 %f32 ; member decorate different field
    %struct6  = OpTypeStruct %u32 %f32 ; different member decoration parameter
    %struct7  = OpTypeStruct %u32 %f32 ; extra decoration on the struct
    %struct10 = OpTypeStruct %u32 %f32 ; no member decoration
  )";
  std::unique_ptr<IRContext> context =
      BuildModule(SPV_ENV_UNIVERSAL_1_1, nullptr, text);
  TypeManager manager(nullptr, context.get());

  ASSERT_EQ(10u, manager.NumTypes());
  // Make sure we get ids correct.
  ASSERT_EQ("uint32", manager.GetType(8)->str());
  ASSERT_EQ("float32", manager.GetType(9)->str());

  // Try all combinations of pairs. Expect to be the same type only when the
  // same id or (1, 4).
  for (const auto id1 : {1, 2, 3, 4, 5, 6, 7, 10}) {
    for (const auto id2 : {1, 2, 3, 4, 5, 6, 7, 10}) {
      if (id1 == id2 || (id1 == 1 && id2 == 4) || (id1 == 4 && id2 == 1)) {
        EXPECT_TRUE(manager.GetType(id1)->IsSame(manager.GetType(id2)))
            << "%struct" << id1 << " is expected to be the same as %struct"
            << id2;
      } else {
        EXPECT_FALSE(manager.GetType(id1)->IsSame(manager.GetType(id2)))
            << "%struct" << id1 << " is expected to be different with %struct"
            << id2;
      }
    }
  }
}

TEST(TypeManager, DecorationEmpty) {
  const std::string text = R"(
    OpDecorate %struct1 Block
    OpMemberDecorate %struct2  0 Offset 0

    %u32 = OpTypeInt 32 0 ; id: 3
    %f32 = OpTypeFloat 32 ; id: 4
    %struct1  = OpTypeStruct %u32 %f32
    %struct2  = OpTypeStruct %f32 %u32
    %struct5  = OpTypeStruct %f32
  )";
  std::unique_ptr<IRContext> context =
      BuildModule(SPV_ENV_UNIVERSAL_1_1, nullptr, text);
  TypeManager manager(nullptr, context.get());

  ASSERT_EQ(5u, manager.NumTypes());
  // Make sure we get ids correct.
  ASSERT_EQ("uint32", manager.GetType(3)->str());
  ASSERT_EQ("float32", manager.GetType(4)->str());

  // %struct1 with decoration on itself
  EXPECT_FALSE(manager.GetType(1)->decoration_empty());
  // %struct2 with decoration on its member
  EXPECT_FALSE(manager.GetType(2)->decoration_empty());
  EXPECT_TRUE(manager.GetType(3)->decoration_empty());
  EXPECT_TRUE(manager.GetType(4)->decoration_empty());
  // %struct5 has no decorations
  EXPECT_TRUE(manager.GetType(5)->decoration_empty());
}

TEST(TypeManager, BeginEndForEmptyModule) {
  const std::string text = "";
  std::unique_ptr<IRContext> context =
      BuildModule(SPV_ENV_UNIVERSAL_1_1, nullptr, text);
  TypeManager manager(nullptr, context.get());
  ASSERT_EQ(0u, manager.NumTypes());

  EXPECT_EQ(manager.begin(), manager.end());
}

TEST(TypeManager, BeginEnd) {
  const std::string text = R"(
    %void1   = OpTypeVoid
    %void2   = OpTypeVoid
    %bool    = OpTypeBool
    %u32     = OpTypeInt 32 0
    %f64     = OpTypeFloat 64
  )";
  std::unique_ptr<IRContext> context =
      BuildModule(SPV_ENV_UNIVERSAL_1_1, nullptr, text);
  TypeManager manager(nullptr, context.get());
  ASSERT_EQ(5u, manager.NumTypes());

  EXPECT_NE(manager.begin(), manager.end());
  for (const auto& t : manager) {
    switch (t.first) {
      case 1:
      case 2:
        EXPECT_EQ("void", t.second->str());
        break;
      case 3:
        EXPECT_EQ("bool", t.second->str());
        break;
      case 4:
        EXPECT_EQ("uint32", t.second->str());
        break;
      case 5:
        EXPECT_EQ("float64", t.second->str());
        break;
      default:
        EXPECT_TRUE(false && "unreachable");
        break;
    }
  }
}

TEST(TypeManager, LookupType) {
  const std::string text = R"(
%void = OpTypeVoid
%uint = OpTypeInt 32 0
%int  = OpTypeInt 32 1
%vec2 = OpTypeVector %int 2
)";

  std::unique_ptr<IRContext> context =
      BuildModule(SPV_ENV_UNIVERSAL_1_2, nullptr, text);
  EXPECT_NE(context, nullptr);
  TypeManager manager(nullptr, context.get());

  Void voidTy;
  EXPECT_EQ(manager.GetId(&voidTy), 1u);

  Integer uintTy(32, false);
  EXPECT_EQ(manager.GetId(&uintTy), 2u);

  Integer intTy(32, true);
  EXPECT_EQ(manager.GetId(&intTy), 3u);

  Integer intTy2(32, true);
  Vector vecTy(&intTy2, 2u);
  EXPECT_EQ(manager.GetId(&vecTy), 4u);
}

TEST(TypeManager, RemoveId) {
  const std::string text = R"(
OpCapability Shader
OpCapability Linkage
OpMemoryModel Logical GLSL450
%1 = OpTypeInt 32 0
%2 = OpTypeInt 32 1
  )";

  std::unique_ptr<IRContext> context =
      BuildModule(SPV_ENV_UNIVERSAL_1_2, nullptr, text,
                  SPV_TEXT_TO_BINARY_OPTION_PRESERVE_NUMERIC_IDS);
  EXPECT_NE(context, nullptr);

  context->get_type_mgr()->RemoveId(1u);
  ASSERT_EQ(context->get_type_mgr()->GetType(1u), nullptr);
  ASSERT_NE(context->get_type_mgr()->GetType(2u), nullptr);

  context->get_type_mgr()->RemoveId(2u);
  ASSERT_EQ(context->get_type_mgr()->GetType(1u), nullptr);
  ASSERT_EQ(context->get_type_mgr()->GetType(2u), nullptr);
}

TEST(TypeManager, RemoveIdNonDuplicateAmbiguousType) {
  const std::string text = R"(
OpCapability Shader
OpCapability Linkage
OpMemoryModel Logical GLSL450
%1 = OpTypeInt 32 0
%2 = OpTypeStruct %1
  )";

  std::unique_ptr<IRContext> context =
      BuildModule(SPV_ENV_UNIVERSAL_1_2, nullptr, text,
                  SPV_TEXT_TO_BINARY_OPTION_PRESERVE_NUMERIC_IDS);
  EXPECT_NE(context, nullptr);

  Integer u32(32, false);
  Struct st({&u32});
  ASSERT_EQ(context->get_type_mgr()->GetId(&st), 2u);
  context->get_type_mgr()->RemoveId(2u);
  ASSERT_EQ(context->get_type_mgr()->GetType(2u), nullptr);
  ASSERT_EQ(context->get_type_mgr()->GetId(&st), 0u);
}

TEST(TypeManager, RemoveIdDuplicateAmbiguousType) {
  const std::string text = R"(
OpCapability Shader
OpCapability Linkage
OpMemoryModel Logical GLSL450
%1 = OpTypeInt 32 0
%2 = OpTypeStruct %1
%3 = OpTypeStruct %1
  )";

  std::unique_ptr<IRContext> context =
      BuildModule(SPV_ENV_UNIVERSAL_1_2, nullptr, text,
                  SPV_TEXT_TO_BINARY_OPTION_PRESERVE_NUMERIC_IDS);
  EXPECT_NE(context, nullptr);

  Integer u32(32, false);
  Struct st({&u32});
  uint32_t id = context->get_type_mgr()->GetId(&st);
  ASSERT_NE(id, 0u);
  uint32_t toRemove = id == 2u ? 2u : 3u;
  uint32_t toStay = id == 2u ? 3u : 2u;
  context->get_type_mgr()->RemoveId(toRemove);
  ASSERT_EQ(context->get_type_mgr()->GetType(toRemove), nullptr);
  ASSERT_EQ(context->get_type_mgr()->GetId(&st), toStay);
}

TEST(TypeManager, RemoveIdDoesntUnmapOtherTypes) {
  const std::string text = R"(
OpCapability Shader
OpMemoryModel Logical GLSL450
%1 = OpTypeInt 32 0
%2 = OpTypeStruct %1
%3 = OpTypeStruct %1
  )";

  std::unique_ptr<IRContext> context =
      BuildModule(SPV_ENV_UNIVERSAL_1_2, nullptr, text,
                  SPV_TEXT_TO_BINARY_OPTION_PRESERVE_NUMERIC_IDS);
  EXPECT_NE(context, nullptr);

  Integer u32(32, false);
  Struct st({&u32});

  EXPECT_EQ(1u, context->get_type_mgr()->GetId(&u32));
  uint32_t id = context->get_type_mgr()->GetId(&st);
  ASSERT_NE(id, 0u);
  uint32_t toRemove = id == 2u ? 3u : 2u;
  uint32_t toStay = id == 2u ? 2u : 3u;
  context->get_type_mgr()->RemoveId(toRemove);
  ASSERT_EQ(context->get_type_mgr()->GetType(toRemove), nullptr);
  ASSERT_EQ(context->get_type_mgr()->GetId(&st), toStay);
}

TEST(TypeManager, GetTypeAndPointerType) {
  const std::string text = R"(
OpCapability Shader
OpCapability Linkage
OpMemoryModel Logical GLSL450
%1 = OpTypeInt 32 0
%2 = OpTypeStruct %1
  )";

  std::unique_ptr<IRContext> context =
      BuildModule(SPV_ENV_UNIVERSAL_1_2, nullptr, text,
                  SPV_TEXT_TO_BINARY_OPTION_PRESERVE_NUMERIC_IDS);
  EXPECT_NE(context, nullptr);

  Integer u32(32, false);
  Pointer u32Ptr(&u32, spv::StorageClass::Function);
  Struct st({&u32});
  Pointer stPtr(&st, spv::StorageClass::Input);

  auto pair = context->get_type_mgr()->GetTypeAndPointerType(
      3u, spv::StorageClass::Function);
  ASSERT_EQ(nullptr, pair.first);
  ASSERT_EQ(nullptr, pair.second);

  pair = context->get_type_mgr()->GetTypeAndPointerType(
      1u, spv::StorageClass::Function);
  ASSERT_TRUE(pair.first->IsSame(&u32));
  ASSERT_TRUE(pair.second->IsSame(&u32Ptr));

  pair = context->get_type_mgr()->GetTypeAndPointerType(
      2u, spv::StorageClass::Input);
  ASSERT_TRUE(pair.first->IsSame(&st));
  ASSERT_TRUE(pair.second->IsSame(&stPtr));
}

TEST(TypeManager, DuplicateType) {
  const std::string text = R"(
OpCapability Shader
OpMemoryModel Logical GLSL450
%1 = OpTypeInt 32 0
%2 = OpTypeInt 32 0
  )";

  std::unique_ptr<IRContext> context =
      BuildModule(SPV_ENV_UNIVERSAL_1_2, nullptr, text,
                  SPV_TEXT_TO_BINARY_OPTION_PRESERVE_NUMERIC_IDS);
  EXPECT_NE(context, nullptr);

  const Type* type1 = context->get_type_mgr()->GetType(1u);
  const Type* type2 = context->get_type_mgr()->GetType(2u);
  EXPECT_NE(type1, nullptr);
  EXPECT_NE(type2, nullptr);
  EXPECT_EQ(*type1, *type2);
}

TEST(TypeManager, MultipleStructs) {
  const std::string text = R"(
OpCapability Shader
OpMemoryModel Logical GLSL450
OpDecorate %3 Constant
%1 = OpTypeInt 32 0
%2 = OpTypeStruct %1
%3 = OpTypeStruct %1
  )";

  std::unique_ptr<IRContext> context =
      BuildModule(SPV_ENV_UNIVERSAL_1_2, nullptr, text,
                  SPV_TEXT_TO_BINARY_OPTION_PRESERVE_NUMERIC_IDS);
  EXPECT_NE(context, nullptr);

  const Type* type1 = context->get_type_mgr()->GetType(2u);
  const Type* type2 = context->get_type_mgr()->GetType(3u);
  EXPECT_NE(type1, nullptr);
  EXPECT_NE(type2, nullptr);
  EXPECT_FALSE(type1->IsSame(type2));
}

TEST(TypeManager, RemovingIdAvoidsUseAfterFree) {
  const std::string text = R"(
OpCapability Shader
OpMemoryModel Logical GLSL450
%1 = OpTypeInt 32 0
%2 = OpTypeStruct %1
  )";

  std::unique_ptr<IRContext> context =
      BuildModule(SPV_ENV_UNIVERSAL_1_2, nullptr, text,
                  SPV_TEXT_TO_BINARY_OPTION_PRESERVE_NUMERIC_IDS);
  EXPECT_NE(context, nullptr);

  Integer u32(32, false);
  Struct st({&u32});
  const Type* type = context->get_type_mgr()->GetType(2u);
  EXPECT_NE(type, nullptr);
  context->get_type_mgr()->RemoveId(1u);
  EXPECT_TRUE(type->IsSame(&st));
}

TEST(TypeManager, RegisterAndRemoveId) {
  const std::string text = R"(
OpCapability Shader
OpMemoryModel Logical GLSL450
%1 = OpTypeInt 32 0
)";

  std::unique_ptr<IRContext> context =
      BuildModule(SPV_ENV_UNIVERSAL_1_2, nullptr, text,
                  SPV_TEXT_TO_BINARY_OPTION_PRESERVE_NUMERIC_IDS);
  EXPECT_NE(context, nullptr);

  uint32_t id = 2u;
  {
    // Ensure that u32 goes out of scope.
    Integer u32(32, false);
    Struct st({&u32});
    context->get_type_mgr()->RegisterType(id, st);
  }

  context->get_type_mgr()->RemoveId(id);
  EXPECT_EQ(nullptr, context->get_type_mgr()->GetType(id));
}

TEST(TypeManager, RegisterAndRemoveIdAllTypes) {
  const std::string text = R"(
OpCapability Shader
OpMemoryModel Logical GLSL450
)";

  std::unique_ptr<IRContext> context =
      BuildModule(SPV_ENV_UNIVERSAL_1_2, nullptr, text,
                  SPV_TEXT_TO_BINARY_OPTION_PRESERVE_NUMERIC_IDS);
  EXPECT_NE(context, nullptr);

  std::vector<std::unique_ptr<Type>> types = GenerateAllTypes();
  uint32_t id = 0u;
  for (auto& t : types) {
    context->get_type_mgr()->RegisterType(++id, *t);
    EXPECT_EQ(*t, *context->get_type_mgr()->GetType(id));
    EXPECT_EQ(id, context->get_type_mgr()->GetId(t.get()));
  }
  types.clear();

  for (; id > 0; --id) {
    context->get_type_mgr()->RemoveId(id);
    EXPECT_EQ(nullptr, context->get_type_mgr()->GetType(id));
  }
}

TEST(TypeManager, RegisterAndRemoveIdWithDecorations) {
  const std::string text = R"(
OpCapability Shader
OpMemoryModel Logical GLSL450
%1 = OpTypeInt 32 0
)";

  std::unique_ptr<IRContext> context =
      BuildModule(SPV_ENV_UNIVERSAL_1_2, nullptr, text,
                  SPV_TEXT_TO_BINARY_OPTION_PRESERVE_NUMERIC_IDS);
  EXPECT_NE(context, nullptr);

  uint32_t id = 2u;
  {
    Integer u32(32, false);
    Struct st({&u32, &u32});
    st.AddDecoration({10});
    st.AddDecoration({11});
    st.AddMemberDecoration(0, {{35, 4}});
    st.AddMemberDecoration(1, {{35, 4}});
    st.AddMemberDecoration(1, {{36, 5}});
    context->get_type_mgr()->RegisterType(id, st);
    EXPECT_EQ(st, *context->get_type_mgr()->GetType(id));
  }

  context->get_type_mgr()->RemoveId(id);
  EXPECT_EQ(nullptr, context->get_type_mgr()->GetType(id));
}

TEST(TypeManager, GetTypeInstructionInt) {
  const std::string text = R"(
; CHECK: OpTypeInt 32 0
; CHECK: OpTypeInt 16 1
OpCapability Shader
OpCapability Int16
OpCapability Linkage
OpMemoryModel Logical GLSL450
  )";

  std::unique_ptr<IRContext> context =
      BuildModule(SPV_ENV_UNIVERSAL_1_2, nullptr, text);
  EXPECT_NE(context, nullptr);

  Integer uint_32(32, false);
  context->get_type_mgr()->GetTypeInstruction(&uint_32);

  Integer int_16(16, true);
  context->get_type_mgr()->GetTypeInstruction(&int_16);

  Match(text, context.get());
}

TEST(TypeManager, GetTypeInstructionDuplicateInts) {
  const std::string text = R"(
; CHECK: OpTypeInt 32 0
; CHECK-NOT: OpType
OpCapability Shader
OpCapability Linkage
OpMemoryModel Logical GLSL450
  )";

  std::unique_ptr<IRContext> context =
      BuildModule(SPV_ENV_UNIVERSAL_1_2, nullptr, text);
  EXPECT_NE(context, nullptr);

  Integer uint_32(32, false);
  uint32_t id = context->get_type_mgr()->GetTypeInstruction(&uint_32);

  Integer other(32, false);
  EXPECT_EQ(context->get_type_mgr()->GetTypeInstruction(&other), id);

  Match(text, context.get());
}

TEST(TypeManager, GetTypeInstructionAllTypes) {
  const std::string text = R"(
; CHECK: [[uint:%\w+]] = OpTypeInt 32 0
; CHECK: [[input_ptr:%\w+]] = OpTypePointer Input [[uint]]
; CHECK: [[uniform_ptr:%\w+]] = OpTypePointer Uniform [[uint]]
; CHECK: [[uint2:%\w+]] = OpConstant [[uint]] 2
; CHECK: [[uint8:%\w+]] = OpConstant [[uint]] 8
; CHECK: [[uint4:%\w+]] = OpConstant [[uint]] 4
; CHECK: [[uint_arr4:%\w+]] = OpTypeArray [[uint]] [[uint4]]
; CHECK: [[uint24:%\w+]] = OpConstant [[uint]] 24
; CHECK: [[uint42:%\w+]] = OpConstant [[uint]] 42
; CHECK: [[uint_arr4_44:%\w+]] = OpConstantComposite [[uint_arr4]] [[uint4]] [[uint4]] [[uint4]] [[uint4]]
; CHECK: [[uint100:%\w+]] = OpConstant [[uint]] 100
; CHECK: [[void:%\w+]] = OpTypeVoid
; CHECK: [[bool:%\w+]] = OpTypeBool
; CHECK: [[s32:%\w+]] = OpTypeInt 32 1
; CHECK: OpTypeInt 64 1
; CHECK: [[u64:%\w+]] = OpTypeInt 64 0
; CHECK: [[f32:%\w+]] = OpTypeFloat 32
; CHECK: OpTypeFloat 64
; CHECK: OpTypeVector [[s32]] 2
; CHECK: [[v3s32:%\w+]] = OpTypeVector [[s32]] 3
; CHECK: OpTypeVector [[u64]] 4
; CHECK: [[v3f32:%\w+]] = OpTypeVector [[f32]] 3
; CHECK: OpTypeMatrix [[v3s32]] 3
; CHECK: OpTypeMatrix [[v3s32]] 4
; CHECK: OpTypeMatrix [[v3f32]] 4
; CHECK: [[image1:%\w+]] = OpTypeImage [[s32]] 2D 0 0 0 0 Rg8 ReadOnly
; CHECK: OpTypeImage [[s32]] 2D 0 1 0 0 Rg8 ReadOnly
; CHECK: OpTypeImage [[s32]] 3D 0 1 0 0 Rg8 ReadOnly
; CHECK: [[image2:%\w+]] = OpTypeImage [[void]] 3D 0 1 0 1 Rg8 ReadWrite
; CHECK: OpTypeSampler
; CHECK: OpTypeSampledImage [[image1]]
; CHECK: OpTypeSampledImage [[image2]]
; CHECK: OpTypeArray [[f32]] [[uint100]]
; CHECK: [[a42f32:%\w+]] = OpTypeArray [[f32]] [[uint42]]
; CHECK: OpTypeArray [[u64]] [[uint24]]
; CHECK: OpTypeRuntimeArray [[v3f32]]
; CHECK: [[rav3s32:%\w+]] = OpTypeRuntimeArray [[v3s32]]
; CHECK: OpTypeStruct [[s32]]
; CHECK: [[sts32f32:%\w+]] = OpTypeStruct [[s32]] [[f32]]
; CHECK: OpTypeStruct [[u64]] [[a42f32]] [[rav3s32]]
; CHECK: OpTypeOpaque ""
; CHECK: OpTypeOpaque "hello"
; CHECK: OpTypeOpaque "world"
; CHECK: OpTypePointer Input [[f32]]
; CHECK: OpTypePointer Function [[sts32f32]]
; CHECK: OpTypePointer Function [[a42f32]]
; CHECK: OpTypeFunction [[void]]
; CHECK: OpTypeFunction [[void]] [[bool]]
; CHECK: OpTypeFunction [[void]] [[bool]] [[s32]]
; CHECK: OpTypeFunction [[s32]] [[bool]] [[s32]]
; CHECK: OpTypeEvent
; CHECK: OpTypeDeviceEvent
; CHECK: OpTypeReserveId
; CHECK: OpTypeQueue
; CHECK: OpTypePipe ReadWrite
; CHECK: OpTypePipe ReadOnly
; CHECK: OpTypeForwardPointer [[input_ptr]] Input
; CHECK: OpTypeForwardPointer [[uniform_ptr]] Input
; CHECK: OpTypeForwardPointer [[uniform_ptr]] Uniform
; CHECK: OpTypePipeStorage
; CHECK: OpTypeNamedBarrier
; CHECK: OpTypeAccelerationStructureKHR
; CHECK: OpTypeCooperativeMatrixNV [[f32]] [[uint24]] [[uint24]] [[uint24]]
; CHECK: OpTypeCooperativeMatrixKHR [[f32]] [[uint8]] [[uint8]] [[uint8]] [[uint2]]
; CHECK: OpTypeRayQueryKHR
; CHECK: OpTypeHitObjectNV
; CHECK: [[tensor_f32:%\w+]] = OpTypeTensorARM [[f32]]
; CHECK: [[tensor_f32_ranked:%\w+]] = OpTypeTensorARM [[f32]] [[uint4]]
; CHECK: [[tensor_f32_shaped:%\w+]] = OpTypeTensorARM [[f32]] [[uint4]] [[uint_arr4_44]]
; CHECK: OpTypeGraphARM 0 [[tensor_f32]]
; CHECK: OpTypeGraphARM 1 [[tensor_f32_ranked]] [[tensor_f32_ranked]]
; CHECK: OpTypeGraphARM 1 [[tensor_f32_shaped]] [[tensor_f32_shaped]]
OpCapability Shader
OpCapability Int64
OpCapability Linkage
OpMemoryModel Logical GLSL450
%uint = OpTypeInt 32 0
%1 = OpTypePointer Input %uint
%2 = OpTypePointer Uniform %uint
%1000 = OpConstant %uint 0
%1001 = OpConstant %uint 1
%1002 = OpConstant %uint 2
%8 = OpConstant %uint 8
%4 = OpConstant %uint 4
%5 = OpTypeArray %uint %4
%24 = OpConstant %uint 24
%42 = OpConstant %uint 42
%44 = OpConstantComposite %5 %4 %4 %4 %4
%100 = OpConstant %uint 100
%1003 = OpConstantFalse %bool
  )";

  std::unique_ptr<IRContext> context =
      BuildModule(SPV_ENV_UNIVERSAL_1_2, nullptr, text,
                  SPV_TEXT_TO_BINARY_OPTION_PRESERVE_NUMERIC_IDS);
  EXPECT_NE(context, nullptr);

  std::vector<std::unique_ptr<Type>> types = GenerateAllTypes();
  for (auto& t : types) {
    context->get_type_mgr()->GetTypeInstruction(t.get());
  }

  Match(text, context.get(), false);
}

TEST(TypeManager, GetTypeInstructionWithDecorations) {
  const std::string text = R"(
; CHECK: OpDecorate [[struct:%\w+]] CPacked
; CHECK: OpMemberDecorate [[struct]] 1 Offset 4
; CHECK: [[uint:%\w+]] = OpTypeInt 32 0
; CHECK: [[struct]] = OpTypeStruct [[uint]] [[uint]]
OpCapability Shader
OpCapability Kernel
OpCapability Linkage
OpMemoryModel Logical GLSL450
%uint = OpTypeInt 32 0
  )";

  std::unique_ptr<IRContext> context =
      BuildModule(SPV_ENV_UNIVERSAL_1_2, nullptr, text,
                  SPV_TEXT_TO_BINARY_OPTION_PRESERVE_NUMERIC_IDS);
  EXPECT_NE(context, nullptr);

  Integer u32(32, false);
  Struct st({&u32, &u32});
  st.AddDecoration({10});
  st.AddMemberDecoration(1, {{35, 4}});
  (void)context->get_def_use_mgr();
  context->get_type_mgr()->GetTypeInstruction(&st);

  Match(text, context.get());
}

TEST(TypeManager, GetPointerToAmbiguousType1) {
  const std::string text = R"(
; CHECK: [[struct1:%\w+]] = OpTypeStruct
; CHECK: [[struct2:%\w+]] = OpTypeStruct
; CHECK: OpTypePointer Function [[struct2]]
; CHECK: OpTypePointer Function [[struct1]]
OpCapability Shader
OpCapability Kernel
OpCapability Linkage
OpMemoryModel Logical GLSL450
%uint = OpTypeInt 32 0
%1 = OpTypeStruct %uint
%2 = OpTypeStruct %uint
%3 = OpTypePointer Function %2
  )";

  std::unique_ptr<IRContext> context =
      BuildModule(SPV_ENV_UNIVERSAL_1_2, nullptr, text,
                  SPV_TEXT_TO_BINARY_OPTION_PRESERVE_NUMERIC_IDS);
  EXPECT_NE(context, nullptr);

  context->get_type_mgr()->FindPointerToType(1, spv::StorageClass::Function);
  Match(text, context.get());
}

TEST(TypeManager, GetPointerToAmbiguousType2) {
  const std::string text = R"(
; CHECK: [[struct1:%\w+]] = OpTypeStruct
; CHECK: [[struct2:%\w+]] = OpTypeStruct
; CHECK: OpTypePointer Function [[struct1]]
; CHECK: OpTypePointer Function [[struct2]]
OpCapability Shader
OpCapability Kernel
OpCapability Linkage
OpMemoryModel Logical GLSL450
%uint = OpTypeInt 32 0
%1 = OpTypeStruct %uint
%2 = OpTypeStruct %uint
%3 = OpTypePointer Function %1
  )";

  std::unique_ptr<IRContext> context =
      BuildModule(SPV_ENV_UNIVERSAL_1_2, nullptr, text,
                  SPV_TEXT_TO_BINARY_OPTION_PRESERVE_NUMERIC_IDS);
  EXPECT_NE(context, nullptr);

  context->get_type_mgr()->FindPointerToType(2, spv::StorageClass::Function);
  Match(text, context.get());
}

// Structures containing circular type references
// (from https://github.com/KhronosGroup/SPIRV-Tools/issues/5623).
TEST(TypeManager, CircularPointerToStruct) {
  const std::string text = R"(
               OpCapability VariablePointers
               OpCapability PhysicalStorageBufferAddresses
               OpCapability Int64
               OpCapability Shader
               OpExtension "SPV_KHR_variable_pointers"
               OpExtension "SPV_KHR_physical_storage_buffer"
               OpMemoryModel PhysicalStorageBuffer64 GLSL450
               OpEntryPoint Fragment %1 "main"
               OpExecutionMode %1 OriginUpperLeft
               OpExecutionMode %1 DepthReplacing
               OpDecorate %1200 ArrayStride 24
               OpMemberDecorate %600 0 Offset 0
               OpMemberDecorate %800 0 Offset 0
               OpMemberDecorate %120 0 Offset 16
               OpTypeForwardPointer %1200 PhysicalStorageBuffer
                 %600 = OpTypeStruct %1200
                 %800 = OpTypeStruct %1200
                 %120 = OpTypeStruct %800
                %1200 = OpTypePointer PhysicalStorageBuffer %120
  )";

  std::unique_ptr<IRContext> context =
      BuildModule(SPV_ENV_UNIVERSAL_1_1, nullptr, text,
                  SPV_TEXT_TO_BINARY_OPTION_PRESERVE_NUMERIC_IDS);
  TypeManager manager(nullptr, context.get());
  uint32_t id = manager.FindPointerToType(600, spv::StorageClass::Function);
  EXPECT_EQ(id, 1201);
}

TEST(TypeManager, AttachLinkageDecoration) {
  const std::string text = R"(
      OpCapability Shader
      OpCapability Linkage
      OpMemoryModel Logical GLSL450
      OpDecorate %1000 LinkageAttributes "_1000" Export
       %800 = OpTypeInt 32 0
      %1000 = OpTypeStruct %800
      %1200 = OpTypeStruct %800
  )";

  std::unique_ptr<IRContext> context =
      BuildModule(SPV_ENV_UNIVERSAL_1_5, nullptr, text,
                  SPV_TEXT_TO_BINARY_OPTION_PRESERVE_NUMERIC_IDS);
  TypeManager manager(nullptr, context.get());

  constexpr uint32_t source_id = 1000u;
  constexpr uint32_t target_id = 1200u;
  std::vector<Instruction*> decorations =
      context->get_decoration_mgr()->GetDecorationsFor(source_id, true);
  Type* type = context->get_type_mgr()->GetType(target_id);
  for (auto dec : decorations) {
    manager.AttachDecoration(*dec, type);
  }
  EXPECT_FALSE(type->decoration_empty());
  EXPECT_TRUE(
      type->HasSameDecorations(context->get_type_mgr()->GetType(source_id)));
}

}  // namespace
}  // namespace analysis
}  // namespace opt
}  // namespace spvtools