//===- ThinLTOBitcodeWriter.cpp - Bitcode writing pass for ThinLTO --------===//
//
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
//===----------------------------------------------------------------------===//

#include "llvm/Transforms/IPO/ThinLTOBitcodeWriter.h"
#include "llvm/Analysis/BasicAliasAnalysis.h"
#include "llvm/Analysis/BlockFrequencyInfo.h"
#include "llvm/Analysis/ModuleSummaryAnalysis.h"
#include "llvm/Analysis/ProfileSummaryInfo.h"
#include "llvm/Bitcode/BitcodeWriter.h"
#include "llvm/IR/Constants.h"
#include "llvm/IR/DebugInfo.h"
#include "llvm/IR/Instructions.h"
#include "llvm/IR/Intrinsics.h"
#include "llvm/IR/Module.h"
#include "llvm/IR/PassManager.h"
#include "llvm/Support/raw_ostream.h"
#include "llvm/Transforms/IPO.h"
#include "llvm/Transforms/IPO/FunctionAttrs.h"
#include "llvm/Transforms/IPO/FunctionImport.h"
#include "llvm/Transforms/IPO/LowerTypeTests.h"
#include "llvm/Transforms/Utils/Cloning.h"
#include "llvm/Transforms/Utils/ModuleUtils.h"
using namespace llvm;

