//===----------------------------------------------------------------------===//
//
// 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 "WebAssembly.h"
#include "WebAssemblyAsmPrinter.h"
#include "WebAssemblyExceptionInfo.h"
#include "WebAssemblyTargetMachine.h"
#include "llvm/CodeGen/AtomicExpand.h"
#include "llvm/CodeGen/FuncletLayout.h"
#include "llvm/CodeGen/GlobalISel/IRTranslator.h"
#include "llvm/CodeGen/GlobalISel/InstructionSelect.h"
#include "llvm/CodeGen/GlobalISel/Legalizer.h"
#include "llvm/CodeGen/GlobalISel/RegBankSelect.h"
#include "llvm/CodeGen/IndirectBrExpand.h"
#include "llvm/CodeGen/MachineBlockPlacement.h"
#include "llvm/CodeGen/MachineCopyPropagation.h"
#include "llvm/CodeGen/MachineLateInstrsCleanup.h"
#include "llvm/CodeGen/PatchableFunction.h"
#include "llvm/CodeGen/PostRAMachineSink.h"
#include "llvm/CodeGen/PostRASchedulerList.h"
#include "llvm/CodeGen/RegisterCoalescerPass.h"
#include "llvm/CodeGen/RemoveLoadsIntoFakeUses.h"
#include "llvm/CodeGen/ShrinkWrap.h"
#include "llvm/CodeGen/UnreachableBlockElim.h"
#include "llvm/IR/PassInstrumentation.h"
#include "llvm/MC/MCStreamer.h"
#include "llvm/Passes/CodeGenPassBuilder.h"
#include "llvm/Passes/PassBuilder.h"
#include "llvm/Support/CodeGen.h"
#include "llvm/Support/Error.h"
#include "llvm/Target/CGPassBuilderOption.h"
#include "llvm/Transforms/Utils/LowerGlobalDtors.h"

using namespace llvm;

namespace WebAssembly {
extern cl::opt<bool> WasmDisableExplicitLocals;
} // namespace WebAssembly

using llvm::WebAssembly::WasmDisableExplicitLocals;

namespace {

class WebAssemblyCodeGenPassBuilder : public CodeGenPassBuilder {
  using Base = CodeGenPassBuilder;

  WebAssemblyTargetMachine &getTM() const {
    return static_cast<WebAssemblyTargetMachine &>(TM);
  }

public:
  explicit WebAssemblyCodeGenPassBuilder(WebAssemblyTargetMachine &TM,
                                         const CGPassBuilderOption &Opts,
                                         PassInstrumentationCallbacks *PIC)
      : CodeGenPassBuilder(TM, Opts, PIC) {
    disablePass<MachineLateInstrsCleanupPass, MachineCopyPropagationPass,
                PostRAMachineSinkingPass, PostRASchedulerPass,
                FuncletLayoutPass, StackMapLivenessPass, PatchableFunctionPass,
                ShrinkWrapPass, RemoveLoadsIntoFakeUsesPass,
                MachineBlockPlacementPass>();

    // Currently RegisterCoalesce degrades wasm debug info quality by a
    // significant margin. As a quick fix, disable this for -O1, which is often
    // used for debugging large applications. Disabling this increases code size
    // of Emscripten core benchmarks by ~5%, which is acceptable for -O1, which
    // is usually not used for production builds.
    // TODO Investigate why RegisterCoalesce degrades debug info quality and fix
    // it properly
    if (getOptLevel() == CodeGenOptLevel::Less)
      disablePass<RegisterCoalescerPass>();
  }

  void addIRPasses(PassManagerWrapper &PMW) override;
  void addISelPrepare(PassManagerWrapper &PMW) override;

  Error addInstSelector(PassManagerWrapper &PMW) override;

  Error addIRTranslator(PassManagerWrapper &PMW) override;
  void addPreLegalizeMachineIR(PassManagerWrapper &PMW) override;
  Error addLegalizeMachineIR(PassManagerWrapper &PMW) override;
  void addPreRegBankSelect(PassManagerWrapper &PMW) override;
  Error addRegBankSelect(PassManagerWrapper &PMW) override;
  Error addGlobalInstructionSelect(PassManagerWrapper &PMW) override;

