//===-- lib/CodeGen/MachineInstrBundle.cpp --------------------------------===//
//
// 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/CodeGen/MachineInstrBundle.h"
#include "llvm/ADT/SetVector.h"
#include "llvm/ADT/SmallSet.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/CodeGen/MachineFunction.h"
#include "llvm/CodeGen/MachineFunctionAnalysisManager.h"
#include "llvm/CodeGen/MachineFunctionPass.h"
#include "llvm/CodeGen/MachineInstrBuilder.h"
#include "llvm/CodeGen/Passes.h"
#include "llvm/CodeGen/TargetInstrInfo.h"
#include "llvm/CodeGen/TargetRegisterInfo.h"
#include "llvm/CodeGen/TargetSubtargetInfo.h"
#include "llvm/InitializePasses.h"
#include "llvm/Pass.h"
#include <utility>
using namespace llvm;

static bool unpackBundles(MachineFunction &MF,
                          std::function<bool(const MachineFunction &)> Ftor) {
  if (Ftor && !Ftor(MF))
    return false;

  bool Changed = false;
  for (MachineBasicBlock &MBB : MF) {
    for (MachineBasicBlock::instr_iterator MII = MBB.instr_begin(),
           MIE = MBB.instr_end(); MII != MIE; ) {
      MachineInstr *MI = &*MII;

      // Remove BUNDLE instruction and the InsideBundle flags from bundled
      // instructions.
      if (MI->isBundle()) {
        while (++MII != MIE && MII->isBundledWithPred()) {
          MII->unbundleFromPred();
          for (MachineOperand &MO  : MII->operands()) {
            if (MO.isReg() && MO.isInternalRead())
              MO.setIsInternalRead(false);
          }
        }
        MI->eraseFromParent();

        Changed = true;
        continue;
      }

      ++MII;
    }
  }

  return Changed;
}

namespace {

class UnpackMachineBundlesLegacy : public MachineFunctionPass {
public:
  static char ID; // Pass identification
  UnpackMachineBundlesLegacy(
      std::function<bool(const MachineFunction &)> Ftor = nullptr)
      : MachineFunctionPass(ID), PredicateFtor(std::move(Ftor)) {}

  bool runOnMachineFunction(MachineFunction &MF) override;

private:
  std::function<bool(const MachineFunction &)> PredicateFtor;
};
} // end anonymous namespace

PreservedAnalyses
UnpackMachineBundlesPass::run(MachineFunction &MF,
                              MachineFunctionAnalysisManager &MFAM) {
  if (unpackBundles(MF, PredicateFtor))
    return PreservedAnalyses::none();
  return PreservedAnalyses::all();
}

char UnpackMachineBundlesLegacy::ID = 0;
char &llvm::UnpackMachineBundlesID = UnpackMachineBundlesLegacy::ID;
INITIALIZE_PASS(UnpackMachineBundlesLegacy, "unpack-mi-bundles",
                "Unpack machine instruction bundles", false, false)

bool UnpackMachineBundlesLegacy::runOnMachineFunction(MachineFunction &MF) {
  return unpackBundles(MF, PredicateFtor);
}

FunctionPass *llvm::createUnpackMachineBundlesLegacy(
    std::function<bool(const MachineFunction &)> Ftor) {
  return new UnpackMachineBundlesLegacy(std::move(Ftor));
}

/// Return the first DebugLoc that has line number information, given a
/// range of instructions. The search range is from FirstMI to LastMI
/// (exclusive). Otherwise return the first DILocation or an empty location if
/// there are none.
static DebugLoc getDebugLoc(MachineBasicBlock::instr_iterator FirstMI,
                            MachineBasicBlock::instr_iterator LastMI) {
  DebugLoc DL;
  for (auto MII = FirstMI; MII != LastMI; ++MII) {
    if (DebugLoc MIIDL = MII->getDebugLoc()) {
      if (MIIDL.getLine() != 0)
        return MIIDL;
      DL = MIIDL.get();
    }
  }
  return DL;
}

