#include "LLVMWrapper.h"

#include "llvm-c/Core.h"
#include "llvm/ADT/ArrayRef.h"
#include "llvm/ADT/DenseSet.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/Analysis/Lint.h"
#include "llvm/Analysis/TargetLibraryInfo.h"
#if LLVM_VERSION_GE(22, 0)
#include "llvm/Analysis/RuntimeLibcallInfo.h"
#endif
#include "llvm/Bitcode/BitcodeWriter.h"
#include "llvm/Bitcode/BitcodeWriterPass.h"
#include "llvm/CodeGen/CommandFlags.h"
#include "llvm/IR/AssemblyAnnotationWriter.h"
#include "llvm/IR/AutoUpgrade.h"
#include "llvm/IR/LegacyPassManager.h"
#include "llvm/IR/PassManager.h"
#include "llvm/IR/Verifier.h"
#include "llvm/IRPrinter/IRPrintingPasses.h"
#include "llvm/LTO/LTO.h"
#include "llvm/MC/MCInstrInfo.h"
#include "llvm/MC/MCSubtargetInfo.h"
#include "llvm/MC/TargetRegistry.h"
#include "llvm/Object/ObjectFile.h"
#include "llvm/Passes/PassBuilder.h"
#if LLVM_VERSION_GE(22, 0)
#include "llvm/Plugins/PassPlugin.h"
#else
#include "llvm/Passes/PassPlugin.h"
#endif
#include "llvm/Passes/StandardInstrumentations.h"
#include "llvm/Support/CBindingWrapping.h"
#include "llvm/Support/FileSystem.h"
#include "llvm/Support/Program.h"
#include "llvm/Support/TimeProfiler.h"
#include "llvm/Support/VirtualFileSystem.h"
#include "llvm/Target/TargetMachine.h"
#include "llvm/TargetParser/Host.h"
#include "llvm/Transforms/IPO/FunctionImport.h"
#include "llvm/Transforms/IPO/Internalize.h"
#include "llvm/Transforms/IPO/LowerTypeTests.h"
#include "llvm/Transforms/IPO/ThinLTOBitcodeWriter.h"
#include "llvm/Transforms/Instrumentation/AddressSanitizer.h"
#include "llvm/Transforms/Instrumentation/DataFlowSanitizer.h"
#include "llvm/Transforms/Instrumentation/HWAddressSanitizer.h"
#include "llvm/Transforms/Instrumentation/InstrProfiling.h"
#include "llvm/Transforms/Instrumentation/MemorySanitizer.h"
#include "llvm/Transforms/Instrumentation/RealtimeSanitizer.h"
#include "llvm/Transforms/Instrumentation/ThreadSanitizer.h"
#include "llvm/Transforms/Scalar/AnnotationRemarks.h"
#if LLVM_VERSION_GE(23, 0)
#include "llvm/Transforms/Utils/AssignGUID.h"
#endif
#include "llvm/Transforms/Utils/CanonicalizeAliases.h"
#include "llvm/Transforms/Utils/FunctionImportUtils.h"
#include "llvm/Transforms/Utils/NameAnonGlobals.h"
#include <set>
#include <string>
#include <vector>

// Conditional includes prevent clang-format from fully sorting the list,
// so if any are needed, keep them separate down here.

using namespace llvm;

static codegen::RegisterCodeGenFlags CGF;

typedef struct LLVMOpaquePass *LLVMPassRef;
typedef struct LLVMOpaqueTargetMachine *LLVMTargetMachineRef;

DEFINE_STDCXX_CONVERSION_FUNCTIONS(Pass, LLVMPassRef)
DEFINE_STDCXX_CONVERSION_FUNCTIONS(TargetMachine, LLVMTargetMachineRef)

extern "C" void LLVMRustTimeTraceProfilerInitialize() {
  timeTraceProfilerInitialize(
      /* TimeTraceGranularity */ 0,
      /* ProcName */ "rustc");
}

extern "C" void LLVMRustTimeTraceProfilerFinishThread() {
  timeTraceProfilerFinishThread();
}

extern "C" void LLVMRustTimeTraceProfilerFinish(const char *FileName) {
  auto FN = StringRef(FileName);
  std::error_code EC;
  auto OS = raw_fd_ostream(FN, EC, sys::fs::CD_CreateAlways);

  timeTraceProfilerWrite(OS);
  timeTraceProfilerCleanup();
}

extern "C" bool LLVMRustHasFeature(LLVMTargetMachineRef TM,
                                   const char *Feature) {
  TargetMachine *Target = unwrap(TM);
#if LLVM_VERSION_GE(23, 0)
  const MCSubtargetInfo &MCInfo = Target->getMCSubtargetInfo();
#else
  const MCSubtargetInfo &MCInfo = *Target->getMCSubtargetInfo();
#endif
  return MCInfo.checkFeatures(std::string("+") + Feature);
}

/// Check whether the target has a specific assembly mnemonic like `ret` or
/// `nop`.
/// This should be fast enough but if its not we have to look into another
/// method of checking.
extern "C" bool LLVMRustTargetHasMnemonic(LLVMTargetMachineRef TM,
                                          const char *Mnemonic) {
  TargetMachine *Target = unwrap(TM);
  const MCInstrInfo *MII = Target->getMCInstrInfo();
  StringRef MnemonicRef(Mnemonic);

  for (unsigned i = 0; i < MII->getNumOpcodes(); i++) {
    StringRef Name = MII->getName(i);
    if (Name.equals_insensitive(MnemonicRef)) {
      return true;
    }
  }
  return false;
}

enum class LLVMRustCodeModel {
  Tiny,
  Small,
  Kernel,
  Medium,
  Large,
  None,
};

static std::optional<CodeModel::Model> fromRust(LLVMRustCodeModel Model) {
  switch (Model) {
  case LLVMRustCodeModel::Tiny:
    return CodeModel::Tiny;
  case LLVMRustCodeModel::Small:
    return CodeModel::Small;
  case LLVMRustCodeModel::Kernel:
    return CodeModel::Kernel;
  case LLVMRustCodeModel::Medium:
    return CodeModel::Medium;
  case LLVMRustCodeModel::Large:
    return CodeModel::Large;
  case LLVMRustCodeModel::None:
    return std::nullopt;
  default:
    report_fatal_error("Bad CodeModel.");
  }
}

enum class LLVMRustCodeGenOptLevel {
  None,
  Less,
  Default,
  Aggressive,
};

using CodeGenOptLevelEnum = llvm::CodeGenOptLevel;

static CodeGenOptLevelEnum fromRust(LLVMRustCodeGenOptLevel Level) {
  switch (Level) {
  case LLVMRustCodeGenOptLevel::None:
    return CodeGenOptLevelEnum::None;
  case LLVMRustCodeGenOptLevel::Less:
    return CodeGenOptLevelEnum::Less;
  case LLVMRustCodeGenOptLevel::Default:
    return CodeGenOptLevelEnum::Default;
  case LLVMRustCodeGenOptLevel::Aggressive:
    return CodeGenOptLevelEnum::Aggressive;
  default:
    report_fatal_error("Bad CodeGenOptLevel.");
  }
}

enum class LLVMRustPassBuilderOptLevel {
  O0,
  O1,
  O2,
  O3,
  Os,
  Oz,
};

static OptimizationLevel fromRust(LLVMRustPassBuilderOptLevel Level) {
  switch (Level) {
  case LLVMRustPassBuilderOptLevel::O0:
    return OptimizationLevel::O0;
  case LLVMRustPassBuilderOptLevel::O1:
    return OptimizationLevel::O1;
  case LLVMRustPassBuilderOptLevel::O2:
    return OptimizationLevel::O2;
  case LLVMRustPassBuilderOptLevel::O3:
    return OptimizationLevel::O3;
#if LLVM_VERSION_GE(23, 0)
  case LLVMRustPassBuilderOptLevel::Os:
    return OptimizationLevel::O2;
  case LLVMRustPassBuilderOptLevel::Oz:
    return OptimizationLevel::O2;
#else
  case LLVMRustPassBuilderOptLevel::Os:
    return OptimizationLevel::Os;
  case LLVMRustPassBuilderOptLevel::Oz:
    return OptimizationLevel::Oz;
#endif
  default:
    report_fatal_error("Bad PassBuilderOptLevel.");
  }
}

enum class LLVMRustRelocModel {
  Static,
  PIC,
  DynamicNoPic,
  ROPI,
  RWPI,
  ROPIRWPI,
};