  Error addRegAssignAndRewriteFast(PassManagerWrapper &PMW) override;
  Expected<bool>
  addRegAssignAndRewriteOptimized(PassManagerWrapper &PMW) override;
  void addPreEmitPass(PassManagerWrapper &PMW) override;
  void addAsmPrinterBegin(PassManagerWrapper &PMW) override;
  void addAsmPrinter(PassManagerWrapper &PMW) override;
  void addAsmPrinterEnd(PassManagerWrapper &PMW) override;
};

void WebAssemblyCodeGenPassBuilder::addIRPasses(PassManagerWrapper &PMW) {
  // Add signatures to prototype-less function declarations
  flushFPMsToMPM(PMW);
  addModulePass(WebAssemblyAddMissingPrototypesPass(), PMW);

  // Lower .llvm.global_dtors into .llvm.global_ctors with __cxa_atexit calls.
  addModulePass(LowerGlobalDtorsPass(), PMW);

  // Fix function bitcasts, as WebAssembly requires caller and callee signatures
  // to match.
  addModulePass(WebAssemblyFixFunctionBitcastsPass(), PMW);

  // Optimize "returned" function attributes.
  if (getOptLevel() != CodeGenOptLevel::None)
    addFunctionPass(WebAssemblyOptimizeReturnedPass(), PMW);

  // Wasm SjLj shares Emscripten SjLj's transformation, so it is handled here.
  // Without Wasm or Emscripten EH, this also lowers invokes.
  flushFPMsToMPM(PMW);
  addModulePass(WebAssemblyLowerEmscriptenEHSjLjPass(), PMW);

  // Expand indirectbr instructions to switches.
  addFunctionPass(IndirectBrExpandPass(TM), PMW);

  // Try to expand `vecreduce_{and, or}` into `{any, all}_true`.
  addFunctionPass(WebAssemblyReduceToAnyAllTruePass(getTM()), PMW);

  Base::addIRPasses(PMW);
}

void WebAssemblyCodeGenPassBuilder::addISelPrepare(PassManagerWrapper &PMW) {
  // We need to move reference type allocas to WASM_ADDRESS_SPACE_VAR so that
  // loads and stores are promoted to local.gets/local.sets.
  addFunctionPass(WebAssemblyRefTypeMem2LocalPass(), PMW);
  // Lower atomics and TLS if necessary
  flushFPMsToMPM(PMW);
  addModulePass(WebAssemblyCoalesceFeaturesAndStripAtomicsPass(getTM()), PMW);

  // This is a no-op if atomics are not used in the module
  addFunctionPass(AtomicExpandPass(TM), PMW);

  Base::addISelPrepare(PMW);
}

Error WebAssemblyCodeGenPassBuilder::addInstSelector(PassManagerWrapper &PMW) {
  addMachineFunctionPass(WebAssemblyISelDAGToDAGPass(getTM(), getOptLevel()),
                         PMW);

  // Run the argument-move pass immediately after the ScheduleDAG scheduler
  // so that we can fix up the ARGUMENT instructions before anything else
  // sees them in the wrong place.
  addMachineFunctionPass(WebAssemblyArgumentMovePass(), PMW);

  // Set the p2align operands. This information is present during ISel, however
  // it's inconvenient to collect. Collect it now, and update the immediate
  // operands.
  addMachineFunctionPass(WebAssemblySetP2AlignOperandsPass(), PMW);

  // Eliminate range checks and add default targets to br_table instructions.
  addMachineFunctionPass(WebAssemblyFixBrTableDefaultsPass(), PMW);

  // unreachable is terminator, non-terminator instruction after it is not
  // allowed.
  addMachineFunctionPass(WebAssemblyCleanCodeAfterTrapPass(), PMW);

  return Error::success();
}

Error WebAssemblyCodeGenPassBuilder::addIRTranslator(PassManagerWrapper &PMW) {
  addMachineFunctionPass(IRTranslatorPass(getOptLevel()), PMW);
  return Error::success();
}

void WebAssemblyCodeGenPassBuilder::addPreLegalizeMachineIR(
    PassManagerWrapper &PMW) {
  if (getOptLevel() != CodeGenOptLevel::None)
    addMachineFunctionPass(WebAssemblyPreLegalizerCombinerPass(), PMW);
}

Error WebAssemblyCodeGenPassBuilder::addLegalizeMachineIR(
    PassManagerWrapper &PMW) {
  addMachineFunctionPass(LegalizerPass(), PMW);
  return Error::success();
}

void WebAssemblyCodeGenPassBuilder::addPreRegBankSelect(
    PassManagerWrapper &PMW) {
  if (getOptLevel() != CodeGenOptLevel::None)
    addMachineFunctionPass(WebAssemblyPostLegalizerCombinerPass(), PMW);
}

Error WebAssemblyCodeGenPassBuilder::addRegBankSelect(PassManagerWrapper &PMW) {
  addMachineFunctionPass(RegBankSelectPass(), PMW);
  return Error::success();
}

Error WebAssemblyCodeGenPassBuilder::addGlobalInstructionSelect(
    PassManagerWrapper &PMW) {
  addMachineFunctionPass(InstructionSelectPass(getOptLevel()), PMW);

  if (isGlobalISelAbortEnabled()) {
    addMachineFunctionPass(WebAssemblyArgumentMovePass(), PMW);
    addMachineFunctionPass(WebAssemblySetP2AlignOperandsPass(), PMW);
    addMachineFunctionPass(WebAssemblyFixBrTableDefaultsPass(), PMW);
    addMachineFunctionPass(WebAssemblyCleanCodeAfterTrapPass(), PMW);
  }

  return Error::success();
}

Error WebAssemblyCodeGenPassBuilder::addRegAssignAndRewriteFast(
    PassManagerWrapper &PMW) {
  return Error::success();
}

Expected<bool> WebAssemblyCodeGenPassBuilder::addRegAssignAndRewriteOptimized(
    PassManagerWrapper &PMW) {
  return false;
}

void WebAssemblyCodeGenPassBuilder::addPreEmitPass(PassManagerWrapper &PMW) {
  Base::addPreEmitPass(PMW);

  // Nullify DBG_VALUE_LISTs that we cannot handle.
  addMachineFunctionPass(WebAssemblyNullifyDebugValueListsPass(), PMW);

  // Remove any unreachable blocks that may be left floating around.
  // Rare, but possible. Needed for WebAssemblyFixIrreducibleControlFlow.
  addMachineFunctionPass(UnreachableMachineBlockElimPass(), PMW);

  // Eliminate multiple-entry loops.
  addMachineFunctionPass(WebAssemblyFixIrreducibleControlFlowPass(), PMW);

  // Do various transformations for exception handling.
  // Every CFG-changing optimizations should come before this.
  addMachineFunctionPass(WebAssemblyLateEHPreparePass(), PMW);

  // Now that we have a prologue and epilogue and all frame indices are
  // rewritten, eliminate SP and FP. This allows them to be stackified,
  // colored, and numbered with the rest of the registers.
  addMachineFunctionPass(WebAssemblyReplacePhysRegsPass(), PMW);

  // Preparations and optimizations related to register stackification.
  if (getOptLevel() != CodeGenOptLevel::None) {
    // Depend on LiveIntervals and perform some optimizations on it.
    addMachineFunctionPass(WebAssemblyOptimizeLiveIntervalsPass(), PMW);

    // Prepare memory intrinsic calls for register stackifying.
    addMachineFunctionPass(WebAssemblyMemIntrinsicResultsPass(), PMW);
  }

  // Mark registers as representing wasm's value stack. This is a key
  // code-compression technique in WebAssembly. We run this pass (and
  // MemIntrinsicResults above) very late, so that it sees as much code as
  // possible, including code emitted by PEI and expanded by late tail
  // duplication.
  addMachineFunctionPass(WebAssemblyRegStackifyPass(getOptLevel()), PMW);

  if (getOptLevel() != CodeGenOptLevel::None) {
    // Run the register coloring pass to reduce the total number of registers.
    // This runs after stackification so that it doesn't consider registers
    // that become stackified.
    addMachineFunctionPass(WebAssemblyRegColoringPass(), PMW);
  }

  // Sort the blocks of the CFG into topological order, a prerequisite for
  // BLOCK and LOOP markers.
  addMachineFunctionPass(WebAssemblyCFGSortPass(), PMW);

  // Insert BLOCK and LOOP markers.
  addMachineFunctionPass(WebAssemblyCFGStackifyPass(), PMW);

  // Insert explicit local.get and local.set operators.
  if (!WasmDisableExplicitLocals)
    addMachineFunctionPass(WebAssemblyExplicitLocalsPass(), PMW);

  // Lower br_unless into br_if.
  addMachineFunctionPass(WebAssemblyLowerBrUnlessPass(), PMW);

  // Perform the very last peephole optimizations on the code.
  if (getOptLevel() != CodeGenOptLevel::None)
    addMachineFunctionPass(WebAssemblyPeepholePass(), PMW);

  // Create a mapping from LLVM CodeGen virtual registers to wasm registers.
  addMachineFunctionPass(WebAssemblyRegNumberingPass(), PMW);

  // Fix debug_values whose defs have been stackified.
  if (!WasmDisableExplicitLocals)
    addMachineFunctionPass(WebAssemblyDebugFixupPass(), PMW);

  // Collect information to prepare for MC lowering / asm printing.
  flushFPMsToMPM(PMW);
  addModulePass(WebAssemblyMCLowerPrePass(), PMW);
}

void WebAssemblyCodeGenPassBuilder::addAsmPrinterBegin(
    PassManagerWrapper &PMW) {
  addModulePass(WebAssemblyAsmPrinterBeginPass(), PMW, /*Force=*/true);
}

void WebAssemblyCodeGenPassBuilder::addAsmPrinter(PassManagerWrapper &PMW) {
  addMachineFunctionPass(WebAssemblyAsmPrinterPass(), PMW);
}

void WebAssemblyCodeGenPassBuilder::addAsmPrinterEnd(PassManagerWrapper &PMW) {
  addModulePass(WebAssemblyAsmPrinterEndPass(), PMW);
}

} // namespace

void WebAssemblyTargetMachine::registerPassBuilderCallbacks(PassBuilder &PB){
#define GET_PASS_REGISTRY "WebAssemblyPassRegistry.def"
#include "llvm/Passes/TargetPassRegistry.inc"
}

Error WebAssemblyTargetMachine::buildCodeGenPipeline(
    ModulePassManager &MPM, ModuleAnalysisManager &MAM, raw_pwrite_stream &Out,
    raw_pwrite_stream *DwoOut, CodeGenFileType FileType,
    const CGPassBuilderOption &Opt, MCContext &Ctx,
    PassInstrumentationCallbacks *PIC) {
  auto CGPB = WebAssemblyCodeGenPassBuilder(*this, Opt, PIC);
  return CGPB.buildPipeline(MPM, MAM, Out, DwoOut, FileType, Ctx);
}