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

Copyright (C) 2017-2021 Intel Corporation

SPDX-License-Identifier: MIT

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

#include "TypesLegalizationPass.hpp"
#include "Compiler/IGCPassSupport.h"
#include "Probe/Assertion.h"

#include "common/LLVMWarningsPush.hpp"
#include "llvmWrapper/Support/Alignment.h"
#include <llvm/IR/InstIterator.h>
#include <llvmWrapper/IR/Instructions.h>
#include "common/LLVMWarningsPop.hpp"

using namespace llvm;

// Register pass to igc-opt
#define PASS_FLAG "types-legalization-pass"
#define PASS_DESCRIPTION "Types Legalization pass"
#define PASS_CFG_ONLY false
#define PASS_ANALYSIS false
#define MAX_ARRAY_SIZE_THRESHOLD 8

IGC_INITIALIZE_PASS_BEGIN( TypesLegalizationPass,PASS_FLAG,PASS_DESCRIPTION,PASS_CFG_ONLY,PASS_ANALYSIS )
IGC_INITIALIZE_PASS_END( TypesLegalizationPass,PASS_FLAG,PASS_DESCRIPTION,PASS_CFG_ONLY,PASS_ANALYSIS )
char TypesLegalizationPass::ID = 0;

TypesLegalizationPass::TypesLegalizationPass()
  : FunctionPass( TypesLegalizationPass::ID ) {
  initializeTypesLegalizationPassPass( *PassRegistry::getPassRegistry() );
}

void TypesLegalizationPass::visitExtractValueInst( ExtractValueInst &ev ) {
  m_ExtractValueInst.push_back(&ev);
}

void TypesLegalizationPass::visitStoreInst( StoreInst &storeInst ) {
  Value *arg = storeInst.getOperand( 0 );
  if(arg != NULL)
  {
    Type *type = arg->getType();

    if((type != NULL) && type->isAggregateType()) {
      m_StoreInst.push_back( &storeInst );
    }
  }
}

void TypesLegalizationPass::visitPHINode( PHINode &phi ) {
  Type *type = phi.getType();

  if((type != NULL) && type->isAggregateType()) {
    m_PhiNodes.push_back( &phi );
  }
}

///////////////////////////////////////////////////////////////////////
/// @brief Creates an alloca instruction at the beginning of current
/// function
/// @param phi instruction to create the alloca for
AllocaInst* TypesLegalizationPass::CreateAlloca( Instruction *inst ) {
  PHINode *phi = dyn_cast<PHINode>(inst);
  if( phi != NULL)
  {
    IRBuilder<> builder( phi );
    Function *curFunc = builder.GetInsertBlock()->getParent();
    BasicBlock &firstBB = curFunc->getEntryBlock();
    BasicBlock::iterator firstInsertPoint = firstBB.getFirstInsertionPt();

    builder.SetInsertPoint( &(*firstInsertPoint) );
    Type *allocaType = phi->getType();
    return builder.CreateAlloca( allocaType );
  }
  else
  {
    return NULL;
  }
}

///////////////////////////////////////////////////////////////////////
/// @brief Creates GEP instruction.
/// @param  builder llvm IR builder to use
/// @param  Ty element type
/// @param  ptr memory pointer to
/// @param  indices set of constant indices to use in GEP
/// @return Value* - GetElementPtrInst* instruction.
Value* TypesLegalizationPass::CreateGEP(IGCLLVM::IRBuilder<>& builder, Type* Ty, Value* ptr, SmallVector<unsigned, 8>& indices) {
    SmallVector< Value*, 8> gepIndices;
    gepIndices.reserve(indices.size() + 1);
    gepIndices.push_back(builder.getInt32(0));
    for (unsigned idx : indices) {
        gepIndices.push_back(builder.getInt32(idx));
    }
    return builder.CreateGEP(Ty, ptr, gepIndices);
}

///////////////////////////////////////////////////////////////////////
/// @brief Resolves a PHI node to alloca store and load.
/// @param phi instruction to resolve