static Reloc::Model fromRust(LLVMRustRelocModel RustReloc) {
  switch (RustReloc) {
  case LLVMRustRelocModel::Static:
    return Reloc::Static;
  case LLVMRustRelocModel::PIC:
    return Reloc::PIC_;
  case LLVMRustRelocModel::DynamicNoPic:
    return Reloc::DynamicNoPIC;
  case LLVMRustRelocModel::ROPI:
    return Reloc::ROPI;
  case LLVMRustRelocModel::RWPI:
    return Reloc::RWPI;
  case LLVMRustRelocModel::ROPIRWPI:
    return Reloc::ROPI_RWPI;
  }
  report_fatal_error("Bad RelocModel.");
}

enum class LLVMRustFloatABI {
  Default,
  Soft,
  Hard,
};

static FloatABI::ABIType fromRust(LLVMRustFloatABI RustFloatAbi) {
  switch (RustFloatAbi) {
  case LLVMRustFloatABI::Default:
    return FloatABI::Default;
  case LLVMRustFloatABI::Soft:
    return FloatABI::Soft;
  case LLVMRustFloatABI::Hard:
    return FloatABI::Hard;
  }
  report_fatal_error("Bad FloatABI.");
}

// Must match the layout of `rustc_codegen_llvm::llvm::ffi::CompressionKind`.
enum class LLVMRustCompressionKind {
  None = 0,
  Zlib = 1,
  Zstd = 2,
};

static llvm::DebugCompressionType fromRust(LLVMRustCompressionKind Kind) {
  switch (Kind) {
  case LLVMRustCompressionKind::None:
    return llvm::DebugCompressionType::None;
  case LLVMRustCompressionKind::Zlib:
    if (!llvm::compression::zlib::isAvailable()) {
      report_fatal_error("LLVMRustCompressionKind::Zlib not available");
    }
    return llvm::DebugCompressionType::Zlib;
  case LLVMRustCompressionKind::Zstd:
    if (!llvm::compression::zstd::isAvailable()) {
      report_fatal_error("LLVMRustCompressionKind::Zstd not available");
    }
    return llvm::DebugCompressionType::Zstd;
  }
  report_fatal_error("bad LLVMRustCompressionKind");
}

extern "C" void LLVMRustPrintTargetCPUs(LLVMTargetMachineRef TM,
                                        RustStringRef OutStr) {
  ArrayRef<SubtargetSubTypeKV> CPUTable =
#if LLVM_VERSION_GE(23, 0)
      unwrap(TM)->getMCSubtargetInfo().getAllProcessorDescriptions();
#else
      unwrap(TM)->getMCSubtargetInfo()->getAllProcessorDescriptions();
#endif
  auto OS = RawRustStringOstream(OutStr);

  // Just print a bare list of target CPU names, and let Rust-side code handle
  // the full formatting of `--print=target-cpus`.
  for (auto &CPU : CPUTable) {
#if LLVM_VERSION_GE(23, 0)
    OS << CPU.key() << "\n";
#else
    OS << CPU.Key << "\n";
#endif
  }
}

extern "C" size_t LLVMRustGetTargetFeaturesCount(LLVMTargetMachineRef TM) {
  const TargetMachine *Target = unwrap(TM);
#if LLVM_VERSION_GE(23, 0)
  const MCSubtargetInfo &MCInfo = Target->getMCSubtargetInfo();
#else
  const MCSubtargetInfo &MCInfo = *Target->getMCSubtargetInfo();
#endif
  const ArrayRef<SubtargetFeatureKV> FeatTable =
      MCInfo.getAllProcessorFeatures();
  return FeatTable.size();
}

extern "C" void LLVMRustGetTargetFeature(LLVMTargetMachineRef TM, size_t Index,
                                         const char **Feature,
                                         const char **Desc) {
  const TargetMachine *Target = unwrap(TM);
#if LLVM_VERSION_GE(23, 0)
  const MCSubtargetInfo &MCInfo = Target->getMCSubtargetInfo();
#else
  const MCSubtargetInfo &MCInfo = *Target->getMCSubtargetInfo();
#endif
  const ArrayRef<SubtargetFeatureKV> FeatTable =
      MCInfo.getAllProcessorFeatures();
#if LLVM_VERSION_GE(23, 0)
  const SubtargetFeatureKV &Feat = FeatTable[Index];
  *Feature = Feat.key();
  *Desc = Feat.desc();
#else
  const SubtargetFeatureKV Feat = FeatTable[Index];
  *Feature = Feat.Key;
  *Desc = Feat.Desc;
#endif
}

extern "C" const char *LLVMRustGetHostCPUName(size_t *OutLen) {
  StringRef Name = sys::getHostCPUName();
  *OutLen = Name.size();
  return Name.data();
}

extern "C" LLVMTargetMachineRef LLVMRustCreateTargetMachine(
    const char *TripleStr, const char *CPU, const char *Feature,
    const char *ABIStr, LLVMRustCodeModel RustCM, LLVMRustRelocModel RustReloc,
    LLVMRustCodeGenOptLevel RustOptLevel, LLVMRustFloatABI RustFloatABIType,
    bool FunctionSections, bool DataSections, bool UniqueSectionNames,
    bool TrapUnreachable, bool Singlethread, bool VerboseAsm,
    bool EmitStackSizeSection, bool RelaxELFRelocations, bool UseInitArray,
    const char *SplitDwarfFile, const char *OutputObjFile,
    LLVMRustCompressionKind DebugInfoCompression, bool UseEmulatedTls,
    bool UseWasmEH, uint64_t LargeDataThreshold) {

  auto OptLevel = fromRust(RustOptLevel);
  auto RM = fromRust(RustReloc);
  auto CM = fromRust(RustCM);
  auto FloatABIType = fromRust(RustFloatABIType);

  std::string Error;
  auto Trip = Triple(Triple::normalize(TripleStr));
  const llvm::Target *TheTarget = TargetRegistry::lookupTarget(Trip, Error);
  if (TheTarget == nullptr) {
    LLVMRustSetLastError(Error.c_str());
    return nullptr;
  }

  TargetOptions Options = codegen::InitTargetOptionsFromCodeGenFlags(Trip);

#if LLVM_VERSION_LT(24, 0)
  Options.FloatABIType = FloatABIType;
#endif
  Options.DataSections = DataSections;
  Options.FunctionSections = FunctionSections;
  Options.UniqueSectionNames = UniqueSectionNames;
  Options.MCOptions.AsmVerbose = VerboseAsm;
  // Always preserve comments that were written by the user
  Options.MCOptions.PreserveAsmComments = true;
  Options.MCOptions.ABIName = ABIStr;
  if (SplitDwarfFile) {
    Options.MCOptions.SplitDwarfFile = SplitDwarfFile;
  }
  if (OutputObjFile) {
    Options.ObjectFilenameForDebug = OutputObjFile;
  }
  // To avoid fatal errors, make sure the Rust-side code only passes a
  // compression kind that is known to be supported by this build of LLVM, via
  // `LLVMRustLLVMHasZlibCompression` and `LLVMRustLLVMHasZstdCompression`.
  Options.MCOptions.CompressDebugSections = fromRust(DebugInfoCompression);
  Options.MCOptions.X86RelaxRelocations = RelaxELFRelocations;
  Options.UseInitArray = UseInitArray;
  Options.EmulatedTLS = UseEmulatedTls;

  if (TrapUnreachable) {
    // Tell LLVM to codegen `unreachable` into an explicit trap instruction.
    // This limits the extent of possible undefined behavior in some cases, as
    // it prevents control flow from "falling through" into whatever code
    // happens to be laid out next in memory.
    Options.TrapUnreachable = true;
    // But don't emit traps after other traps or no-returns unnecessarily.
    // ...except for when targeting WebAssembly, because the NoTrapAfterNoreturn
    // option causes bugs in the LLVM WebAssembly backend. You should be able to
    // remove this check when Rust's minimum supported LLVM version is >= 18
    // https://github.com/llvm/llvm-project/pull/65876
    //
    // Also keep traps after noreturn calls on Windows, because the trap is
    // needed to keep the return address within the calling function's
    // .pdata range. Without it, RtlLookupFunctionEntry resolves the wrong
    // function and SEH unwinding (used for backtraces) terminates early.
    // See https://github.com/rust-lang/rust/issues/140489
    if (!Trip.isWasm() && !Trip.isOSWindows()) {
      Options.NoTrapAfterNoreturn = true;
    }
  }

  if (Singlethread) {
    Options.ThreadModel = ThreadModel::Single;
  }

  if (UseWasmEH)
    Options.ExceptionModel = ExceptionHandling::Wasm;

  Options.EmitStackSizeSection = EmitStackSizeSection;

  TargetMachine *TM = TheTarget->createTargetMachine(Trip, CPU, Feature,
                                                     Options, RM, CM, OptLevel);

  if (LargeDataThreshold != 0) {
    TM->setLargeDataThreshold(LargeDataThreshold);
  }

  return wrap(TM);
}