/// Check if target reg is contained in given lists, which are:
/// LocalDefsV as given list for virtual regs
/// LocalDefsP as given list for physical regs, in BitVector[RegUnit] form
static bool containsReg(SmallSetVector<Register, 32> LocalDefsV,
                        const BitVector &LocalDefsP, Register Reg,
                        const TargetRegisterInfo *TRI) {
  if (Reg.isPhysical()) {
    for (MCRegUnit Unit : TRI->regunits(Reg.asMCReg()))
      if (!LocalDefsP[static_cast<unsigned>(Unit)])
        return false;

    return true;
  }
  return LocalDefsV.contains(Reg);
}

/// finalizeBundle - Finalize a machine instruction bundle which includes
/// a sequence of instructions starting from FirstMI to LastMI (exclusive).
/// This routine adds a BUNDLE instruction to represent the bundle, it adds
/// IsInternalRead markers to MachineOperands which are defined inside the
/// bundle, and it copies externally visible defs and uses to the BUNDLE
/// instruction.
void llvm::finalizeBundle(MachineBasicBlock &MBB,
                          MachineBasicBlock::instr_iterator FirstMI,
                          MachineBasicBlock::instr_iterator LastMI) {
  assert(FirstMI != LastMI && "Empty bundle?");
  MIBundleBuilder Bundle(MBB, FirstMI, LastMI);

  MachineFunction &MF = *MBB.getParent();
  const TargetInstrInfo *TII = MF.getSubtarget().getInstrInfo();
  const TargetRegisterInfo *TRI = MF.getSubtarget().getRegisterInfo();

  MachineInstrBuilder MIB =
      BuildMI(MF, getDebugLoc(FirstMI, LastMI), TII->get(TargetOpcode::BUNDLE));
  Bundle.prepend(MIB);

  SmallSetVector<Register, 32> LocalDefs;
  BitVector LocalDefsP(TRI->getNumRegUnits());
  SmallSet<Register, 8> DeadDefSet;
  SmallSetVector<Register, 8> ExternUses;
  SmallSet<Register, 8> KilledUseSet;
  SmallSet<Register, 8> UndefUseSet;
  SmallVector<std::pair<Register, Register>> TiedOperands;
  SmallVector<MachineInstr *> MemMIs;
  for (auto MII = FirstMI; MII != LastMI; ++MII) {
    // Debug instructions have no effects to track.
    if (MII->isDebugInstr())
      continue;

    for (MachineOperand &MO : MII->all_uses()) {
      Register Reg = MO.getReg();
      if (!Reg)
        continue;

      if (containsReg(LocalDefs, LocalDefsP, Reg, TRI)) {
        MO.setIsInternalRead();
        if (MO.isKill()) {
          // Internal def is now killed.
          DeadDefSet.insert(Reg);
        }
      } else {
        if (ExternUses.insert(Reg)) {
          if (MO.isUndef())
            UndefUseSet.insert(Reg);
        }
        if (MO.isKill()) {
          // External def is now killed.
          KilledUseSet.insert(Reg);
        }
        if (MO.isTied() && Reg.isVirtual()) {
          // Record tied operand constraints that involve virtual registers so
          // that bundles that are formed pre-register allocation reflect the
          // relevant constraints.
          unsigned TiedIdx = MII->findTiedOperandIdx(MO.getOperandNo());
          MachineOperand &TiedMO = MII->getOperand(TiedIdx);
          Register DefReg = TiedMO.getReg();
          TiedOperands.emplace_back(DefReg, Reg);
        }
      }
    }

    for (MachineOperand &MO : MII->all_defs()) {
      Register Reg = MO.getReg();
      if (!Reg)
        continue;

      if (LocalDefs.insert(Reg)) {
        if (!MO.isDead() && Reg.isPhysical()) {
          for (MCRegUnit Unit : TRI->regunits(Reg.asMCReg()))
            LocalDefsP.set(static_cast<unsigned>(Unit));
        }
      } else {
        if (!MO.isDead()) {
          // Re-defined inside the bundle, it's no longer dead.
          DeadDefSet.erase(Reg);
        }
      }
      if (MO.isDead())
        DeadDefSet.insert(Reg);
    }

    // Set FrameSetup/FrameDestroy for the bundle. If any of the instructions
    // got the property, then also set it on the bundle.
    if (MII->getFlag(MachineInstr::FrameSetup))
      MIB.setMIFlag(MachineInstr::FrameSetup);
    if (MII->getFlag(MachineInstr::FrameDestroy))
      MIB.setMIFlag(MachineInstr::FrameDestroy);

    if (MII->mayLoadOrStore())
      MemMIs.push_back(&*MII);
  }

  for (Register Reg : LocalDefs) {
    // If it's not live beyond end of the bundle, mark it dead.
    bool isDead = DeadDefSet.contains(Reg);
    MIB.addReg(Reg, getDefRegState(true) | getDeadRegState(isDead) |
                        getImplRegState(true));
  }

  for (Register Reg : ExternUses) {
    bool isKill = KilledUseSet.contains(Reg);
    bool isUndef = UndefUseSet.contains(Reg);
    MIB.addReg(Reg, getKillRegState(isKill) | getUndefRegState(isUndef) |
                        getImplRegState(true));
  }

  for (auto [DefReg, UseReg] : TiedOperands) {
    unsigned DefIdx =
        std::distance(LocalDefs.begin(), llvm::find(LocalDefs, DefReg));
    unsigned UseIdx =
        std::distance(ExternUses.begin(), llvm::find(ExternUses, UseReg));
    assert(DefIdx < LocalDefs.size());
    assert(UseIdx < ExternUses.size());
    MIB->tieOperands(DefIdx, LocalDefs.size() + UseIdx);
  }

  MIB->cloneMergedMemRefs(MF, MemMIs);
}