void TypesLegalizationPass::ResolvePhiNode( PHINode *phi ) {
  AllocaInst *allocaPtr = TypesLegalizationPass::CreateAlloca( phi );
  if(allocaPtr != NULL)
  {
    for(unsigned i = 0; i < phi->getNumOperands(); ++i)
    {
      IGCLLVM::IRBuilder<> builder( phi->getIncomingBlock( i )->getTerminator() );
      Value* source = phi->getOperand( i );
      StoreInst* newStore = builder.CreateStore( source,allocaPtr );
      m_StoreInst.push_back( newStore );
    }
    IGCLLVM::IRBuilder<> builder( phi );
    BasicBlock *block = builder.GetInsertBlock();
    builder.SetInsertPoint( &(*block->getFirstInsertionPt()) );
    Value* newLoad = builder.CreateLoad(allocaPtr->getAllocatedType(), allocaPtr);
    phi->replaceAllUsesWith( newLoad );
    phi->eraseFromParent();
  }
}

///////////////////////////////////////////////////////////////////////
/// @brief Resolves the use of composite types (i.e. structures or
/// arrays of vectors) in store and PHI nodes.
/// PHI node are resolved to alloca store and loads. Store instructions
/// of composite values are resolved to GEP + store of individual
/// elements of composite type.
bool TypesLegalizationPass::LegalizeTypes() {
  bool modified = false;

  // Handle PHI nodes before store instruction as the code creates
  // new store instructions.
  for(PHINode *phi : m_PhiNodes) {
    TypesLegalizationPass::ResolvePhiNode( phi );
    modified = true;
  }

  for(ExtractValueInst *ev : m_ExtractValueInst) {
    TypesLegalizationPass::ResolveExtractValue( ev );
    modified = true;
  }

  for(StoreInst *st : m_StoreInst) {
    TypesLegalizationPass::ResolveStoreInst( st );
    modified = true;
  }
  return modified;
}

///////////////////////////////////////////////////////////////////////
/// @brief Checks if an array contains null constants
/// @param  storeInst is an instruction that stores array
/// @param  indices leading to current structure filed or array elemen
/// @param  numelems is a number of elements in the array
bool TypesLegalizationPass::CheckNullArray(Instruction* storeInst) {
    Value* op0 = storeInst->getOperand(0);
    Type* type = op0->getType();

    bool isElemVectorType = type->getArrayElementType()->isVectorTy();
    bool isConstantAggregateZero = isa<ConstantAggregateZero>(op0);

    if (isElemVectorType || !isConstantAggregateZero)
      return false;

    return true;
}