namespace {

// Promote each local-linkage entity defined by ExportM and used by ImportM by
// changing visibility and appending the given ModuleId.
void promoteInternals(Module &ExportM, Module &ImportM, StringRef ModuleId,
                      SetVector<GlobalValue *> *PromoteExtra = nullptr) {
  DenseMap<const Comdat *, Comdat *> RenamedComdats;
  for (auto &ExportGV : ExportM.global_values()) {
    if (!ExportGV.hasLocalLinkage())
      continue;

    auto Name = ExportGV.getName();
    GlobalValue *ImportGV = nullptr;
    const bool MustPromote = PromoteExtra && PromoteExtra->count(&ExportGV);
    if (!MustPromote) {
      ImportGV = ImportM.getNamedValue(Name);
      if (!ImportGV)
        continue;
      ImportGV->removeDeadConstantUsers();
      if (ImportGV->use_empty()) {
        ImportGV->eraseFromParent();
        continue;
      }
    }

    std::string OldName = Name.str();
    std::string NewName = (Name + ModuleId).str();

    if (const auto *C = ExportGV.getComdat())
      if (C->getName() == Name)
        RenamedComdats.try_emplace(C, ExportM.getOrInsertComdat(NewName));

    Constant *Aliasee = &ExportGV;
    while (auto *GA = dyn_cast<GlobalAlias>(Aliasee))
      Aliasee = GA->getAliasee();

    // We must use the function's value type (FunctionType), not ptr - hence
    // ExportGV.getValueType() rather than getType(). Otherwise, when an
    // internal coroutine is imported into another module, IRMover sees a
    // non-function value type for the unimported alias and materializes it as
    // an external GlobalVariable rather than a Function declaration. That
    // violates the verifier requirement that the coroutine argument of
    // @llvm.coro.id must refer to a function. We could "pierce through" by
    // stripping pointers, and cases other than coro do that, but this is
    // cleaner.
    auto *ExternalAlias = GlobalAlias::create(
        ExportGV.getValueType(), ExportGV.getAddressSpace(),
        GlobalValue::ExternalLinkage, NewName, Aliasee, &ExportM);
    ExternalAlias->setVisibility(GlobalValue::HiddenVisibility);

    if (MustPromote) {
      ExportGV.replaceUsesWithIf(
          ExternalAlias, [](Use &U) { return !isa<GlobalAlias>(U.getUser()); });
      PromoteExtra->remove(&ExportGV);
      PromoteExtra->insert(ExternalAlias);
    }

    if (ImportGV) {
      ImportGV->setName(NewName);
      ImportGV->setVisibility(GlobalValue::HiddenVisibility);
      ImportGV->reassignGUID();
    }
  }

  if (!RenamedComdats.empty())
    for (auto &GO : ExportM.global_objects())
      if (auto *C = GO.getComdat()) {
        auto Replacement = RenamedComdats.find(C);
        if (Replacement != RenamedComdats.end())
          GO.setComdat(Replacement->second);
      }
}

// Promote all internal (i.e. distinct) type ids used by the module by replacing
// them with external type ids formed using the module id.
//
// Note that this needs to be done before we clone the module because each clone
// will receive its own set of distinct metadata nodes.
void promoteTypeIds(Module &M, StringRef ModuleId) {
  DenseMap<Metadata *, Metadata *> LocalToGlobal;
  auto ExternalizeTypeId = [&](CallInst *CI, unsigned ArgNo) {
    Metadata *MD =
        cast<MetadataAsValue>(CI->getArgOperand(ArgNo))->getMetadata();

    if (isa<MDNode>(MD) && cast<MDNode>(MD)->isDistinct()) {
      Metadata *&GlobalMD = LocalToGlobal[MD];
      if (!GlobalMD) {
        std::string NewName = (Twine(LocalToGlobal.size()) + ModuleId).str();
        GlobalMD = MDString::get(M.getContext(), NewName);
      }

      CI->setArgOperand(ArgNo,
                        MetadataAsValue::get(M.getContext(), GlobalMD));
    }
  };

  if (Function *TypeTestFunc =
          Intrinsic::getDeclarationIfExists(&M, Intrinsic::type_test)) {
    for (const Use &U : TypeTestFunc->uses()) {
      auto CI = cast<CallInst>(U.getUser());
      ExternalizeTypeId(CI, 1);
    }
  }

  if (Function *PublicTypeTestFunc =
          Intrinsic::getDeclarationIfExists(&M, Intrinsic::public_type_test)) {
    for (const Use &U : PublicTypeTestFunc->uses()) {
      auto CI = cast<CallInst>(U.getUser());
      ExternalizeTypeId(CI, 1);
    }
  }

  if (Function *TypeCheckedLoadFunc =
          Intrinsic::getDeclarationIfExists(&M, Intrinsic::type_checked_load)) {
    for (const Use &U : TypeCheckedLoadFunc->uses()) {
      auto CI = cast<CallInst>(U.getUser());
      ExternalizeTypeId(CI, 2);
    }
  }

  if (Function *TypeCheckedLoadRelativeFunc = Intrinsic::getDeclarationIfExists(
          &M, Intrinsic::type_checked_load_relative)) {
    for (const Use &U : TypeCheckedLoadRelativeFunc->uses()) {
      auto CI = cast<CallInst>(U.getUser());
      ExternalizeTypeId(CI, 2);
    }
  }

  for (GlobalObject &GO : M.global_objects()) {
    SmallVector<MDNode *, 1> MDs;
    GO.getMetadata(LLVMContext::MD_type, MDs);

    GO.eraseMetadata(LLVMContext::MD_type);
    for (auto *MD : MDs) {
      auto I = LocalToGlobal.find(MD->getOperand(1));
      if (I == LocalToGlobal.end()) {
        GO.addMetadata(LLVMContext::MD_type, *MD);
        continue;
      }
      GO.addMetadata(
          LLVMContext::MD_type,
          *MDNode::get(M.getContext(), {MD->getOperand(0), I->second}));
    }

    SmallVector<MDNode *, 1> CGMDs;
    GO.getMetadata(LLVMContext::MD_callgraph, CGMDs);

    GO.eraseMetadata(LLVMContext::MD_callgraph);
    for (auto *MD : CGMDs) {
      if (MD->getNumOperands() == 1) {
        auto I = LocalToGlobal.find(MD->getOperand(0));
        if (I == LocalToGlobal.end()) {
          GO.addMetadata(LLVMContext::MD_callgraph, *MD);
          continue;
        }
        GO.addMetadata(LLVMContext::MD_callgraph,
                       *MDNode::get(M.getContext(), {I->second}));
      }
    }
  }
}

// Drop unused globals, and drop type information from function declarations.
// FIXME: If we made functions typeless then there would be no need to do this.
void simplifyExternals(Module &M) {
  FunctionType *EmptyFT =
      FunctionType::get(Type::getVoidTy(M.getContext()), false);

  for (Function &F : llvm::make_early_inc_range(M)) {
    if (F.isDeclaration() && F.use_empty()) {
      F.eraseFromParent();
      continue;
    }

    if (!F.isDeclaration() || F.getFunctionType() == EmptyFT ||
        // Changing the type of an intrinsic may invalidate the IR.
        F.getName().starts_with("llvm."))
      continue;

    Function *NewF = Function::Create(EmptyFT, GlobalValue::ExternalLinkage,
                                      F.getAddressSpace(), "", &M);
    NewF->copyAttributesFrom(&F);
    // Only copy function attribtues.
    NewF->setAttributes(AttributeList::get(M.getContext(),
                                           AttributeList::FunctionIndex,
                                           F.getAttributes().getFnAttrs()));
    NewF->takeName(&F);
    NewF->setMetadata(LLVMContext::MD_guid,
                      F.getMetadata(LLVMContext::MD_guid));
    F.replaceAllUsesWith(NewF);
    F.eraseFromParent();
  }

  for (GlobalIFunc &I : llvm::make_early_inc_range(M.ifuncs())) {
    if (I.use_empty())
      I.eraseFromParent();
    else
      assert(I.getResolverFunction() && "ifunc misses its resolver function");
  }

  for (GlobalVariable &GV : llvm::make_early_inc_range(M.globals())) {
    if (GV.isDeclaration() && GV.use_empty()) {
      GV.eraseFromParent();
      continue;
    }
  }
}

static void
filterModule(Module *M,
             function_ref<bool(const GlobalValue *)> ShouldKeepDefinition) {
  std::vector<GlobalValue *> V;
  for (GlobalValue &GV : M->global_values())
    if (!ShouldKeepDefinition(&GV))
      V.push_back(&GV);

  for (GlobalValue *GV : V)
    if (!convertToDeclaration(*GV))
      GV->eraseFromParent();
}

void forEachVirtualFunction(Constant *C, function_ref<void(Function *)> Fn) {
  if (auto *F = dyn_cast<Function>(C))
    return Fn(F);
  if (isa<GlobalValue>(C))
    return;
  for (Value *Op : C->operands())
    forEachVirtualFunction(cast<Constant>(Op), Fn);
}

// Clone any @llvm[.compiler].used over to the new module and append
// values whose defs were cloned into that module.
static void cloneUsedGlobalVariables(const Module &SrcM, Module &DestM,
                                     bool CompilerUsed) {
  SmallVector<GlobalValue *, 4> Used, NewUsed;
  // First collect those in the llvm[.compiler].used set.
  collectUsedGlobalVariables(SrcM, Used, CompilerUsed);
  // Next build a set of the equivalent values defined in DestM.
  for (auto *V : Used) {
    auto *GV = DestM.getNamedValue(V->getName());
    if (GV && !GV->isDeclaration())
      NewUsed.push_back(GV);
  }
  // Finally, add them to a llvm[.compiler].used variable in DestM.
  if (CompilerUsed)
    appendToCompilerUsed(DestM, NewUsed);
  else
    appendToUsed(DestM, NewUsed);
}

#ifndef NDEBUG
static bool enableUnifiedLTO(Module &M) {
  bool UnifiedLTO = false;
  if (auto *MD =
          mdconst::extract_or_null<ConstantInt>(M.getModuleFlag("UnifiedLTO")))
    UnifiedLTO = MD->getZExtValue();
  return UnifiedLTO;
}
#endif

bool mustEmitToMergedModule(const GlobalValue *GV) {
  // The __cfi_check definition is filled in by the CrossDSOCFI pass which
  // runs only in the merged module.
  return GV->getName() == "__cfi_check";
}

// If it's possible to split M into regular and thin LTO parts, do so and write
// a multi-module bitcode file with the two parts to OS. Otherwise, write only a
// regular LTO bitcode file to OS.
void splitAndWriteThinLTOBitcode(
    raw_ostream &OS, raw_ostream *ThinLinkOS,
    function_ref<AAResults &(Function &)> AARGetter,
    function_ref<const BlockFrequencyInfo &(Function &)> BFIGetter, Module &M,
    const bool ShouldPreserveUseListOrder) {
  std::string ModuleId = getUniqueModuleId(&M);
  if (ModuleId.empty()) {
    assert(!enableUnifiedLTO(M));
    // We couldn't generate a module ID for this module, write it out as a
    // regular LTO module with an index for summary-based dead stripping.
    ProfileSummaryInfo PSI(M);
    M.addModuleFlag(Module::Error, "ThinLTO", uint32_t(0));
    ModuleSummaryIndex Index = buildModuleSummaryIndex(M, nullptr, &PSI);
    WriteBitcodeToFile(M, OS, ShouldPreserveUseListOrder, &Index,
                       /*UnifiedLTO=*/false);

    if (ThinLinkOS)
      // We don't have a ThinLTO part, but still write the module to the
      // ThinLinkOS if requested so that the expected output file is produced.
      WriteBitcodeToFile(M, *ThinLinkOS, ShouldPreserveUseListOrder, &Index,
                         /*UnifiedLTO=*/false);

    return;
  }

  promoteTypeIds(M, ModuleId);

  // Collect the set of virtual functions that are eligible for virtual constant
  // propagation. Each eligible function must not access memory, must return
  // an integer of width <=64 bits, must take at least one argument, must not
  // use its first argument (assumed to be "this") and all arguments other than
  // the first one must be of <=64 bit integer type.
  //
  // Note that we test whether this copy of the function is readnone, rather
  // than testing function attributes, which must hold for any copy of the
  // function, even a less optimized version substituted at link time. This is
  // sound because the virtual constant propagation optimizations effectively
  // inline all implementations of the virtual function into each call site,
  // rather than using function attributes to perform local optimization.
  DenseSet<const Function *> EligibleVirtualFns;
  // If any member of a comdat lives in MergedM, put all members of that
  // comdat in MergedM to keep the comdat together.
  DenseSet<const Comdat *> MergedMComdats;
  for (GlobalVariable &GV : M.globals())
    if (!GV.isDeclaration() && lowertypetests::hasTypeMetadata(GV)) {
      if (const auto *C = GV.getComdat())
        MergedMComdats.insert(C);
      forEachVirtualFunction(GV.getInitializer(), [&](Function *F) {
        auto *RT = dyn_cast<IntegerType>(F->getReturnType());
        if (!RT || RT->getBitWidth() > 64 || F->arg_empty() ||
            !F->arg_begin()->use_empty())
          return;
        for (auto &Arg : drop_begin(F->args())) {
          auto *ArgT = dyn_cast<IntegerType>(Arg.getType());
          if (!ArgT || ArgT->getBitWidth() > 64)
            return;
        }
        if (!F->isDeclaration() &&
            computeFunctionBodyMemoryAccess(*F, AARGetter(*F))
                .doesNotAccessMemory())
          EligibleVirtualFns.insert(F);
      });
    }

  ValueToValueMapTy VMap;
  std::unique_ptr<Module> MergedM(
      CloneModule(M, VMap, [&](const GlobalValue *GV) -> bool {
        if (const auto *C = GV->getComdat())
          if (MergedMComdats.count(C))
            return true;
        if (mustEmitToMergedModule(GV))
          return true;
        if (auto *F = dyn_cast<Function>(GV))
          return EligibleVirtualFns.count(F);
        if (auto *GVar =
                dyn_cast_or_null<GlobalVariable>(GV->getAliaseeObject()))
          return lowertypetests::hasTypeMetadata(*GVar);
        return false;
      }));
  StripDebugInfo(*MergedM);
  MergedM->removeModuleInlineAsm();

  // Clone any llvm.*used globals to ensure the included values are
  // not deleted.
  cloneUsedGlobalVariables(M, *MergedM, /*CompilerUsed*/ false);
  cloneUsedGlobalVariables(M, *MergedM, /*CompilerUsed*/ true);

  for (Function &F : *MergedM)
    if (!F.isDeclaration() && !mustEmitToMergedModule(&F)) {
      // Reset the linkage of all functions eligible for virtual constant
      // propagation. The canonical definitions live in the thin LTO module so
      // that they can be imported.
      F.setLinkage(GlobalValue::AvailableExternallyLinkage);
      F.setComdat(nullptr);
    }

  SetVector<GlobalValue *> CfiFunctions = lowertypetests::findCfiFunctions(M);

  // Remove all globals with type metadata, globals with comdats that live in
  // MergedM, and aliases pointing to such globals from the thin LTO module.
  filterModule(&M, [&](const GlobalValue *GV) {
    if (auto *GVar = dyn_cast_or_null<GlobalVariable>(GV->getAliaseeObject()))
      if (lowertypetests::hasTypeMetadata(*GVar))
        return false;
    if (const auto *C = GV->getComdat())
      if (MergedMComdats.count(C))
        return false;
    if (mustEmitToMergedModule(GV))
      return false;
    return true;
  });

  // CfiFunctions contains only symbols from M. promoteInternals tries to find
  // match values from its first argument (the "exporting module") in
  // CfiFunctions. So we only need CfiFunctions for the second promotion (M ->
  // MergedM)
  promoteInternals(*MergedM, M, ModuleId, nullptr);
  promoteInternals(M, *MergedM, ModuleId, &CfiFunctions);

  // FIXME: Try to re-use PSI from the original module here.
  ProfileSummaryInfo PSI(M);

  lowertypetests::createCfiMetadata(*MergedM, M, CfiFunctions.getArrayRef(),
                                    PSI, BFIGetter);

  simplifyExternals(*MergedM);

  // FIXME: Try to re-use BSI from the original module here.
  ModuleSummaryIndex Index = buildModuleSummaryIndex(M, nullptr, &PSI);

  // Mark the merged module as requiring full LTO. We still want an index for
  // it though, so that it can participate in summary-based dead stripping.
  MergedM->addModuleFlag(Module::Error, "ThinLTO", uint32_t(0));
  ModuleSummaryIndex MergedMIndex =
      buildModuleSummaryIndex(*MergedM, nullptr, &PSI);

  SmallVector<char, 0> Buffer;

  BitcodeWriter W(Buffer);
  // Save the module hash produced for the full bitcode, which will
  // be used in the backends, and use that in the minimized bitcode
  // produced for the full link.
  ModuleHash ModHash = {{0}};
  W.writeModule(M, ShouldPreserveUseListOrder, &Index,
                /*GenerateHash=*/true, &ModHash);
  W.writeModule(*MergedM, ShouldPreserveUseListOrder, &MergedMIndex);
  W.writeSymtab();
  W.writeStrtab();
  OS << Buffer;

  // If a minimized bitcode module was requested for the thin link, only
  // the information that is needed by thin link will be written in the
  // given OS (the merged module will be written as usual).
  if (ThinLinkOS) {
    Buffer.clear();
    BitcodeWriter W2(Buffer);
    StripDebugInfo(M);
    W2.writeThinLinkBitcode(M, Index, ModHash);
    W2.writeModule(*MergedM, /*ShouldPreserveUseListOrder=*/false,
                   &MergedMIndex);
    W2.writeSymtab();
    W2.writeStrtab();
    *ThinLinkOS << Buffer;
  }
}

// Check if the LTO Unit splitting has been enabled.
bool enableSplitLTOUnit(Module &M) {
  bool EnableSplitLTOUnit = false;
  if (auto *MD = mdconst::extract_or_null<ConstantInt>(
          M.getModuleFlag("EnableSplitLTOUnit")))
    EnableSplitLTOUnit = MD->getZExtValue();
  return EnableSplitLTOUnit;
}

// Returns whether this module needs to be split (if splitting is enabled).
bool requiresSplit(Module &M) {
  for (auto &GO : M.global_objects()) {
    if (GO.hasMetadata(LLVMContext::MD_type))
      return true;
    if (mustEmitToMergedModule(&GO))
      return true;
  }
  return false;
}

bool writeThinLTOBitcode(
    raw_ostream &OS, raw_ostream *ThinLinkOS,
    function_ref<AAResults &(Function &)> AARGetter,
    function_ref<const BlockFrequencyInfo &(Function &)> BFIGetter, Module &M,
    const ModuleSummaryIndex *Index, const bool ShouldPreserveUseListOrder) {
  std::unique_ptr<ModuleSummaryIndex> NewIndex = nullptr;
  // See if this module needs to be split. If so, we try to split it
  // or at least promote type ids to enable WPD.
  if (requiresSplit(M)) {
    if (enableSplitLTOUnit(M)) {
      splitAndWriteThinLTOBitcode(OS, ThinLinkOS, AARGetter, BFIGetter, M,
                                  ShouldPreserveUseListOrder);
      return true;
    }
    // Promote type ids as needed for index-based WPD.
    std::string ModuleId = getUniqueModuleId(&M);
    if (!ModuleId.empty()) {
      promoteTypeIds(M, ModuleId);
      // Need to rebuild the index so that it contains type metadata
      // for the newly promoted type ids.
      // FIXME: Probably should not bother building the index at all
      // in the caller of writeThinLTOBitcode (which does so via the
      // ModuleSummaryIndexAnalysis pass), since we have to rebuild it
      // anyway whenever there is type metadata (here or in
      // splitAndWriteThinLTOBitcode). Just always build it once via the
      // buildModuleSummaryIndex when Module(s) are ready.
      ProfileSummaryInfo PSI(M);
      NewIndex = std::make_unique<ModuleSummaryIndex>(
          buildModuleSummaryIndex(M, nullptr, &PSI));
      Index = NewIndex.get();
    }
  }

  // Write it out as an unsplit ThinLTO module.

  // Save the module hash produced for the full bitcode, which will
  // be used in the backends, and use that in the minimized bitcode
  // produced for the full link.
  ModuleHash ModHash = {{0}};
  WriteBitcodeToFile(M, OS, ShouldPreserveUseListOrder, Index,
                     /*GenerateHash=*/true, &ModHash);
  // If a minimized bitcode module was requested for the thin link, only
  // the information that is needed by thin link will be written in the
  // given OS.
  if (ThinLinkOS && Index)
    writeThinLinkBitcodeToFile(M, *ThinLinkOS, *Index, ModHash);
  return false;
}

} // anonymous namespace

PreservedAnalyses
llvm::ThinLTOBitcodeWriterPass::run(Module &M, ModuleAnalysisManager &AM) {
  FunctionAnalysisManager &FAM =
      AM.getResult<FunctionAnalysisManagerModuleProxy>(M).getManager();

  bool Changed = writeThinLTOBitcode(
      OS, ThinLinkOS,
      [&FAM](Function &F) -> AAResults & {
        return FAM.getResult<AAManager>(F);
      },
      [&FAM](Function &F) -> const BlockFrequencyInfo & {
        return FAM.getResult<BlockFrequencyAnalysis>(F);
      },
      M, &AM.getResult<ModuleSummaryIndexAnalysis>(M),
      ShouldPreserveUseListOrder);

  return Changed ? PreservedAnalyses::none() : PreservedAnalyses::all();
}