// Unfortunately, the LLVM C API doesn't provide a way to create the
// TargetLibraryInfo pass, so we use this method to do so.
extern "C" void LLVMRustAddLibraryInfo(LLVMTargetMachineRef T,
                                       LLVMPassManagerRef PMR, LLVMModuleRef M,
                                       bool DisableSimplifyLibCalls) {
  auto TargetTriple = Triple(unwrap(M)->getTargetTriple());
  TargetOptions *Options = &unwrap(T)->Options;
  auto TLII = TargetLibraryInfoImpl(TargetTriple);
  if (DisableSimplifyLibCalls)
    TLII.disableAllFunctions();
  unwrap(PMR)->add(new TargetLibraryInfoWrapperPass(TLII));
#if LLVM_VERSION_GE(24, 0)
  unwrap(PMR)->add(new RuntimeLibraryInfoWrapper(
      Options->ExceptionModel, Options->EABIVersion, Options->MCOptions.ABIName,
      Options->VecLib));
#elif LLVM_VERSION_GE(22, 0)
  unwrap(PMR)->add(new RuntimeLibraryInfoWrapper(
      TargetTriple, Options->ExceptionModel, Options->FloatABIType,
      Options->EABIVersion, Options->MCOptions.ABIName, Options->VecLib));
#endif
}

extern "C" void LLVMRustSetLLVMOptions(int Argc, char **Argv) {
  // Initializing the command-line options more than once is not allowed. So,
  // check if they've already been initialized. (This could happen if we're
  // being called from rustpkg, for example). If the arguments change, then
  // that's just kinda unfortunate.
  static bool Initialized = false;
  if (Initialized)
    return;
  Initialized = true;
  cl::ParseCommandLineOptions(Argc, Argv);
}

enum class LLVMRustFileType {
  AssemblyFile,
  ObjectFile,
};

static CodeGenFileType fromRust(LLVMRustFileType Type) {
  switch (Type) {
  case LLVMRustFileType::AssemblyFile:
    return CodeGenFileType::AssemblyFile;
  case LLVMRustFileType::ObjectFile:
    return CodeGenFileType::ObjectFile;
  default:
    report_fatal_error("Bad FileType.");
  }
}

extern "C" LLVMRustResult
LLVMRustWriteOutputFile(LLVMTargetMachineRef Target, LLVMPassManagerRef PMR,
                        LLVMModuleRef M, const char *Path, const char *DwoPath,
                        LLVMRustFileType RustFileType, bool VerifyIR) {
  llvm::legacy::PassManager *PM = unwrap<llvm::legacy::PassManager>(PMR);
  auto FileType = fromRust(RustFileType);

  std::string ErrorInfo;
  std::error_code EC;
  auto OS = raw_fd_ostream(Path, EC, sys::fs::OF_None);
  if (EC)
    ErrorInfo = EC.message();
  if (ErrorInfo != "") {
    LLVMRustSetLastError(ErrorInfo.c_str());
    return LLVMRustResult::Failure;
  }

  auto BOS = buffer_ostream(OS);
  if (DwoPath) {
    auto DOS = raw_fd_ostream(DwoPath, EC, sys::fs::OF_None);
    EC.clear();
    if (EC)
      ErrorInfo = EC.message();
    if (ErrorInfo != "") {
      LLVMRustSetLastError(ErrorInfo.c_str());
      return LLVMRustResult::Failure;
    }
    auto DBOS = buffer_ostream(DOS);
    unwrap(Target)->addPassesToEmitFile(*PM, BOS, &DBOS, FileType, !VerifyIR);
    PM->run(*unwrap(M));
  } else {
    unwrap(Target)->addPassesToEmitFile(*PM, BOS, nullptr, FileType, !VerifyIR);
    PM->run(*unwrap(M));
  }

  // Apparently `addPassesToEmitFile` adds a pointer to our on-the-stack output
  // stream (OS), so the only real safe place to delete this is here? Don't we
  // wish this was written in Rust?
  LLVMDisposePassManager(PMR);
  return LLVMRustResult::Success;
}

extern "C" typedef void (*LLVMRustSelfProfileBeforePassCallback)(
    void *,        // LlvmSelfProfiler
    const char *,  // pass name
    const char *); // IR name
extern "C" typedef void (*LLVMRustSelfProfileAfterPassCallback)(
    void *); // LlvmSelfProfiler