/// finalizeBundle - Same functionality as the previous finalizeBundle except
/// the last instruction in the bundle is not provided as an input. This is
/// used in cases where bundles are pre-determined by marking instructions
/// with 'InsideBundle' marker. It returns the MBB instruction iterator that
/// points to the end of the bundle.
MachineBasicBlock::instr_iterator
llvm::finalizeBundle(MachineBasicBlock &MBB,
                     MachineBasicBlock::instr_iterator FirstMI) {
  MachineBasicBlock::instr_iterator E = MBB.instr_end();
  MachineBasicBlock::instr_iterator LastMI = std::next(FirstMI);
  while (LastMI != E && LastMI->isInsideBundle())
    ++LastMI;
  finalizeBundle(MBB, FirstMI, LastMI);
  return LastMI;
}

/// finalizeBundles - Finalize instruction bundles in the specified
/// MachineFunction. Return true if any bundles are finalized.
bool llvm::finalizeBundles(MachineFunction &MF) {
  bool Changed = false;
  for (MachineBasicBlock &MBB : MF) {
    MachineBasicBlock::instr_iterator MII = MBB.instr_begin();
    MachineBasicBlock::instr_iterator MIE = MBB.instr_end();
    if (MII == MIE)
      continue;
    assert(!MII->isInsideBundle() &&
           "First instr cannot be inside bundle before finalization!");

    for (++MII; MII != MIE; ) {
      if (!MII->isInsideBundle())
        ++MII;
      else {
        MII = finalizeBundle(MBB, std::prev(MII));
        Changed = true;
      }
    }
  }

  return Changed;
}