///////////////////////////////////////////////////////////////////////
/// @brief Finds a structure field or array element pointed by indices
/// @param  ip insert point for new instructions
/// @param  val is a structure or array to extract the elements from
/// @param  indices specify array element or structure field to get
/// @return    Value* array element or structure field
Value *
TypesLegalizationPass::ResolveValue( Instruction *ip,Value *val,SmallVector<unsigned,8> &indices ) {

  if(InsertValueInst *civ = dyn_cast<InsertValueInst>(val)) {
    ArrayRef<unsigned> civIndices = civ->getIndices();
    if(indices.size() < civIndices.size()) {
      // We have less indices collected than it is needed to reference scalar value
      // inserted here. We will skip this instruction and anticipate there is further
      // ExtractValue further that will specify missing indices.
      return nullptr;
    }
    unsigned i = 0;
    for(; i < civIndices.size(); ++i) {
      if(civIndices[i] != indices[i]) {
        break;
      }
    }
    if(i == civIndices.size()) {
      if(i == indices.size()) {
        return civ->getInsertedValueOperand();
      }
      SmallVector<unsigned,8> ri( &indices[i],indices.end() );
      return ResolveValue( civ,civ->getInsertedValueOperand(),ri );
    }
    return ResolveValue( civ,civ->getAggregateOperand(),indices );
  }
  else if(LoadInst* ld = dyn_cast<LoadInst>(val))
  {
    IGCLLVM::IRBuilder<> builder( ld );
    Value* gep = CreateGEP( builder, ld->getType(), ld->getOperand( 0 ),indices );
    auto alignment = IGCLLVM::getAlignmentValue(ld);
    unsigned pointerTypeSize = ld->getType()->getScalarSizeInBits() / 8;
    Type* gepTy = IGCLLVM::getGEPIndexedType(ld->getType(), indices);
    if ( alignment && (alignment_t)pointerTypeSize == alignment )
      return builder.CreateAlignedLoad(gepTy, gep, IGCLLVM::getAlign(alignment));
    return builder.CreateLoad(gepTy, gep);
  }
  else if(Constant *c = dyn_cast<Constant>(val)) {
    IRBuilder<> builder( ip );
    // Create ExtractValue - it will be folded.
    return builder.CreateExtractValue( c,indices );
  }
  else if(ExtractValueInst* ev =
    dyn_cast<ExtractValueInst>(val))
  {
    auto evi = ev->getIndices();
    SmallVector<unsigned,8> newIndices( evi.begin(),evi.end() );
    newIndices.append( &indices.front(),&indices.back() );

    return ResolveValue( ev,ev->getAggregateOperand(),newIndices );
  }
  else if (auto *II = dyn_cast<IntrinsicInst>(val))
  {
      switch (II->getIntrinsicID())
      {
      case Intrinsic::uadd_with_overflow:
      case Intrinsic::sadd_with_overflow:
      case Intrinsic::ssub_with_overflow:
      case Intrinsic::usub_with_overflow:
      case Intrinsic::smul_with_overflow:
      case Intrinsic::umul_with_overflow:
      {
          // This gets handled in the legalizer.
          IRBuilder<> builder(ip);
          return builder.CreateExtractValue(val, indices);
      }
      default:
          IGC_ASSERT_MESSAGE(0, "Unsupported intrinsic instruction!");
          break;
      }
  }
  else if (isa<CallInst>(val) && (val->getType()->isStructTy() || val->getType()->isArrayTy())) {
    // Handle struct and array types of call instructions return value
    IRBuilder<> builder(ip);
    return builder.CreateExtractValue(val, indices);
  }
  else if (isa<Argument>(val)) {

    IGC_ASSERT_MESSAGE(!val->getType()->isStructTy(), "Illegal IR. Structures are passed as a pointer to a struct with the byval attribute.!");

    if ((val->getType()->isArrayTy()))
    {
      IRBuilder<> builder(ip);
      return builder.CreateExtractValue(val, indices);
    }
  }
  else if (SelectInst* select = dyn_cast<SelectInst>(val))
  {
    IGC_ASSERT(3 == select->getNumOperands());

    Value* condition = select->getOperand(0);
    Value* whenTrue = select->getOperand(1);
    Value* whenFalse = select->getOperand(2);

    whenTrue = ResolveValue(select, whenTrue, indices);
    whenFalse = ResolveValue(select, whenFalse, indices);

    IRBuilder<> builder(select);
    return builder.CreateSelect(condition, whenTrue, whenFalse);
  }

  // What other kind of instruction can we have here?
  IGC_ASSERT_MESSAGE(0, "Unresolved instruction!");

  // Fallback by creating an ExtractValueInstr.
  IRBuilder<> builder( ip );
  return builder.CreateExtractValue( val,indices );
}

///////////////////////////////////////////////////////////////////////
/// @brief Resolves extract value instructions of composite types into
/// sequences of GEP + load
/// @param  evInst instruction being resolved
/// @param  type    Type* type of current strucuter field or array
/// element being resolved
/// @param  indices leading to current structure filed or array element
void TypesLegalizationPass::ResolveExtractValue( ExtractValueInst* evInst )
{
  IRBuilder<>  builder( evInst );
  auto evIndices = evInst->getIndices();
  SmallVector<unsigned,8> indices( evIndices.begin(),evIndices.end() );

  Value* val = ResolveValue( evInst,evInst->getOperand( 0 ),indices );
  if(val)
  {
    evInst->replaceAllUsesWith( val );
    evInst->eraseFromParent();
  }
}