#if LLVM_VERSION_GE(24, 0)
std::string LLVMRustwrappedIrGetName(const llvm::IRUnitRef &WrappedIr) {
  if (const auto *Cast = dyn_cast<Module>(WrappedIr))
    return Cast->getName().str();
  if (const auto *Cast = dyn_cast<Function>(WrappedIr))
    return Cast->getName().str();
  if (const auto *Cast = dyn_cast<Loop>(WrappedIr))
    return Cast->getName().str();
  if (const auto *Cast = dyn_cast<LazyCallGraph::SCC>(WrappedIr))
    return Cast->getName();
#else
std::string LLVMRustwrappedIrGetName(const llvm::Any &WrappedIr) {
  if (const auto *Cast = any_cast<const Module *>(&WrappedIr))
    return (*Cast)->getName().str();
  if (const auto *Cast = any_cast<const Function *>(&WrappedIr))
    return (*Cast)->getName().str();
  if (const auto *Cast = any_cast<const Loop *>(&WrappedIr))
    return (*Cast)->getName().str();
  if (const auto *Cast = any_cast<const LazyCallGraph::SCC *>(&WrappedIr))
    return (*Cast)->getName();
#endif
  return "<UNKNOWN>";
}

void LLVMSelfProfileInitializeCallbacks(
    PassInstrumentationCallbacks &PIC, void *LlvmSelfProfiler,
    LLVMRustSelfProfileBeforePassCallback BeforePassCallback,
    LLVMRustSelfProfileAfterPassCallback AfterPassCallback) {
  PIC.registerBeforeNonSkippedPassCallback(
#if LLVM_VERSION_GE(24, 0)
      [LlvmSelfProfiler, BeforePassCallback](StringRef Pass,
                                             llvm::IRUnitRef Ir) {
#else
      [LlvmSelfProfiler, BeforePassCallback](StringRef Pass, llvm::Any Ir) {
#endif
        std::string PassName = Pass.str();
        std::string IrName = LLVMRustwrappedIrGetName(Ir);
        BeforePassCallback(LlvmSelfProfiler, PassName.c_str(), IrName.c_str());
      });

  PIC.registerAfterPassCallback(
#if LLVM_VERSION_GE(24, 0)
      [LlvmSelfProfiler,
       AfterPassCallback](StringRef Pass, llvm::IRUnitRef IR,
                          const PreservedAnalyses &Preserved) {
#else
      [LlvmSelfProfiler, AfterPassCallback](
          StringRef Pass, llvm::Any IR, const PreservedAnalyses &Preserved) {
#endif
        AfterPassCallback(LlvmSelfProfiler);
      });

  PIC.registerAfterPassInvalidatedCallback(
      [LlvmSelfProfiler,
       AfterPassCallback](StringRef Pass, const PreservedAnalyses &Preserved) {
        AfterPassCallback(LlvmSelfProfiler);
      });

#if LLVM_VERSION_GE(24, 0)
  PIC.registerBeforeAnalysisCallback([LlvmSelfProfiler, BeforePassCallback](
                                         StringRef Pass, llvm::IRUnitRef Ir) {
#else
  PIC.registerBeforeAnalysisCallback(
      [LlvmSelfProfiler, BeforePassCallback](StringRef Pass, llvm::Any Ir) {
#endif
    std::string PassName = Pass.str();
    std::string IrName = LLVMRustwrappedIrGetName(Ir);
    BeforePassCallback(LlvmSelfProfiler, PassName.c_str(), IrName.c_str());
  });

#if LLVM_VERSION_GE(24, 0)
  PIC.registerAfterAnalysisCallback([LlvmSelfProfiler, AfterPassCallback](
                                        StringRef Pass, llvm::IRUnitRef Ir) {
#else
  PIC.registerAfterAnalysisCallback(
      [LlvmSelfProfiler, AfterPassCallback](StringRef Pass, llvm::Any Ir) {
#endif
    AfterPassCallback(LlvmSelfProfiler);
  });
}

enum class LLVMRustOptStage {
  PreLinkNoLTO,
  PreLinkThinLTO,
  PreLinkFatLTO,
  ThinLTO,
  FatLTO,
};

struct LLVMRustSanitizerOptions {
  bool SanitizeAddress;
  bool SanitizeAddressRecover;
  bool SanitizeCFI;
  bool SanitizeDataFlow;
  char **SanitizeDataFlowABIList;
  size_t SanitizeDataFlowABIListLen;
  bool SanitizeKCFI;
  bool SanitizeMemory;
  bool SanitizeMemoryRecover;
  int SanitizeMemoryTrackOrigins;
  bool SanitizerRealtime;
  bool SanitizeThread;
  bool SanitizeHWAddress;
  bool SanitizeHWAddressRecover;
  bool SanitizeKernelAddress;
  bool SanitizeKernelAddressRecover;
  bool SanitizeKernelHWAddress;
  bool SanitizeKernelHWAddressRecover;
};

extern "C" typedef void (*registerEnzymeAndPassPipelineFn)(
    llvm::PassBuilder &PB, bool augment);

/// Forces the bitcode writer to emit full LTO summary instead of thin LTO
/// summary for embedded bitcode under Fat LTO.
///
/// Note the bitcode writer will only emit the full LTO block ID if the
/// "ThinLTO" metadata is defined and explicitly set to zero. Otherwise, the
/// thin LTO block ID will be emitted.
static void forceFullLTOSummary(Module *M) {
  // This function may be called twice, such as if you call it with `-C lto=fat
  // --emit=llvm-bc`, so exit early if if we've already set up the module to
  // emit full LTO summaries.
  if (auto *Existing = M->getModuleFlag("ThinLTO")) {
    auto *Const = mdconst::extract<ConstantInt>(Existing);
    assert(Const->getZExtValue() == 0 &&
           "ThinLTO flag already set to non-zero");
    return;
  }

  auto *Zero = ConstantInt::get(Type::getInt32Ty(M->getContext()), 0);
  M->addModuleFlag(Module::Error, "ThinLTO", Zero);
}

extern "C" LLVMRustResult LLVMRustOptimize(
    LLVMModuleRef ModuleRef, LLVMTargetMachineRef TMRef,
    LLVMRustPassBuilderOptLevel OptLevelRust, LLVMRustOptStage OptStage,
    bool IsLinkerPluginLTO, bool NoPrepopulatePasses, bool VerifyIR,
    bool LintIR, LLVMRustBuffer **ThinLTOBufferRef,
    LLVMRustBuffer **ThinLTOSummaryBufferRef, bool MergeFunctions,
    bool UnrollLoops, bool SLPVectorize, bool LoopVectorize,
    bool DisableSimplifyLibCalls, bool EmitLifetimeMarkers,
    registerEnzymeAndPassPipelineFn EnzymePtr, bool PrintBeforeEnzyme,
    bool PrintAfterEnzyme, bool PrintPasses,
    LLVMRustSanitizerOptions *SanitizerOptions, const char *PGOGenPath,
    const char *PGOUsePath, bool InstrumentCoverage,
    const char *InstrProfileOutput, const char *PGOSampleUsePath,
    bool DebugInfoForProfiling, void *LlvmSelfProfiler,
    LLVMRustSelfProfileBeforePassCallback BeforePassCallback,
    LLVMRustSelfProfileAfterPassCallback AfterPassCallback,
    const char *PostEnzymePasses, size_t PostEnzymePassesLen,
    const char *ExtraPasses, size_t ExtraPassesLen, const char *LLVMPlugins,
    size_t LLVMPluginsLen) {
  Module *TheModule = unwrap(ModuleRef);
  TargetMachine *TM = unwrap(TMRef);
  OptimizationLevel OptLevel = fromRust(OptLevelRust);

  PipelineTuningOptions PTO;
  PTO.LoopUnrolling = UnrollLoops;
  PTO.LoopInterleaving = UnrollLoops;
  PTO.LoopVectorization = LoopVectorize;
  PTO.SLPVectorization = SLPVectorize;
  PTO.MergeFunctions = MergeFunctions;

  PassInstrumentationCallbacks PIC;

  if (LlvmSelfProfiler) {
    LLVMSelfProfileInitializeCallbacks(PIC, LlvmSelfProfiler,
                                       BeforePassCallback, AfterPassCallback);
  }

  std::optional<PGOOptions> PGOOpt;
#if LLVM_VERSION_LT(22, 0)
  auto FS = vfs::getRealFileSystem();
#endif
  if (PGOGenPath) {
    assert(!PGOUsePath && !PGOSampleUsePath);
    PGOOpt = PGOOptions(
#if LLVM_VERSION_GE(22, 0)
        PGOGenPath, "", "", "", PGOOptions::IRInstr, PGOOptions::NoCSAction,
#else
        PGOGenPath, "", "", "", FS, PGOOptions::IRInstr, PGOOptions::NoCSAction,
#endif
        PGOOptions::ColdFuncOpt::Default, DebugInfoForProfiling);
  } else if (PGOUsePath) {
    assert(!PGOSampleUsePath);
    PGOOpt = PGOOptions(
#if LLVM_VERSION_GE(22, 0)
        PGOUsePath, "", "", "", PGOOptions::IRUse, PGOOptions::NoCSAction,
#else
        PGOUsePath, "", "", "", FS, PGOOptions::IRUse, PGOOptions::NoCSAction,
#endif
        PGOOptions::ColdFuncOpt::Default, DebugInfoForProfiling);
  } else if (PGOSampleUsePath) {
    PGOOpt =
#if LLVM_VERSION_GE(22, 0)
        PGOOptions(PGOSampleUsePath, "", "", "", PGOOptions::SampleUse,
#else
        PGOOptions(PGOSampleUsePath, "", "", "", FS, PGOOptions::SampleUse,
#endif
                   PGOOptions::NoCSAction, PGOOptions::ColdFuncOpt::Default,
                   DebugInfoForProfiling);
  } else if (DebugInfoForProfiling) {
    PGOOpt = PGOOptions(
#if LLVM_VERSION_GE(22, 0)
        "", "", "", "", PGOOptions::NoAction, PGOOptions::NoCSAction,
#else
        "", "", "", "", FS, PGOOptions::NoAction, PGOOptions::NoCSAction,
#endif
        PGOOptions::ColdFuncOpt::Default, DebugInfoForProfiling);
  }

  auto PB = PassBuilder(TM, PTO, PGOOpt, &PIC);
  LoopAnalysisManager LAM;
  FunctionAnalysisManager FAM;
  CGSCCAnalysisManager CGAM;
  ModuleAnalysisManager MAM;

  StandardInstrumentations SI(TheModule->getContext(),
                              /*DebugLogging=*/false);
  SI.registerCallbacks(PIC, &MAM);

  if (LLVMPluginsLen) {
    auto PluginsStr = StringRef(LLVMPlugins, LLVMPluginsLen);
    SmallVector<StringRef> Plugins;
    PluginsStr.split(Plugins, ',', -1, false);
    for (auto PluginPath : Plugins) {
      auto Plugin = PassPlugin::Load(PluginPath.str());
      if (!Plugin) {
        auto Err = Plugin.takeError();
        auto ErrMsg = llvm::toString(std::move(Err));
        LLVMRustSetLastError(ErrMsg.c_str());
        return LLVMRustResult::Failure;
      }
      Plugin->registerPassBuilderCallbacks(PB);
    }
  }

  FAM.registerPass([&] { return PB.buildDefaultAAPipeline(); });

  Triple TargetTriple(TheModule->getTargetTriple());
  std::unique_ptr<TargetLibraryInfoImpl> TLII(
      new TargetLibraryInfoImpl(TargetTriple));
  if (DisableSimplifyLibCalls)
    TLII->disableAllFunctions();
  FAM.registerPass([&] { return TargetLibraryAnalysis(*TLII); });

  PB.registerModuleAnalyses(MAM);
  PB.registerCGSCCAnalyses(CGAM);
  PB.registerFunctionAnalyses(FAM);
  PB.registerLoopAnalyses(LAM);
  PB.crossRegisterProxies(LAM, FAM, CGAM, MAM);

  // We manually collect pipeline callbacks so we can apply them at O0, where
  // the PassBuilder does not create a pipeline.
  std::vector<std::function<void(ModulePassManager &, OptimizationLevel)>>
      PipelineStartEPCallbacks;
  std::vector<std::function<void(ModulePassManager &, OptimizationLevel,
                                 ThinOrFullLTOPhase)>>
      OptimizerLastEPCallbacks;

  if (!IsLinkerPluginLTO && SanitizerOptions && SanitizerOptions->SanitizeCFI &&
      !NoPrepopulatePasses) {
    PipelineStartEPCallbacks.push_back(
        [](ModulePassManager &MPM, OptimizationLevel Level) {
          MPM.addPass(LowerTypeTestsPass(
              /*ExportSummary=*/nullptr,
              /*ImportSummary=*/nullptr));
        });
  }

  if (VerifyIR) {
    PipelineStartEPCallbacks.push_back(
        [VerifyIR](ModulePassManager &MPM, OptimizationLevel Level) {
          MPM.addPass(VerifierPass());
        });
  }

  if (LintIR) {
    PipelineStartEPCallbacks.push_back([](ModulePassManager &MPM,
                                          OptimizationLevel Level) {
      MPM.addPass(
          createModuleToFunctionPassAdaptor(LintPass(/*AbortOnError=*/true)));
    });
  }

  if (InstrumentCoverage) {
    PipelineStartEPCallbacks.push_back(
        [InstrProfileOutput](ModulePassManager &MPM, OptimizationLevel Level) {
          InstrProfOptions Options;
          if (InstrProfileOutput) {
            Options.InstrProfileOutput = InstrProfileOutput;
          }
          // cargo run tests in multhreading mode by default
          // so use atomics for coverage counters
          Options.Atomic = true;
          MPM.addPass(InstrProfilingLoweringPass(Options, false));
        });
  }

  if (SanitizerOptions) {
    if (SanitizerOptions->SanitizeDataFlow) {
      std::vector<std::string> ABIListFiles(
          SanitizerOptions->SanitizeDataFlowABIList,
          SanitizerOptions->SanitizeDataFlowABIList +
              SanitizerOptions->SanitizeDataFlowABIListLen);
      OptimizerLastEPCallbacks.push_back(
          [ABIListFiles](ModulePassManager &MPM, OptimizationLevel Level,
                         ThinOrFullLTOPhase phase) {
            MPM.addPass(DataFlowSanitizerPass(ABIListFiles));
          });
    }

    if (SanitizerOptions->SanitizeMemory) {
      MemorySanitizerOptions Options(
          SanitizerOptions->SanitizeMemoryTrackOrigins,
          SanitizerOptions->SanitizeMemoryRecover,
          /*CompileKernel=*/false,
          /*EagerChecks=*/true);
      OptimizerLastEPCallbacks.push_back([Options](ModulePassManager &MPM,
                                                   OptimizationLevel Level,
                                                   ThinOrFullLTOPhase phase) {
        MPM.addPass(MemorySanitizerPass(Options));
      });
    }

    if (SanitizerOptions->SanitizeThread) {
      OptimizerLastEPCallbacks.push_back([](ModulePassManager &MPM,
                                            OptimizationLevel Level,
                                            ThinOrFullLTOPhase phase) {
        MPM.addPass(ModuleThreadSanitizerPass());
        MPM.addPass(createModuleToFunctionPassAdaptor(ThreadSanitizerPass()));
      });
    }

    if (SanitizerOptions->SanitizeAddress ||
        SanitizerOptions->SanitizeKernelAddress) {
      OptimizerLastEPCallbacks.push_back([SanitizerOptions,
                                          TM](ModulePassManager &MPM,
                                              OptimizationLevel Level,
                                              ThinOrFullLTOPhase phase) {
        auto CompileKernel = SanitizerOptions->SanitizeKernelAddress;
        AddressSanitizerOptions opts = AddressSanitizerOptions{
            CompileKernel,
            SanitizerOptions->SanitizeAddressRecover ||
                SanitizerOptions->SanitizeKernelAddressRecover,
            /*UseAfterScope=*/true,
            AsanDetectStackUseAfterReturnMode::Runtime,
        };
        MPM.addPass(
            AddressSanitizerPass(opts,
                                 /*UseGlobalGC*/ true,
                                 // UseOdrIndicator should be false on windows
                                 // machines https://reviews.llvm.org/D137227
                                 !TM->getTargetTriple().isOSWindows()));
      });
    }
    if (SanitizerOptions->SanitizeHWAddress ||
        SanitizerOptions->SanitizeKernelHWAddress) {
      OptimizerLastEPCallbacks.push_back(
          [SanitizerOptions](ModulePassManager &MPM, OptimizationLevel Level,
                             ThinOrFullLTOPhase phase) {
            HWAddressSanitizerOptions opts(
                SanitizerOptions->SanitizeKernelHWAddress,
                SanitizerOptions->SanitizeHWAddressRecover ||
                    SanitizerOptions->SanitizeKernelHWAddressRecover,
                /*DisableOptimization=*/false);
            MPM.addPass(HWAddressSanitizerPass(opts));
          });
    }
    if (SanitizerOptions->SanitizerRealtime) {
      OptimizerLastEPCallbacks.push_back([](ModulePassManager &MPM,
                                            OptimizationLevel Level,
                                            ThinOrFullLTOPhase phase) {
        MPM.addPass(RealtimeSanitizerPass());
      });
    }
  }

  ModulePassManager MPM;
  bool NeedThinLTOBufferPasses = true;
  auto ThinLTOBuffer = std::make_unique<LLVMRustBuffer>();
  auto ThinLTOSummaryBuffer = std::make_unique<LLVMRustBuffer>();
  raw_string_ostream ThinLTODataOS(ThinLTOBuffer->data);
  raw_string_ostream ThinLinkDataOS(ThinLTOSummaryBuffer->data);
  bool IsLTO = OptStage == LLVMRustOptStage::ThinLTO ||
               OptStage == LLVMRustOptStage::FatLTO;
  if (PostEnzymePassesLen) {
    if (auto Err = PB.parsePassPipeline(
            MPM, StringRef(PostEnzymePasses, PostEnzymePassesLen))) {
      std::string ErrMsg = toString(std::move(Err));
      LLVMRustSetLastError(ErrMsg.c_str());
      return LLVMRustResult::Failure;
    }
  } else {
    if (!NoPrepopulatePasses) {
      for (const auto &C : PipelineStartEPCallbacks)
        PB.registerPipelineStartEPCallback(C);
      for (const auto &C : OptimizerLastEPCallbacks)
        PB.registerOptimizerLastEPCallback(C);

      // The pre-link pipelines don't support O0 and require using
      // buildO0DefaultPipeline() instead. At the same time, the LTO pipelines
      // do support O0 and using them is required.
      if (OptLevel == OptimizationLevel::O0 && !IsLTO) {
        // We manually schedule ThinLTOBufferPasses below, so don't pass the
        // value to enable it here.
        MPM = PB.buildO0DefaultPipeline(OptLevel);
      } else {
        switch (OptStage) {
        case LLVMRustOptStage::PreLinkNoLTO:
          if (ThinLTOBufferRef) {
            // This is similar to LLVM's `buildFatLTODefaultPipeline`, where the
            // bitcode for embedding is obtained after performing
            // `ThinLTOPreLinkDefaultPipeline`.
            MPM.addPass(PB.buildThinLTOPreLinkDefaultPipeline(OptLevel));
            MPM.addPass(ThinLTOBitcodeWriterPass(
                ThinLTODataOS,
                ThinLTOSummaryBufferRef ? &ThinLinkDataOS : nullptr));
            *ThinLTOBufferRef = ThinLTOBuffer.release();
            if (ThinLTOSummaryBufferRef) {
              *ThinLTOSummaryBufferRef = ThinLTOSummaryBuffer.release();
            }
            MPM.addPass(PB.buildModuleOptimizationPipeline(
                OptLevel, ThinOrFullLTOPhase::None));
            MPM.addPass(
                createModuleToFunctionPassAdaptor(AnnotationRemarksPass()));
          } else {
            MPM = PB.buildPerModuleDefaultPipeline(OptLevel);
          }
          break;
        case LLVMRustOptStage::PreLinkThinLTO:
        case LLVMRustOptStage::PreLinkFatLTO:
          MPM = PB.buildThinLTOPreLinkDefaultPipeline(OptLevel);
          NeedThinLTOBufferPasses = false;
          break;
        case LLVMRustOptStage::ThinLTO:
          // FIXME: Does it make sense to pass the ModuleSummaryIndex?
          // It only seems to be needed for C++ specific optimizations.
          MPM = PB.buildThinLTODefaultPipeline(OptLevel, nullptr);
          break;
        case LLVMRustOptStage::FatLTO:
          MPM = PB.buildLTODefaultPipeline(OptLevel, nullptr);
          NeedThinLTOBufferPasses = false;
          break;
        }
      }
    } else {
      // We're not building any of the default pipelines but we still want to
      // add the verifier, instrumentation, etc passes if they were requested
      for (const auto &C : PipelineStartEPCallbacks)
        C(MPM, OptLevel);
      for (const auto &C : OptimizerLastEPCallbacks)
        C(MPM, OptLevel, ThinOrFullLTOPhase::None);
    }

    if (ExtraPassesLen) {
      if (auto Err = PB.parsePassPipeline(
              MPM, StringRef(ExtraPasses, ExtraPassesLen))) {
        std::string ErrMsg = toString(std::move(Err));
        LLVMRustSetLastError(ErrMsg.c_str());
        return LLVMRustResult::Failure;
      }
    }

    if (NeedThinLTOBufferPasses) {
      MPM.addPass(CanonicalizeAliasesPass());
      MPM.addPass(NameAnonGlobalPass());
#if LLVM_VERSION_GE(23, 0)
      MPM.addPass(AssignGUIDPass());
#endif
    }
    // For `-Copt-level=0`, and the pre-link fat/thin LTO stages.
    if (ThinLTOBufferRef && *ThinLTOBufferRef == nullptr) {
      // thin lto summaries prevent fat lto, so emit a full summary instead if
      // fat lto is requested. See PR #136840 for background information.
      if (OptStage != LLVMRustOptStage::PreLinkFatLTO) {
        MPM.addPass(ThinLTOBitcodeWriterPass(
            ThinLTODataOS,
            ThinLTOSummaryBufferRef ? &ThinLinkDataOS : nullptr));
      } else {
        forceFullLTOSummary(TheModule);
        MPM.addPass(BitcodeWriterPass(ThinLTODataOS,
                                      /*ShouldPreserveUseListOrder=*/false,
                                      /*EmitSummaryIndex=*/true));
      }
      *ThinLTOBufferRef = ThinLTOBuffer.release();
      if (ThinLTOSummaryBufferRef) {
        *ThinLTOSummaryBufferRef = ThinLTOSummaryBuffer.release();
      }
    }

    // now load "-enzyme" pass:
    // With dlopen, ENZYME macro may not be defined, so check EnzymePtr directly
    // In the case of debug builds with multiple codegen units, we might not
    // have all function definitions available during the early compiler
    // invocations. We therefore wait for the final lto step to run Enzyme.
    if (EnzymePtr && IsLTO) {

      if (PrintBeforeEnzyme) {
        // Handle the Rust flag `-Zautodiff=PrintModBefore`.
        std::string Banner = "Module before EnzymeNewPM";
        MPM.addPass(PrintModulePass(outs(), Banner, true, false));
      }

      EnzymePtr(PB, false);
      if (auto Err = PB.parsePassPipeline(MPM, "enzyme")) {
        std::string ErrMsg = toString(std::move(Err));
        LLVMRustSetLastError(ErrMsg.c_str());
        return LLVMRustResult::Failure;
      }

      if (PrintAfterEnzyme) {
        // Handle the Rust flag `-Zautodiff=PrintModAfter`.
        std::string Banner = "Module after EnzymeNewPM";
        MPM.addPass(PrintModulePass(outs(), Banner, true, false));
      }
    }
  }

  if (PrintPasses) {
    // Print all passes from the PM:
    std::string Pipeline;
    raw_string_ostream SOS(Pipeline);
    MPM.printPipeline(SOS, [&PIC](StringRef ClassName) {
      auto PassName = PIC.getPassNameForClassName(ClassName);
      return PassName.empty() ? ClassName : PassName;
    });
    outs() << Pipeline;
    outs() << "\n";
  }

  // Upgrade all calls to old intrinsics first.
  for (Module::iterator I = TheModule->begin(), E = TheModule->end(); I != E;)
    UpgradeCallsToIntrinsic(&*I++); // must be post-increment, as we remove

  MPM.run(*TheModule, MAM);
  return LLVMRustResult::Success;
}

// Callback to demangle function name
// Parameters:
// * name to be demangled
// * name len
// * output buffer
// * output buffer len
// Returns len of demangled string, or 0 if demangle failed.
typedef size_t (*DemangleFn)(const char *, size_t, char *, size_t);

namespace {

class RustAssemblyAnnotationWriter : public AssemblyAnnotationWriter {
  DemangleFn Demangle;
  std::vector<char> Buf;

public:
  RustAssemblyAnnotationWriter(DemangleFn Demangle) : Demangle(Demangle) {}

  // Return empty string if demangle failed
  // or if name does not need to be demangled
  StringRef CallDemangle(StringRef name) {
    if (!Demangle) {
      return StringRef();
    }

    if (Buf.size() < name.size() * 2) {
      // Semangled name usually shorter than mangled,
      // but allocate twice as much memory just in case
      Buf.resize(name.size() * 2);
    }

    auto R = Demangle(name.data(), name.size(), Buf.data(), Buf.size());
    if (!R) {
      // Demangle failed.
      return StringRef();
    }

    auto Demangled = StringRef(Buf.data(), R);
    if (Demangled == name) {
      // Do not print anything if demangled name is equal to mangled.
      return StringRef();
    }

    return Demangled;
  }

  void emitFunctionAnnot(const Function *F,
                         formatted_raw_ostream &OS) override {
    StringRef Demangled = CallDemangle(F->getName());
    if (Demangled.empty()) {
      return;
    }

    OS << "; " << Demangled << "\n";
  }

  void emitInstructionAnnot(const Instruction *I,
                            formatted_raw_ostream &OS) override {
    const char *Name;
    const Value *Value;
    if (const CallInst *CI = dyn_cast<CallInst>(I)) {
      Name = "call";
      Value = CI->getCalledOperand();
    } else if (const InvokeInst *II = dyn_cast<InvokeInst>(I)) {
      Name = "invoke";
      Value = II->getCalledOperand();
    } else {
      // Could demangle more operations, e. g.
      // `store %place, @function`.
      return;
    }

    if (!Value->hasName()) {
      return;
    }

    StringRef Demangled = CallDemangle(Value->getName());
    if (Demangled.empty()) {
      return;
    }

    OS << "; " << Name << " " << Demangled << "\n";
  }
};

} // namespace

extern "C" LLVMRustResult LLVMRustPrintModule(LLVMModuleRef M, const char *Path,
                                              DemangleFn Demangle) {
  std::string ErrorInfo;
  std::error_code EC;
  auto OS = raw_fd_ostream(Path, EC, sys::fs::OF_None);
  if (EC)
    ErrorInfo = EC.message();
  if (ErrorInfo != "") {
    LLVMRustSetLastError(ErrorInfo.c_str());
    return LLVMRustResult::Failure;
  }

  auto AAW = RustAssemblyAnnotationWriter(Demangle);
  auto FOS = formatted_raw_ostream(OS);
  unwrap(M)->print(FOS, &AAW);

  return LLVMRustResult::Success;
}

extern "C" void LLVMRustPrintPasses() {
  PassBuilder PB;
  PB.printPassNames(outs());
}

extern "C" void LLVMRustRunRestrictionPass(LLVMModuleRef M, char **Symbols,
                                           size_t Len) {
  auto PreserveFunctions = [=](const GlobalValue &GV) {
    // Preserve symbols exported from Rust modules.
    for (size_t I = 0; I < Len; I++) {
      if (GV.getName() == Symbols[I]) {
        return true;
      }
    }
    return false;
  };

  internalizeModule(*unwrap(M), PreserveFunctions);
}

extern "C" void
LLVMRustSetDataLayoutFromTargetMachine(LLVMModuleRef Module,
                                       LLVMTargetMachineRef TMR) {
  TargetMachine *Target = unwrap(TMR);
  unwrap(Module)->setDataLayout(Target->createDataLayout());
}

extern "C" void LLVMRustSetModulePICLevel(LLVMModuleRef M) {
  unwrap(M)->setPICLevel(PICLevel::Level::BigPIC);
}

extern "C" void LLVMRustSetModulePIELevel(LLVMModuleRef M) {
  unwrap(M)->setPIELevel(PIELevel::Level::Large);
}

extern "C" void LLVMRustSetModuleCodeModel(LLVMModuleRef M,
                                           LLVMRustCodeModel Model) {
  auto CM = fromRust(Model);
  if (!CM)
    return;
  unwrap(M)->setCodeModel(*CM);
}

extern "C" void LLVMRustSetModuleLargeDataThreshold(LLVMModuleRef M,
                                                    uint64_t Threshold) {
  unwrap(M)->setLargeDataThreshold(Threshold);
}

// Here you'll find an implementation of ThinLTO as used by the Rust compiler
// right now. This ThinLTO support is only enabled on "recent ish" versions of
// LLVM, and otherwise it's just blanket rejected from other compilers.
//
// Most of this implementation is straight copied from LLVM. At the time of
// this writing it wasn't *quite* suitable to reuse more code from upstream
// for our purposes, but we should strive to upstream this support once it's
// ready to go! I figure we may want a bit of testing locally first before
// sending this upstream to LLVM. I hear though they're quite eager to receive
// feedback like this!
//
// If you're reading this code and wondering "what in the world" or you're
// working "good lord by LLVM upgrade is *still* failing due to these bindings"
// then fear not! (ok maybe fear a little). All code here is mostly based
// on `lib/LTO/ThinLTOCodeGenerator.cpp` in LLVM.
//
// You'll find that the general layout here roughly corresponds to the `run`
// method in that file as well as `ProcessThinLTOModule`. Functions are
// specifically commented below as well, but if you're updating this code
// or otherwise trying to understand it, the LLVM source will be useful in
// interpreting the mysteries within.
//
// Otherwise I'll apologize in advance, it probably requires a relatively
// significant investment on your part to "truly understand" what's going on
// here. Not saying I do myself, but it took me a while staring at LLVM's source
// and various online resources about ThinLTO to make heads or tails of all
// this.

// This is a shared data structure which *must* be threadsafe to share
// read-only amongst threads. This also corresponds basically to the arguments
// of the `ProcessThinLTOModule` function in the LLVM source.
struct LLVMRustThinLTOData {
  // The combined index that is the global analysis over all modules we're
  // performing ThinLTO for. This is mostly managed by LLVM.
  ModuleSummaryIndex Index;

  // All modules we may look at, stored as in-memory serialized versions. This
  // is later used when inlining to ensure we can extract any module to inline
  // from.
  StringMap<MemoryBufferRef> ModuleMap;

  // A set that we manage of everything we *don't* want internalized. Note that
  // this includes all transitive references right now as well, but it may not
  // always!
  DenseSet<GlobalValue::GUID> GUIDPreservedSymbols;

  // Not 100% sure what these are, but they impact what's internalized and
  // what's inlined across modules, I believe.
  FunctionImporter::ImportListsTy ImportLists;
  DenseMap<StringRef, FunctionImporter::ExportSetTy> ExportLists;
  DenseMap<StringRef, GVSummaryMapTy> ModuleToDefinedGVSummaries;
  StringMap<std::map<GlobalValue::GUID, GlobalValue::LinkageTypes>> ResolvedODR;

  LLVMRustThinLTOData() : Index(/* HaveGVs = */ false) {}
};

// Just an argument to the `LLVMRustCreateThinLTOData` function below.
struct LLVMRustThinLTOModule {
  const char *identifier;
  const char *data;
  size_t len;
};

// This is copied from `lib/LTO/ThinLTOCodeGenerator.cpp`, not sure what it
// does.
static const GlobalValueSummary *getFirstDefinitionForLinker(
    ArrayRef<std::unique_ptr<GlobalValueSummary>> GVSummaryList) {
  auto StrongDefForLinker = llvm::find_if(
      GVSummaryList, [](const std::unique_ptr<GlobalValueSummary> &Summary) {
        auto Linkage = Summary->linkage();
        return !GlobalValue::isAvailableExternallyLinkage(Linkage) &&
               !GlobalValue::isWeakForLinker(Linkage);
      });
  if (StrongDefForLinker != GVSummaryList.end())
    return StrongDefForLinker->get();

  auto FirstDefForLinker = llvm::find_if(
      GVSummaryList, [](const std::unique_ptr<GlobalValueSummary> &Summary) {
        auto Linkage = Summary->linkage();
        return !GlobalValue::isAvailableExternallyLinkage(Linkage);
      });
  if (FirstDefForLinker == GVSummaryList.end())
    return nullptr;
  return FirstDefForLinker->get();
}

// The main entry point for creating the global ThinLTO analysis. The structure
// here is basically the same as before threads are spawned in the `run`
// function of `lib/LTO/ThinLTOCodeGenerator.cpp`.
extern "C" LLVMRustThinLTOData *
LLVMRustCreateThinLTOData(LLVMRustThinLTOModule *modules, size_t num_modules,
                          const char **preserved_symbols, size_t num_symbols) {
  auto Ret = std::make_unique<LLVMRustThinLTOData>();

  // Load each module's summary and merge it into one combined index
  for (size_t i = 0; i < num_modules; i++) {
    auto module = &modules[i];
    auto buffer = StringRef(module->data, module->len);
    auto mem_buffer = MemoryBufferRef(buffer, module->identifier);

    Ret->ModuleMap[module->identifier] = mem_buffer;

    if (Error Err = readModuleSummaryIndex(mem_buffer, Ret->Index)) {
      LLVMRustSetLastError(toString(std::move(Err)).c_str());
      return nullptr;
    }
  }

  // Collect for each module the list of function it defines (GUID -> Summary)
  Ret->Index.collectDefinedGVSummariesPerModule(
      Ret->ModuleToDefinedGVSummaries);

  // Convert the preserved symbols set from string to GUID, this is then needed
  // for internalization.
  for (size_t i = 0; i < num_symbols; i++) {
    auto GUID =
        GlobalValue::getGUIDAssumingExternalLinkage(preserved_symbols[i]);
    Ret->GUIDPreservedSymbols.insert(GUID);
  }

  // Collect the import/export lists for all modules from the call-graph in the
  // combined index
  //
  // This is copied from `lib/LTO/ThinLTOCodeGenerator.cpp`
  auto deadIsPrevailing = [&](GlobalValue::GUID G) {
    return PrevailingType::Unknown;
  };
  // We don't have a complete picture in our use of ThinLTO, just our immediate
  // crate, so we need `ImportEnabled = false` to limit internalization.
  // Otherwise, we sometimes lose `static` values -- see #60184.
  computeDeadSymbolsWithConstProp(Ret->Index, Ret->GUIDPreservedSymbols,
                                  deadIsPrevailing,
                                  /* ImportEnabled = */ false);
  // Resolve LinkOnce/Weak symbols, this has to be computed early be cause it
  // impacts the caching.
  //
  // This is copied from `lib/LTO/ThinLTOCodeGenerator.cpp` with some of this
  // being lifted from `lib/LTO/LTO.cpp` as well
  DenseMap<GlobalValue::GUID, const GlobalValueSummary *> PrevailingCopy;
  for (auto &I : Ret->Index) {
#if LLVM_VERSION_GE(22, 0)
    const auto &SummaryList = I.second.getSummaryList();
#else
    const auto &SummaryList = I.second.SummaryList;
#endif
    if (SummaryList.size() > 1)
      PrevailingCopy[I.first] = getFirstDefinitionForLinker(SummaryList);
  }
  auto isPrevailing = [&](GlobalValue::GUID GUID, const GlobalValueSummary *S) {
    const auto &Prevailing = PrevailingCopy.find(GUID);
    if (Prevailing == PrevailingCopy.end())
      return true;
    return Prevailing->second == S;
  };
  ComputeCrossModuleImport(Ret->Index, Ret->ModuleToDefinedGVSummaries,
                           isPrevailing, Ret->ImportLists, Ret->ExportLists);

  auto recordNewLinkage = [&](StringRef ModuleIdentifier,
                              GlobalValue::GUID GUID,
                              GlobalValue::LinkageTypes NewLinkage) {
    Ret->ResolvedODR[ModuleIdentifier][GUID] = NewLinkage;
  };

  // Uses FromPrevailing visibility scheme which works for many binary
  // formats. We probably could and should use ELF visibility scheme for many of
  // our targets, however.
  lto::Config conf;
  thinLTOResolvePrevailingInIndex(conf, Ret->Index, isPrevailing,
                                  recordNewLinkage, Ret->GUIDPreservedSymbols);

  // Here we calculate an `ExportedGUIDs` set for use in the `isExported`
  // callback below. This callback below will dictate the linkage for all
  // summaries in the index, and we basically just only want to ensure that dead
  // symbols are internalized. Otherwise everything that's already external
  // linkage will stay as external, and internal will stay as internal.
  std::set<GlobalValue::GUID> ExportedGUIDs;
  for (auto &List : Ret->Index) {
#if LLVM_VERSION_GE(22, 0)
    const auto &SummaryList = List.second.getSummaryList();
#else
    const auto &SummaryList = List.second.SummaryList;
#endif
    for (auto &GVS : SummaryList) {
      if (GlobalValue::isLocalLinkage(GVS->linkage()))
        continue;
      auto GUID = GVS->getOriginalName();
      if (GVS->flags().Live)
        ExportedGUIDs.insert(GUID);
    }
  }
  auto isExported = [&](StringRef ModuleIdentifier, ValueInfo VI) {
    const auto &ExportList = Ret->ExportLists.find(ModuleIdentifier);
    return (ExportList != Ret->ExportLists.end() &&
            ExportList->second.count(VI)) ||
           ExportedGUIDs.count(VI.getGUID());
  };
  thinLTOInternalizeAndPromoteInIndex(Ret->Index, isExported, isPrevailing);

  return Ret.release();
}

extern "C" void LLVMRustFreeThinLTOData(LLVMRustThinLTOData *Data) {
  delete Data;
}

// Below are the various passes that happen *per module* when doing ThinLTO.
//
// In other words, these are the functions that are all run concurrently
// with one another, one per module. The passes here correspond to the analysis
// passes in `lib/LTO/ThinLTOCodeGenerator.cpp`, currently found in the
// `ProcessThinLTOModule` function. Here they're split up into separate steps
// so rustc can save off the intermediate bytecode between each step.

static bool clearDSOLocalOnDeclarations(Module &Mod, TargetMachine &TM) {
  // When linking an ELF shared object, dso_local should be dropped. We
  // conservatively do this for -fpic.
  bool ClearDSOLocalOnDeclarations = TM.getTargetTriple().isOSBinFormatELF() &&
                                     TM.getRelocationModel() != Reloc::Static &&
                                     Mod.getPIELevel() == PIELevel::Default;
  return ClearDSOLocalOnDeclarations;
}

extern "C" void LLVMRustPrepareThinLTORename(const LLVMRustThinLTOData *Data,
                                             LLVMModuleRef M,
                                             LLVMTargetMachineRef TM) {
  Module &Mod = *unwrap(M);
  TargetMachine &Target = *unwrap(TM);

  bool ClearDSOLocal = clearDSOLocalOnDeclarations(Mod, Target);
  renameModuleForThinLTO(Mod, Data->Index, ClearDSOLocal);
}

extern "C" bool
LLVMRustPrepareThinLTOResolveWeak(const LLVMRustThinLTOData *Data,
                                  LLVMModuleRef M) {
  Module &Mod = *unwrap(M);
  const auto &DefinedGlobals =
      Data->ModuleToDefinedGVSummaries.lookup(Mod.getModuleIdentifier());
  thinLTOFinalizeInModule(Mod, DefinedGlobals, /*PropagateAttrs=*/true);
  return true;
}

extern "C" bool
LLVMRustPrepareThinLTOInternalize(const LLVMRustThinLTOData *Data,
                                  LLVMModuleRef M) {
  Module &Mod = *unwrap(M);
  const auto &DefinedGlobals =
      Data->ModuleToDefinedGVSummaries.lookup(Mod.getModuleIdentifier());
  thinLTOInternalizeModule(Mod, DefinedGlobals);
  return true;
}

extern "C" bool LLVMRustPrepareThinLTOImport(const LLVMRustThinLTOData *Data,
                                             LLVMModuleRef M,
                                             LLVMTargetMachineRef TM) {
  Module &Mod = *unwrap(M);
  TargetMachine &Target = *unwrap(TM);

  const auto &ImportList = Data->ImportLists.lookup(Mod.getModuleIdentifier());
  auto Loader = [&](StringRef Identifier) {
    const auto &Memory = Data->ModuleMap.lookup(Identifier);
    auto &Context = Mod.getContext();
    auto MOrErr = getLazyBitcodeModule(Memory, Context, true, true);

    if (!MOrErr)
      return MOrErr;

    // The rest of this closure is a workaround for
    // https://bugs.llvm.org/show_bug.cgi?id=38184 where during ThinLTO imports
    // we accidentally import wasm custom sections into different modules,
    // duplicating them by in the final output artifact.
    //
    // The issue is worked around here by manually removing the
    // `wasm.custom_sections` named metadata node from any imported module. This
    // we know isn't used by any optimization pass so there's no need for it to
    // be imported.
    //
    // Note that the metadata is currently lazily loaded, so we materialize it
    // here before looking up if there's metadata inside. The `FunctionImporter`
    // will immediately materialize metadata anyway after an import, so this
    // shouldn't be a perf hit.
    if (Error Err = (*MOrErr)->materializeMetadata()) {
      Expected<std::unique_ptr<Module>> Ret(std::move(Err));
      return Ret;
    }

    auto *WasmCustomSections =
        (*MOrErr)->getNamedMetadata("wasm.custom_sections");
    if (WasmCustomSections)
      WasmCustomSections->eraseFromParent();

    // `llvm.ident` named metadata also gets duplicated.
    auto *llvmIdent = (*MOrErr)->getNamedMetadata("llvm.ident");
    if (llvmIdent)
      llvmIdent->eraseFromParent();

    return MOrErr;
  };
  bool ClearDSOLocal = clearDSOLocalOnDeclarations(Mod, Target);
  auto Importer = FunctionImporter(Data->Index, Loader, ClearDSOLocal);
  Expected<bool> Result = Importer.importFunctions(Mod, ImportList);
  if (!Result) {
    LLVMRustSetLastError(toString(Result.takeError()).c_str());
    return false;
  }
  return true;
}

extern "C" LLVMRustBuffer *LLVMRustModuleSerialize(LLVMModuleRef M,
                                                   bool is_thin) {
  auto Ret = std::make_unique<LLVMRustBuffer>();
  {
    auto OS = raw_string_ostream(Ret->data);
    {
      PassBuilder PB;
      LoopAnalysisManager LAM;
      FunctionAnalysisManager FAM;
      CGSCCAnalysisManager CGAM;
      ModuleAnalysisManager MAM;
      PB.registerModuleAnalyses(MAM);
      PB.registerCGSCCAnalyses(CGAM);
      PB.registerFunctionAnalyses(FAM);
      PB.registerLoopAnalyses(LAM);
      PB.crossRegisterProxies(LAM, FAM, CGAM, MAM);

      ModulePassManager MPM;
#if LLVM_VERSION_GE(23, 0)
      MPM.addPass(AssignGUIDPass());
#endif

      if (is_thin) {
        MPM.addPass(ThinLTOBitcodeWriterPass(OS, nullptr));
      } else {
        forceFullLTOSummary(unwrap(M));
        MPM.addPass(BitcodeWriterPass(OS, /*ShouldPreserveUseListOrder=*/false,
                                      /*EmitSummaryIndex=*/true));
      }
      MPM.run(*unwrap(M), MAM);
    }
  }
  return Ret.release();
}

// This is what we used to parse upstream bitcode for actual ThinLTO
// processing. We'll call this once per module optimized through ThinLTO, and
// it'll be called concurrently on many threads.
extern "C" LLVMModuleRef LLVMRustParseBitcodeForLTO(LLVMContextRef Context,
                                                    const char *data,
                                                    size_t len,
                                                    const char *identifier) {
  auto Data = StringRef(data, len);
  auto Buffer = MemoryBufferRef(Data, identifier);
  unwrap(Context)->enableDebugTypeODRUniquing();
  Expected<std::unique_ptr<Module>> SrcOrError =
      parseBitcodeFile(Buffer, *unwrap(Context));
  if (!SrcOrError) {
    LLVMRustSetLastError(toString(SrcOrError.takeError()).c_str());
    return nullptr;
  }
  return wrap(std::move(*SrcOrError).release());
}

// Computes the LTO cache key for the provided 'ModId' in the given 'Data',
// storing the result in 'KeyOut'.
// Currently, this cache key is a SHA-1 hash of anything that could affect
// the result of optimizing this module (e.g. module imports, exports, liveness
// of access globals, etc).
// The precise details are determined by LLVM in `computeLTOCacheKey`, which is
// used during the normal linker-plugin incremental thin-LTO process.
extern "C" void LLVMRustComputeLTOCacheKey(RustStringRef KeyOut,
                                           const char *ModId,
                                           LLVMRustThinLTOData *Data) {
  SmallString<40> Key;
  llvm::lto::Config conf;
  const auto &ImportList = Data->ImportLists.lookup(ModId);
  const auto &ExportList = Data->ExportLists.lookup(ModId);
  const auto &ResolvedODR = Data->ResolvedODR.lookup(ModId);
  const auto &DefinedGlobals = Data->ModuleToDefinedGVSummaries.lookup(ModId);
  DenseSet<GlobalValue::GUID> CfiFunctionDefs;
  DenseSet<GlobalValue::GUID> CfiFunctionDecls;

  // Based on the 'InProcessThinBackend' constructor in LLVM
#if LLVM_VERSION_GE(23, 0)
  CfiFunctionDefs.insert_range(
      Data->Index.cfiFunctionDefs().getExportedThinLTOGUIDs());
  CfiFunctionDecls.insert_range(
      Data->Index.cfiFunctionDecls().getExportedThinLTOGUIDs());
#else
  for (auto &Name : Data->Index.cfiFunctionDefs().symbols())
    CfiFunctionDefs.insert(GlobalValue::getGUIDAssumingExternalLinkage(
        GlobalValue::dropLLVMManglingEscape(Name)));
  for (auto &Name : Data->Index.cfiFunctionDecls().symbols())
    CfiFunctionDecls.insert(GlobalValue::getGUIDAssumingExternalLinkage(
        GlobalValue::dropLLVMManglingEscape(Name)));
#endif

  Key = llvm::computeLTOCacheKey(conf, Data->Index, ModId, ImportList,
                                 ExportList, ResolvedODR, DefinedGlobals,
                                 CfiFunctionDefs, CfiFunctionDecls);

  auto OS = RawRustStringOstream(KeyOut);
  OS << Key.str();
}