VirtRegInfo llvm::AnalyzeVirtRegInBundle(
    MachineInstr &MI, Register Reg,
    SmallVectorImpl<std::pair<MachineInstr *, unsigned>> *Ops) {
  VirtRegInfo RI = {false, false, false};
  for (MIBundleOperands O(MI); O.isValid(); ++O) {
    MachineOperand &MO = *O;
    if (!MO.isReg() || MO.getReg() != Reg)
      continue;

    // Remember each (MI, OpNo) that refers to Reg.
    if (Ops)
      Ops->push_back(std::make_pair(MO.getParent(), O.getOperandNo()));

    // Both defs and uses can read virtual registers.
    if (MO.readsReg()) {
      RI.Reads = true;
      if (MO.isDef())
        RI.Tied = true;
    }

    // Only defs can write.
    if (MO.isDef())
      RI.Writes = true;
    else if (!RI.Tied &&
             MO.getParent()->isRegTiedToDefOperand(O.getOperandNo()))
      RI.Tied = true;
  }
  return RI;
}

std::pair<LaneBitmask, LaneBitmask>
llvm::AnalyzeVirtRegLanesInBundle(const MachineInstr &MI, Register Reg,
                                  const MachineRegisterInfo &MRI,
                                  const TargetRegisterInfo &TRI) {

  LaneBitmask UseMask, DefMask;

  for (const MachineOperand &MO : const_mi_bundle_ops(MI)) {
    if (!MO.isReg() || MO.getReg() != Reg)
      continue;

    unsigned SubReg = MO.getSubReg();
    if (SubReg == 0 && MO.isUse() && !MO.isUndef())
      UseMask |= MRI.getMaxLaneMaskForVReg(Reg);

    LaneBitmask SubRegMask = TRI.getSubRegIndexLaneMask(SubReg);
    if (MO.isDef()) {
      if (!MO.isUndef())
        UseMask |= ~SubRegMask;
      DefMask |= SubRegMask;
    } else if (!MO.isUndef())
      UseMask |= SubRegMask;
  }

  return {UseMask, DefMask};
}

PhysRegInfo llvm::AnalyzePhysRegInBundle(const MachineInstr &MI, Register Reg,
                                         const TargetRegisterInfo *TRI) {
  bool AllDefsDead = true;
  PhysRegInfo PRI = {false, false, false, false, false, false, false, false};

  assert(Reg.isPhysical() && "analyzePhysReg not given a physical register!");
  for (const MachineOperand &MO : const_mi_bundle_ops(MI)) {
    if (MO.isRegMask() && MO.clobbersPhysReg(Reg)) {
      PRI.Clobbered = true;
      continue;
    }

    if (!MO.isReg())
      continue;

    Register MOReg = MO.getReg();
    if (!MOReg || !MOReg.isPhysical())
      continue;

    if (!TRI->regsOverlap(MOReg, Reg))
      continue;

    bool Covered = TRI->isSuperRegisterEq(Reg, MOReg);
    if (MO.readsReg()) {
      PRI.Read = true;
      if (Covered) {
        PRI.FullyRead = true;
        if (MO.isKill())
          PRI.Killed = true;
      }
    } else if (MO.isDef()) {
      PRI.Defined = true;
      if (Covered)
        PRI.FullyDefined = true;
      if (!MO.isDead())
        AllDefsDead = false;
    }
  }

  if (AllDefsDead) {
    if (PRI.FullyDefined || PRI.Clobbered)
      PRI.DeadDef = true;
    else if (PRI.Defined)
      PRI.PartialDeadDef = true;
  }

  return PRI;
}

PreservedAnalyses
llvm::FinalizeBundleTestPass::run(MachineFunction &MF,
                                  MachineFunctionAnalysisManager &) {
  // For testing purposes, bundle the entire contents of each basic block
  // except for terminators.
  for (MachineBasicBlock &MBB : MF)
    finalizeBundle(MBB, MBB.instr_begin(), MBB.getFirstInstrTerminator());
  return getMachineFunctionPassPreservedAnalyses();
}