/*========================== begin_copyright_notice ============================

Copyright (C) 2023 Intel Corporation

SPDX-License-Identifier: MIT

============================= end_copyright_notice ===========================*/

#include "UserAddrSpaceMD.hpp"
#include "IGC/Probe/Assertion.h"

using namespace IGC;

UserAddrSpaceMD::UserAddrSpaceMD(llvm::LLVMContext *ctx) {
  this->ctx = ctx;
  user_addrspace_priv = ctx->getMDKindID("user_as_priv");
  user_addrspace_global = ctx->getMDKindID("user_as_global");
  user_addrspace_local = ctx->getMDKindID("user_as_local");
  user_addrspace_generic = ctx->getMDKindID("user_as_generic");
  user_addrspace_raystack = ctx->getMDKindID("user_as_raystack");

  dummyNode = llvm::MDNode::get(*ctx, llvm::MDString::get(*ctx, ""));
}

void UserAddrSpaceMD::Set(llvm::Instruction *inst, LSC_DOC_ADDR_SPACE type, llvm::MDNode *node) {
  if (inst) {
    llvm::MDNode *chosen_node = node != nullptr ? node : dummyNode;

    user_addrspace chosen_type = -1;

    switch (type) {
    case LSC_DOC_ADDR_SPACE::PRIVATE:
      chosen_type = user_addrspace_priv;
      break;
    case LSC_DOC_ADDR_SPACE::GLOBAL:
      chosen_type = user_addrspace_global;
      break;
    case LSC_DOC_ADDR_SPACE::LOCAL:
      chosen_type = user_addrspace_local;
      break;
    case LSC_DOC_ADDR_SPACE::GENERIC:
      chosen_type = user_addrspace_generic;
      break;
    case LSC_DOC_ADDR_SPACE::RAYSTACK:
      chosen_type = user_addrspace_raystack;
      break;
    default:
      IGC_ASSERT_MESSAGE(0, "Unsupport addrspace type");
      break;
    }
    inst->setMetadata(chosen_type, chosen_node);
  }
}

LSC_DOC_ADDR_SPACE UserAddrSpaceMD::Get(llvm::Instruction *inst) {
  LSC_DOC_ADDR_SPACE addrSpace = LSC_DOC_ADDR_SPACE::INVALID;

  if (inst) {
    if (auto md_addrspace_org = inst->getMetadata(user_addrspace_priv)) {
      addrSpace = LSC_DOC_ADDR_SPACE::PRIVATE;
    } else if (auto md_addrspace_org = inst->getMetadata(user_addrspace_global)) {
      addrSpace = LSC_DOC_ADDR_SPACE::GLOBAL;
    } else if (auto md_addrspace_org = inst->getMetadata(user_addrspace_local)) {
      addrSpace = LSC_DOC_ADDR_SPACE::LOCAL;
    } else if (auto md_addrspace_org = inst->getMetadata(user_addrspace_generic)) {
      addrSpace = LSC_DOC_ADDR_SPACE::GENERIC;
    } else if (auto md_addrspace_org = inst->getMetadata(user_addrspace_raystack)) {
      addrSpace = LSC_DOC_ADDR_SPACE::RAYSTACK;
    }
  }

  return addrSpace;
}

bool UserAddrSpaceMD::Has(llvm::Instruction *inst, LSC_DOC_ADDR_SPACE type) { return Get(inst) == type; }