///////////////////////////////////////////////////////////////////////
/// @brief Resolves store instructions of composite types into
/// sequences of GEP + store
/// @param  storeInst instruction to be resolved
/// @param  type    Type* type of current strucuter field or array
/// element being resolved
/// @param  indices leading to current structure filed or array element
void TypesLegalizationPass::ResolveStoreInst(
  StoreInst *storeInst,Type *type,
  SmallVector<unsigned,8> &indices ) {

  if(type->isStructTy()) {
    for(unsigned field = 0; field < type->getStructNumElements(); field++) {
      indices.push_back( field );
      ResolveStoreInst( storeInst,type->getStructElementType( field ),indices );
      indices.pop_back();
    }
  }
  else if(type->isArrayTy()) {
    auto num = type->getArrayNumElements();
    // Check if recursion depth is 0 (indices.size() == 0) and
    // create memset instead of multiple store instructions that writes nulls
    if (!indices.size() && num >= MAX_ARRAY_SIZE_THRESHOLD && CheckNullArray(storeInst)) {
        IRBuilder<> Builder(storeInst);
        unsigned int size = type->getArrayElementType()->getScalarSizeInBits() / 8;
        auto al = IGCLLVM::getAlignmentValue(storeInst);
        Builder.CreateMemSet(storeInst->getOperand(1), Builder.getInt8(0), Builder.getInt64(num * size), IGCLLVM::getAlign(al), storeInst->isVolatile());
    }
    else {
        for (unsigned elem = 0; elem < type->getArrayNumElements(); ++elem) {
            indices.push_back(elem);
            ResolveStoreInst(storeInst, type->getArrayElementType(), indices);
            indices.pop_back();
        }
    }
  }
  else {
    Value *val =
      ResolveValue( storeInst,storeInst->getOperand( 0 ),indices );
    if (val) {
        IGCLLVM::IRBuilder<> builder(storeInst);
        bool isPackedStruct = false;

        if (StructType* st = dyn_cast<StructType>(storeInst->getValueOperand()->getType()))
            isPackedStruct = st->isPacked();

        Value* pGEP = CreateGEP(builder, storeInst->getValueOperand()->getType(), storeInst->getOperand(1), indices);
        if (isPackedStruct)
            builder.CreateAlignedStore(val, pGEP, IGCLLVM::getAlign(1));
        else
            builder.CreateStore(val, pGEP);
    }
  }
}

///////////////////////////////////////////////////////////////////////
/// @brief Resolves store instructions of composite types into
/// sequences of GEP + store
/// @param  storeInst instruction to be resolved
void TypesLegalizationPass::ResolveStoreInst( StoreInst *storeInst ) {
  SmallVector<unsigned,8> indices;
  Value *arg = storeInst->getOperand( 0 );
  if(arg != NULL)
  {
    ResolveStoreInst(storeInst,arg->getType(),indices);
    storeInst->eraseFromParent();
  }
}

bool TypesLegalizationPass::runOnFunction( Function &function ) {
  this->visit( function );
  bool shaderModified = LegalizeTypes();
  m_StoreInst.clear();
  m_ExtractValueInst.clear();
  m_PhiNodes.clear();
  if (function.hasOptNone()) {
    // type-legalization may create dead load, clean up those loads
    // when DCE would not. Otherwise IGC cannot handle aggregate loads
    SmallVector<LoadInst*, 32> DeadLoads;
    for (inst_iterator I = inst_begin(function), E = inst_end(function);
        I != E; ++I) {
      Instruction * inst = &*I;
      if (LoadInst* Load = dyn_cast<LoadInst>(inst)) {
        if (!Load->isVolatile() && Load->use_empty() &&
          Load->getType()->isAggregateType()) {
          DeadLoads.push_back(Load);
        }
      }
    }
    for (LoadInst* ld : DeadLoads) {
      ld->eraseFromParent();
    }
  }
  return shaderModified;
}
