//=====-- Rematerializer.cpp - MIR rematerialization support ----*- C++ -*-===//
//
// 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
//
//==-----------------------------------------------------------------------===//
//
/// \file
/// Implements helpers for target-independent rematerialization at the MIR
/// level.
//
//===----------------------------------------------------------------------===//

#include "llvm/CodeGen/Rematerializer.h"
#include "llvm/ADT/MapVector.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/ADT/SetVector.h"
#include "llvm/CodeGen/LiveIntervals.h"
#include "llvm/CodeGen/MachineBasicBlock.h"
#include "llvm/CodeGen/MachineOperand.h"
#include "llvm/CodeGen/MachineRegisterInfo.h"
#include "llvm/CodeGen/Register.h"
#include "llvm/CodeGen/TargetRegisterInfo.h"
#include "llvm/Support/Debug.h"
#include "llvm/Support/ErrorHandling.h"
#include <optional>

#define DEBUG_TYPE "rematerializer"

using namespace llvm;
using RegisterIdx = Rematerializer::RegisterIdx;

// Pin the vtable to this file.
void Rematerializer::Listener::anchor() {}

/// Checks whether the value in \p LI at \p UseIdx is identical to \p OVNI (this
/// implies it is also live there). When \p LI has sub-ranges, checks that
/// all sub-ranges intersecting with \p Mask are also live at \p UseIdx.
static bool isIdenticalAtUse(const VNInfo &OVNI, LaneBitmask Mask,
                             SlotIndex UseIdx, const LiveInterval &LI) {
  if (&OVNI != LI.getVNInfoAt(UseIdx))
    return false;

  if (LI.hasSubRanges()) {
    // Check that intersecting subranges are live at user.
    for (const LiveInterval::SubRange &SR : LI.subranges()) {
      if ((SR.LaneMask & Mask).none())
        continue;
      if (!SR.liveAt(UseIdx))
        return false;

      // Early exit if all used lanes are checked. No need to continue.
      Mask &= ~SR.LaneMask;
      if (Mask.none())
        break;
    }
  }
  return true;
}

/// If \p MO is a virtual read register, returns it. Otherwise returns the
/// sentinel register.
static Register getRegDependency(const MachineOperand &MO) {
  if (!MO.isReg() || !MO.readsReg())
    return Register();
  Register Reg = MO.getReg();
  if (Reg.isPhysical()) {
    // By the requirements on trivially rematerializable instructions, a
    // physical register use is either constant or ignorable.
    return Register();
  }
  return Reg;
}

RegisterIdx Rematerializer::rematerializeToRegion(RegisterIdx RootIdx,
                                                  unsigned UseRegion,
                                                  DependencyReuseInfo &DRI) {
  MachineInstr *FirstMI =
      getReg(RootIdx).getRegionUseBounds(UseRegion, LIS).first;
  // If there are no users in the region, rematerialize the register at the very
  // end of the region.
  MachineBasicBlock::iterator InsertPos =
      FirstMI ? FirstMI : Regions[UseRegion].second;
  RegisterIdx NewRegIdx =
      rematerializeToPos(RootIdx, UseRegion, InsertPos, DRI);
  transferRegionUsers(RootIdx, NewRegIdx, UseRegion);
  return NewRegIdx;
}

RegisterIdx
Rematerializer::rematerializeToPos(RegisterIdx RootIdx, unsigned UseRegion,
                                   MachineBasicBlock::iterator InsertPos,
                                   DependencyReuseInfo &DRI) {
  assert(!DRI.DependencyMap.contains(RootIdx));
  LLVM_DEBUG(dbgs() << "Rematerializing " << printID(RootIdx) << '\n');

  SmallVector<RegisterIdx, 2> NewDeps;
  // Copy all dependencies because recursive rematerialization of dependencies
  // may invalidate references to the backing vector of registers.
  SmallVector<RegisterIdx, 2> OldDeps(getReg(RootIdx).Dependencies);
  for (RegisterIdx DepRegIdx : OldDeps) {
    // Recursively rematerialize required dependencies at the same position as
    // the root. Registers form a DAG so the recursion is guaranteed to
    // terminate.
    auto RematIdx = DRI.DependencyMap.find(DepRegIdx);
    RegisterIdx NewDepRegIdx;
    if (RematIdx == DRI.DependencyMap.end())
      NewDepRegIdx = rematerializeToPos(DepRegIdx, UseRegion, InsertPos, DRI);
    else
      NewDepRegIdx = RematIdx->second;
    NewDeps.push_back(NewDepRegIdx);
  }
  RegisterIdx NewIdx =
      rematerializeReg(RootIdx, UseRegion, InsertPos, std::move(NewDeps));
  DRI.DependencyMap.insert({RootIdx, NewIdx});
  return NewIdx;
}

void Rematerializer::transferUser(RegisterIdx FromRegIdx, RegisterIdx ToRegIdx,
                                  unsigned UserRegion, MachineInstr &UserMI) {
  transferUserImpl(FromRegIdx, ToRegIdx, UserMI);
  Regs[FromRegIdx].eraseUser(&UserMI, UserRegion);
  Regs[ToRegIdx].addUser(&UserMI, UserRegion);
  deleteRegIfUnused(FromRegIdx);
}

void Rematerializer::transferRegionUsers(RegisterIdx FromRegIdx,
                                         RegisterIdx ToRegIdx,
                                         unsigned UseRegion) {
  auto &FromRegUsers = Regs[FromRegIdx].Uses;
  auto UsesIt = FromRegUsers.find(UseRegion);
  if (UsesIt == FromRegUsers.end())
    return;

  const SmallDenseSet<MachineInstr *, 4> &RegionUsers = UsesIt->getSecond();
  for (MachineInstr *UserMI : RegionUsers)
    transferUserImpl(FromRegIdx, ToRegIdx, *UserMI);
  Regs[ToRegIdx].addUsers(RegionUsers, UseRegion);
  FromRegUsers.erase(UseRegion);
  deleteRegIfUnused(FromRegIdx);
}

void Rematerializer::transferAllUsers(RegisterIdx FromRegIdx,
                                      RegisterIdx ToRegIdx) {
  Reg &FromReg = Regs[FromRegIdx], &ToReg = Regs[ToRegIdx];
  for (const auto &[UseRegion, RegionUsers] : FromReg.Uses) {
    for (MachineInstr *UserMI : RegionUsers)
      transferUserImpl(FromRegIdx, ToRegIdx, *UserMI);
    ToReg.addUsers(RegionUsers, UseRegion);
  }
  FromReg.Uses.clear();
  deleteRegIfUnused(FromRegIdx);
}

void Rematerializer::transferUserImpl(RegisterIdx FromRegIdx,
                                      RegisterIdx ToRegIdx,
                                      MachineInstr &UserMI) {
  assert(FromRegIdx != ToRegIdx && "identical registers");
  assert(getOriginOrSelf(FromRegIdx) == getOriginOrSelf(ToRegIdx) &&
         "unrelated registers");

  LLVM_DEBUG(dbgs() << "User transfer from " << printID(FromRegIdx) << " to "
                    << printID(ToRegIdx) << ": " << printUser(&UserMI) << '\n');

  UserMI.substituteRegister(getReg(FromRegIdx).getDefReg(),
                            getReg(ToRegIdx).getDefReg(), 0, TRI);
  LISUpdates.insert(FromRegIdx);
  LISUpdates.insert(ToRegIdx);

  // If the user is rematerializable, we must change its dependency to the
  // new register.
  if (RegisterIdx UserRegIdx = getDefRegIdx(UserMI); UserRegIdx != NoReg) {
    // Look for the user's dependency that matches the register.
    for (RegisterIdx &DepRegIdx : Regs[UserRegIdx].Dependencies) {
      if (DepRegIdx == FromRegIdx) {
        DepRegIdx = ToRegIdx;
        return;
      }
    }
    llvm_unreachable("broken dependency");
  }
}

void Rematerializer::updateLiveIntervals() {
  DenseSet<Register> SeenUnrematRegs;
  for (RegisterIdx RegIdx : LISUpdates) {
    const Reg &UpdateReg = getReg(RegIdx);
    assert(UpdateReg.isAlive() && "dead register");

    Register DefReg = UpdateReg.getDefReg();
    if (LIS.hasInterval(DefReg))
      LIS.removeInterval(DefReg);
    // Rematerializable registers have a single definition by construction so
    // re-creating their interval cannot yield a live interval with multiple
    // connected components.
    LIS.createAndComputeVirtRegInterval(DefReg);

    LLVM_DEBUG({
      dbgs() << "Re-computed interval for " << printID(RegIdx) << ": ";
      LIS.getInterval(DefReg).print(dbgs());
      dbgs() << '\n' << printRegUsers(RegIdx);
    });

    // Update intervals for unrematerializable operands.
    for (const auto &[UnrematReg, Mask] : getUnrematableDeps(RegIdx)) {
      if (!SeenUnrematRegs.insert(UnrematReg).second)
        continue;
      LIS.removeInterval(UnrematReg);
      bool NeedSplit = false;

      // Unrematerializable registers may end up with multiple connected
      // components in their live interval after it is re-created. It needs to
      // be split in such cases. We don't track unrematerializable registers by
      // their actual register index (just by operand index) so we do not need
      // to update any state in the rematerializer.
      LiveInterval &LI =
          LIS.createAndComputeVirtRegInterval(UnrematReg, NeedSplit);
      if (NeedSplit) {
        SmallVector<LiveInterval *> SplitLIs;
        LIS.splitSeparateComponents(LI, SplitLIs);
      }
      LLVM_DEBUG(
          dbgs() << "  Re-computed interval for unrematerializable register "
                 << printReg(UnrematReg, &TRI, 0, &MRI) << " with lanemask "
                 << Mask << '\n');
    }
  }
  LISUpdates.clear();
}

bool Rematerializer::isMOIdenticalAtUses(MachineOperand &MO,
                                         ArrayRef<SlotIndex> Uses) const {
  unsigned SubIdx = MO.getSubReg();
  LaneBitmask Mask = SubIdx ? TRI.getSubRegIndexLaneMask(SubIdx)
                            : MRI.getMaxLaneMaskForVReg(MO.getReg());
  return isRegIdenticalAtUses(
      MO.getReg(), Mask,
      LIS.getInstructionIndex(*MO.getParent()).getRegSlot(true), Uses);
}

bool Rematerializer::isRegIdenticalAtUses(Register Reg, LaneBitmask Mask,
                                          SlotIndex RefSlot,
                                          ArrayRef<SlotIndex> Uses) const {
  if (Uses.empty())
    return true;
  const LiveInterval &LI = LIS.getInterval(Reg);
  const VNInfo *DefVN = LI.getVNInfoAt(RefSlot);
  for (SlotIndex Use : Uses) {
    if (!isIdenticalAtUse(*DefVN, Mask, Use, LI))
      return false;
  }
  return true;
}

RegisterIdx Rematerializer::findRematInRegion(RegisterIdx RegIdx,
                                              unsigned Region,
                                              SlotIndex Before) const {
  auto It = Rematerializations.find(getOriginOrSelf(RegIdx));
  if (It == Rematerializations.end())
    return NoReg;
  const RematsOf &Remats = It->getSecond();

  SlotIndex BestSlot;
  RegisterIdx BestRegIdx = NoReg;
  for (RegisterIdx RematRegIdx : Remats) {
    const Reg &RematReg = getReg(RematRegIdx);
    if (RematReg.DefRegion != Region || RematReg.Uses.empty())
      continue;
    SlotIndex RematRegSlot =
        LIS.getInstructionIndex(*RematReg.DefMI).getRegSlot();
    if (RematRegSlot < Before &&
        (BestRegIdx == NoReg || RematRegSlot > BestSlot)) {
      BestSlot = RematRegSlot;
      BestRegIdx = RematRegIdx;
    }
  }
  return BestRegIdx;
}

void Rematerializer::deleteRegIfUnused(RegisterIdx RootIdx) {
  if (!getReg(RootIdx).Uses.empty())
    return;

  // Traverse the root's dependency DAG depth-first to find the set of registers
  // we can delete and a legal order to delete them in.
  SmallVector<RegisterIdx, 4> DepDAG{RootIdx};
  SmallVector<RegisterIdx, 8> DeleteOrder{RootIdx};
  do {
    // A deleted register's dependencies may be deletable too.
    const Reg &DeleteReg = getReg(DepDAG.pop_back_val());
    for (RegisterIdx DepRegIdx : DeleteReg.Dependencies) {
      // All dependencies lose a user (the deleted register).
      Reg &DepReg = Regs[DepRegIdx];
      DepReg.eraseUser(DeleteReg.DefMI, DeleteReg.DefRegion);
      if (DepReg.Uses.empty()) {
        DeleteOrder.push_back(DepRegIdx);
        DepDAG.push_back(DepRegIdx);
      }
    }
  } while (!DepDAG.empty());

  for (RegisterIdx RegIdx : DeleteOrder) {
    Reg &DeleteReg = Regs[RegIdx];

    // It is possible that the defined register we are deleting doesn't have an
    // interval yet if the LIS hasn't been updated since it was created.
    Register DefReg = DeleteReg.getDefReg();
    if (LIS.hasInterval(DefReg))
      LIS.removeInterval(DefReg);
    LISUpdates.erase(RegIdx);

    deleteReg(RegIdx);
    if (isRematerializedRegister(RegIdx)) {
      // Delete rematerialized register from its origin's rematerializations.
      const RegisterIdx OriginIdx = getOriginOf(RegIdx);
      RematsOf &OriginRemats = Rematerializations.at(OriginIdx);
      assert(OriginRemats.contains(RegIdx) && "broken remat<->origin link");
      OriginRemats.erase(RegIdx);
      if (OriginRemats.empty())
        Rematerializations.erase(OriginIdx);
    }
    LLVM_DEBUG(dbgs() << "** Deleted " << printID(RegIdx) << "\n");
  }
}

void Rematerializer::deleteReg(RegisterIdx RegIdx) {
  noteRegWillBeDeleted(RegIdx);

  Reg &DeleteReg = Regs[RegIdx];
  assert(DeleteReg.DefMI && "register was already deleted");
  // It is not possible for the deleted instruction to be the upper region
  // boundary since we don't ever consider them rematerializable.
  MachineBasicBlock::iterator &RegionBegin = Regions[DeleteReg.DefRegion].first;
  if (RegionBegin == DeleteReg.DefMI)
    RegionBegin = std::next(MachineBasicBlock::iterator(DeleteReg.DefMI));
  LIS.RemoveMachineInstrFromMaps(*DeleteReg.DefMI);
  DeleteReg.DefMI->eraseFromParent();
  DeleteReg.DefMI = nullptr;
}

Rematerializer::Rematerializer(MachineFunction &MF,
                               SmallVectorImpl<RegionBoundaries> &Regions,
                               LiveIntervals &LIS)
    : Regions(Regions), MRI(MF.getRegInfo()), LIS(LIS),
      TII(*MF.getSubtarget().getInstrInfo()), TRI(TII.getRegisterInfo()) {
#ifdef EXPENSIVE_CHECKS
  // Check that regions are valid.
  DenseSet<MachineInstr *> SeenMIs;
  for (const auto &[RegionBegin, RegionEnd] : Regions) {
    assert(RegionBegin != RegionEnd && "empty region");
    for (auto MI = RegionBegin; MI != RegionEnd; ++MI) {
      bool IsNewMI = SeenMIs.insert(&*MI).second;
      assert(IsNewMI && "overlapping regions");
      assert(!MI->isTerminator() && "terminator in region");
    }
    if (RegionEnd != RegionBegin->getParent()->end()) {
      bool IsNewMI = SeenMIs.insert(&*RegionEnd).second;
      assert(IsNewMI && "overlapping regions (upper bound)");
    }
  }
#endif
}

bool Rematerializer::analyze() {
  Regs.clear();
  UnrematableDeps.clear();
  Origins.clear();
  Rematerializations.clear();
  RegionMBB.clear();
  RegToIdx.clear();
  LISUpdates.clear();
  if (Regions.empty())
    return false;

  /// Maps all MIs to their parent region. Region terminators are considered
  /// part of the region they terminate.
  DenseMap<MachineInstr *, unsigned> MIRegion;

  // Initialize MI to containing region mapping.
  RegionMBB.reserve(Regions.size());
  for (unsigned I = 0, E = Regions.size(); I < E; ++I) {
    RegionBoundaries Region = Regions[I];
    assert(Region.first != Region.second && "empty cannot be region");
    for (auto MI = Region.first; MI != Region.second; ++MI) {
      assert(!MIRegion.contains(&*MI) && "regions should not intersect");
      MIRegion.insert({&*MI, I});
    }
    MachineBasicBlock &MBB = *Region.first->getParent();
    RegionMBB.push_back(&MBB);

    // A terminator instruction is considered part of the region it terminates.
    if (Region.second != MBB.end()) {
      MachineInstr *RegionTerm = &*Region.second;
      assert(!MIRegion.contains(RegionTerm) && "regions should not intersect");
      MIRegion.insert({RegionTerm, I});
    }
  }

  const unsigned NumVirtRegs = MRI.getNumVirtRegs();
  BitVector SeenRegs(NumVirtRegs);
  for (unsigned I = 0, E = NumVirtRegs; I != E; ++I) {
    if (!SeenRegs[I])
      addRegIfRematerializable(I, MIRegion, SeenRegs);
  }
  assert(Regs.size() == UnrematableDeps.size());

  LLVM_DEBUG({
    for (RegisterIdx I = 0, E = getNumRegs(); I < E; ++I)
      dbgs() << printDependencyDAG(I) << '\n';
  });
  return !Regs.empty();
}

void Rematerializer::addRegIfRematerializable(
    unsigned VirtRegIdx, const DenseMap<MachineInstr *, unsigned> &MIRegion,
    BitVector &SeenRegs) {
  assert(!SeenRegs[VirtRegIdx] && "register already seen");
  Register DefReg = Register::index2VirtReg(VirtRegIdx);
  SeenRegs.set(VirtRegIdx);

  MachineOperand *MO = MRI.getOneDef(DefReg);
  if (!MO)
    return;
  MachineInstr &DefMI = *MO->getParent();
  if (!isMIRematerializable(DefMI))
    return;
  auto DefRegion = MIRegion.find(&DefMI);
  if (DefRegion == MIRegion.end())
    return;

  Reg RematReg;
  RematReg.DefMI = &DefMI;
  RematReg.DefRegion = DefRegion->second;
  unsigned SubIdx = DefMI.getOperand(0).getSubReg();
  RematReg.Mask = SubIdx ? TRI.getSubRegIndexLaneMask(SubIdx)
                         : MRI.getMaxLaneMaskForVReg(DefReg);

  // Collect the candidate's direct users, both rematerializable and
  // unrematerializable. MIs outside provided regions cannot be tracked so the
  // registers they use are not safely rematerializable.
  for (MachineInstr &UseMI : MRI.use_nodbg_instructions(DefReg)) {
    if (auto UseRegion = MIRegion.find(&UseMI); UseRegion != MIRegion.end())
      RematReg.addUser(&UseMI, UseRegion->second);
    else
      return;
  }
  if (RematReg.Uses.empty())
    return;

  // Collect the candidate's dependencies, rematerializable or not. If the same
  // rematerializable register is used multiple times we just need to consider
  // it once.
  SmallSetVector<RegisterIdx, 2> RematDeps;
  SmallMapVector<Register, LaneBitmask, 2> UnrematDeps;
  for (const MachineOperand &MO : DefMI.all_uses()) {
    Register DepReg = getRegDependency(MO);
    if (!DepReg)
      continue;
    unsigned DepRegIdx = DepReg.virtRegIndex();
    if (!SeenRegs[DepRegIdx])
      addRegIfRematerializable(DepRegIdx, MIRegion, SeenRegs);
    if (auto DepIt = RegToIdx.find(DepReg); DepIt != RegToIdx.end()) {
      RematDeps.insert(DepIt->second);
    } else {
      LaneBitmask &CurrentMask =
          UnrematDeps.try_emplace(DepReg, LaneBitmask::getNone()).first->second;
      LaneBitmask Mask = MO.getSubReg()
                             ? TRI.getSubRegIndexLaneMask(MO.getSubReg())
                             : MRI.getMaxLaneMaskForVReg(DepReg);
      CurrentMask |= Mask;
    }
  }

  // The register is rematerializable.
  RematReg.Dependencies = RematDeps.takeVector();
  RegToIdx.insert({DefReg, Regs.size()});
  Regs.push_back(RematReg);
  UnrematableDeps.push_back(UnrematDeps.takeVector());
}

bool Rematerializer::isMIRematerializable(const MachineInstr &MI) const {
  if (!TII.isReMaterializable(MI))
    return false;

  assert(MI.getOperand(0).getReg().isVirtual() && "should be virtual");
  assert(MRI.hasOneDef(MI.getOperand(0).getReg()) && "should have single def");

  for (const MachineOperand &MO : MI.all_uses()) {
    // We can't remat physreg uses, unless it is a constant or an ignorable
    // use (e.g. implicit exec use on VALU instructions)
    if (MO.getReg().isPhysical()) {
      if (MRI.isConstantPhysReg(MO.getReg()) || TII.isIgnorableUse(MO))
        continue;
      return false;
    }
  }

  return true;
}

RegisterIdx Rematerializer::getDefRegIdx(const MachineInstr &MI) const {
  if (!MI.getNumOperands() || !MI.getOperand(0).isReg() ||
      MI.getOperand(0).readsReg())
    return NoReg;
  Register Reg = MI.getOperand(0).getReg();
  auto UserRegIt = RegToIdx.find(Reg);
  if (UserRegIt == RegToIdx.end())
    return NoReg;
  return UserRegIt->second;
}

RegisterIdx
Rematerializer::rematerializeReg(RegisterIdx RegIdx, unsigned UseRegion,
                                 MachineBasicBlock::iterator InsertPos,
                                 SmallVectorImpl<RegisterIdx> &&Dependencies) {
  RegisterIdx NewRegIdx = Regs.size();

  Reg &NewReg = Regs.emplace_back();
  Reg &FromReg = Regs[RegIdx];
  NewReg.Mask = FromReg.Mask;
  NewReg.DefRegion = UseRegion;
  NewReg.Dependencies = std::move(Dependencies);

  // Track rematerialization link between registers. Origins are always
  // registers that existed originally, and rematerializations are always
  // attached to them.
  const RegisterIdx OriginIdx = getOriginOrSelf(RegIdx);
  Origins.push_back(OriginIdx);
  Rematerializations[OriginIdx].insert(NewRegIdx);

  // Use the TII to rematerialize the defining instruction with a new defined
  // register.
  Register NewDefReg = MRI.cloneVirtualRegister(FromReg.getDefReg());
  TII.reMaterialize(*RegionMBB[UseRegion], InsertPos, NewDefReg, 0,
                    *FromReg.DefMI);
  NewReg.DefMI = &*std::prev(InsertPos);
  RegToIdx.insert({NewDefReg, NewRegIdx});
  postRematerialization(RegIdx, NewRegIdx);

  noteRegCreated(NewRegIdx);
  LLVM_DEBUG(dbgs() << "** Rematerialized " << printID(RegIdx) << " as "
                    << printRematReg(NewRegIdx) << '\n');
  return NewRegIdx;
}

void Rematerializer::recreateReg(RegisterIdx RegIdx,
                                 MachineBasicBlock::iterator InsertPos,
                                 Register DefReg) {
  assert(RegToIdx.contains(DefReg) && "unknown defined register");
  assert(RegToIdx.at(DefReg) == RegIdx && "incorrect defined register");
  assert(!getReg(RegIdx).isAlive() && "register is still alive");

  Reg &OriginReg = Regs[RegIdx];

  // Re-establish the link between origin and rematerialization if necessary.
  const bool RecreateOriginalReg = isOriginalRegister(RegIdx);
  if (!RecreateOriginalReg)
    Rematerializations[getOriginOf(RegIdx)].insert(RegIdx);

  // Rematerialize from one of the existing rematerializations or from the
  // origin. We expect at least one to exist, otherwise it would mean the value
  // held by the original register is no longer available anywhere in the MF.
  RegisterIdx ModelRegIdx;
  if (RecreateOriginalReg) {
    assert(Rematerializations.contains(RegIdx) && "expected remats");
    ModelRegIdx = *Rematerializations.at(RegIdx).begin();
  } else {
    assert(getReg(getOriginOf(RegIdx)).isAlive() && "expected alive origin");
    ModelRegIdx = getOriginOf(RegIdx);
  }
  const MachineInstr &ModelDefMI = *getReg(ModelRegIdx).DefMI;

  TII.reMaterialize(*RegionMBB[OriginReg.DefRegion], InsertPos, DefReg, 0,
                    ModelDefMI);
  OriginReg.DefMI = &*std::prev(InsertPos);
  postRematerialization(ModelRegIdx, RegIdx);
  LLVM_DEBUG(dbgs() << "** Recreated " << printID(RegIdx) << " as "
                    << printRematReg(RegIdx) << '\n');
}

void Rematerializer::postRematerialization(RegisterIdx ModelRegIdx,
                                           RegisterIdx RematRegIdx) {

  // The start of the new register's region may have changed.
  Reg &ModelReg = Regs[ModelRegIdx], &RematReg = Regs[RematRegIdx];
  LIS.InsertMachineInstrInMaps(*RematReg.DefMI);
  MachineBasicBlock::iterator &RegionBegin = Regions[RematReg.DefRegion].first;
  if (RegionBegin == std::next(MachineBasicBlock::iterator(RematReg.DefMI)))
    RegionBegin = RematReg.DefMI;

  // Replace dependencies as needed in the rematerialized MI. All dependencies
  // of the latter gain a new user.
  auto ZipedDeps = zip_equal(ModelReg.Dependencies, RematReg.Dependencies);
  for (const auto &[OldDepRegIdx, NewDepRegIdx] : ZipedDeps) {
    LLVM_DEBUG(dbgs() << "  Dependency: " << printID(OldDepRegIdx) << " -> "
                      << printID(NewDepRegIdx) << '\n');

    Reg &NewDepReg = Regs[NewDepRegIdx];
    if (OldDepRegIdx != NewDepRegIdx) {
      Reg &OldDepReg = Regs[OldDepRegIdx];
      RematReg.DefMI->substituteRegister(OldDepReg.getDefReg(),
                                         NewDepReg.getDefReg(), 0, TRI);
      LISUpdates.insert(OldDepRegIdx);
    }
    NewDepReg.addUser(RematReg.DefMI, RematReg.DefRegion);
    LISUpdates.insert(NewDepRegIdx);
  }
}

std::pair<MachineInstr *, MachineInstr *>
Rematerializer::Reg::getRegionUseBounds(unsigned UseRegion,
                                        const LiveIntervals &LIS) const {
  auto It = Uses.find(UseRegion);
  if (It == Uses.end())
    return {nullptr, nullptr};
  const RegionUsers &RegionUsers = It->getSecond();
  assert(!RegionUsers.empty() && "empty userset in region");

  auto User = RegionUsers.begin(), UserEnd = RegionUsers.end();
  MachineInstr *FirstMI = *User, *LastMI = FirstMI;
  SlotIndex FirstIndex = LIS.getInstructionIndex(*FirstMI),
            LastIndex = FirstIndex;

  while (++User != UserEnd) {
    SlotIndex UserIndex = LIS.getInstructionIndex(**User);
    if (UserIndex < FirstIndex) {
      FirstIndex = UserIndex;
      FirstMI = *User;
    } else if (UserIndex > LastIndex) {
      LastIndex = UserIndex;
      LastMI = *User;
    }
  }

  return {FirstMI, LastMI};
}

void Rematerializer::Reg::addUser(MachineInstr *MI, unsigned Region) {
  Uses[Region].insert(MI);
}

void Rematerializer::Reg::addUsers(const RegionUsers &NewUsers,
                                   unsigned Region) {
  Uses[Region].insert_range(NewUsers);
}

void Rematerializer::Reg::eraseUser(MachineInstr *MI, unsigned Region) {
  RegionUsers &RUsers = Uses.at(Region);
  assert(RUsers.contains(MI) && "user not in region");
  if (RUsers.size() == 1)
    Uses.erase(Region);
  else
    RUsers.erase(MI);
}

Printable Rematerializer::printDependencyDAG(RegisterIdx RootIdx) const {
  return Printable([&, RootIdx](raw_ostream &OS) {
    DenseMap<RegisterIdx, unsigned> RegDepths;
    std::function<void(RegisterIdx, unsigned)> WalkTree =
        [&](RegisterIdx RegIdx, unsigned Depth) -> void {
      unsigned MaxDepth = std::max(RegDepths.lookup_or(RegIdx, Depth), Depth);
      RegDepths.emplace_or_assign(RegIdx, MaxDepth);
      for (RegisterIdx DepRegIdx : getReg(RegIdx).Dependencies)
        WalkTree(DepRegIdx, Depth + 1);
    };
    WalkTree(RootIdx, 0);

    // Sort in decreasing depth order to print root at the bottom.
    SmallVector<std::pair<RegisterIdx, unsigned>> Regs(RegDepths.begin(),
                                                       RegDepths.end());
    sort(Regs, [](const auto &LHS, const auto &RHS) {
      return LHS.second > RHS.second;
    });

    OS << printID(RootIdx) << " has " << Regs.size() - 1 << " dependencies\n";
    for (const auto &[RegIdx, Depth] : Regs) {
      OS << indent(Depth, 2) << (Depth ? '|' : '*') << ' '
         << printRematReg(RegIdx, /*SkipRegions=*/Depth) << '\n';
    }
    OS << printRegUsers(RootIdx);
  });
}

Printable Rematerializer::printID(RegisterIdx RegIdx) const {
  return Printable([&, RegIdx](raw_ostream &OS) {
    const Reg &PrintReg = getReg(RegIdx);
    OS << '(' << RegIdx << '/';
    if (!PrintReg.isAlive()) {
      OS << "<dead>";
    } else {
      OS << printReg(PrintReg.getDefReg(), &TRI,
                     PrintReg.DefMI->getOperand(0).getSubReg(), &MRI);
    }
    OS << ")[" << PrintReg.DefRegion << "]";
  });
}

Printable Rematerializer::printRematReg(RegisterIdx RegIdx,
                                        bool SkipRegions) const {
  return Printable([&, RegIdx, SkipRegions](raw_ostream &OS) {
    const Reg &PrintReg = getReg(RegIdx);
    OS << printID(RegIdx);
    if (!SkipRegions) {
      OS << " [" << PrintReg.DefRegion;
      if (!PrintReg.Uses.empty()) {
        assert(PrintReg.isAlive() && "dead register cannot have uses");
        const LiveInterval &LI = LIS.getInterval(PrintReg.getDefReg());
        // First display all regions in which the register is live-through and
        // not used.
        bool First = true;
        for (const auto &[I, Bounds] : enumerate(Regions)) {
          if (PrintReg.Uses.contains(I))
            continue;
          // The register must be live at the live-ins and live-outs of the
          // region.
          MachineBasicBlock::iterator LiveIn =
              skipDebugInstructionsForward(Bounds.first, Bounds.second);
          if (LiveIn == Bounds.second) {
            // The region has no non-debug instructions, it's hard to assess
            // whether the register is live across it without an index.
            continue;
          }
          // LiveIn is inside the range and a non-debug instruction so we know
          // this will also point to a non-debug instruction within the region.
          MachineBasicBlock::iterator LiveOut = skipDebugInstructionsBackward(
              std::prev(Bounds.second), Bounds.first);
          if (LI.liveAt(LIS.getInstructionIndex(*LiveIn)) &&
              LI.liveAt(LIS.getInstructionIndex(*LiveOut).getDeadSlot())) {
            OS << (First ? " - " : ",") << I;
            First = false;
          }
        }
        OS << (First ? " --> " : " -> ");

        // Then display regions in which the register is used.
        auto It = PrintReg.Uses.begin();
        OS << It->first;
        while (++It != PrintReg.Uses.end())
          OS << "," << It->first;
      }
      OS << "] ";
    }
    if (PrintReg.isAlive()) {
      PrintReg.DefMI->print(OS, /*IsStandalone=*/true, /*SkipOpers=*/false,
                            /*SkipDebugLoc=*/false, /*AddNewLine=*/false);
      OS << " @ ";
      LIS.getInstructionIndex(*PrintReg.DefMI).print(OS);
    }
  });
}

Printable Rematerializer::printRegUsers(RegisterIdx RegIdx) const {
  return Printable([&, RegIdx](raw_ostream &OS) {
    for (const auto &[UseRegion, Users] : getReg(RegIdx).Uses) {
      for (MachineInstr *MI : Users)
        OS << "  User " << printUser(MI, UseRegion) << '\n';
    }
  });
}

Printable Rematerializer::printUser(const MachineInstr *MI,
                                    std::optional<unsigned> UseRegion) const {
  return Printable([&, MI, UseRegion](raw_ostream &OS) {
    RegisterIdx RegIdx = getDefRegIdx(*MI);
    if (RegIdx != NoReg) {
      OS << printID(RegIdx);
    } else {
      OS << "(-/-)[";
      if (UseRegion)
        OS << *UseRegion;
      else
        OS << '?';
      OS << ']';
    }
    OS << ' ';
    MI->print(OS, /*IsStandalone=*/true, /*SkipOpers=*/false,
              /*SkipDebugLoc=*/false, /*AddNewLine=*/false);
    OS << " @ ";
    LIS.getInstructionIndex(*MI).print(OS);
  });
}

void Rollbacker::rematerializerNoteRegCreated(const Rematerializer &Remater,
                                              RegisterIdx RegIdx) {
  if (RollingBack)
    return;
  assert(Remater.isRematerializedRegister(RegIdx) && "only remats are created");
  Rematerializations[Remater.getOriginOf(RegIdx)].insert(RegIdx);
}

void Rollbacker::rematerializerNoteRegWillBeDeleted(
    const Rematerializer &Remater, RegisterIdx RegIdx) {
  if (RollingBack)
    return;

  // Find a valid re-creation position after the register's definition.
  MachineInstr *DefMI = Remater.getReg(RegIdx).DefMI;
  MachineBasicBlock *ParentMBB = DefMI->getParent();
  MachineBasicBlock::iterator ValidPos = std::next(DefMI->getIterator());
  while (ValidPos != ParentMBB->end() && isRollbackableMI(*ValidPos, Remater))
    ValidPos = std::next(ValidPos);

  if (Remater.isRematerializedRegister(RegIdx)) {
    // Rematerializations will not be re-created. Previously deleted registers
    // that reference this register's defining instruction as their re-creation
    // position should instead be re-created at a valid position after the
    // deleted MI.
    invalidatePosition(DefMI, ValidPos);
    return;
  }

  // Original registers can be re-created. Add a re-creation position for the
  // definition of the rematerializable register.
  DeadRegs.push_back(DeadReg(RegIdx, Remater));
  const InsertBeforePos InsertPos = makePos(ValidPos, ParentMBB);
  PosToIdx[InsertPos].insert(Positions.size());
  Positions.push_back(InsertPos);
}

void Rollbacker::rollback(Rematerializer &Remater) {
  RollingBack = true;

  // As we re-create registers, map deleted definitions to re-created ones. This
  // allows to replace invalid re-creation positions that reference deleted
  // definitions to valid new positions while restoring original MI order.
  DenseMap<MachineInstr *, MachineInstr *> Replacements;
  unsigned PositionIndex = Positions.size();

  // Re-create deleted registers in reverse order of deletion. Related registers
  // are deleted in reverse def-use order so this ensures we re-create registers
  // in def-use order. This also ensures that re-creation positions that became
  // invalid due to later MI deletions can be corrected as we go.
  for (const DeadReg &Reg : reverse(DeadRegs)) {
    assert(!Remater.getReg(Reg.Idx).isAlive() && "register should be dead");

    // Determine re-creation position for the register's definition.
    MachineBasicBlock::iterator InsertPosition;
    InsertBeforePos Pos = Positions[--PositionIndex];
    if (auto *MBB = dyn_cast<MachineBasicBlock *>(Pos)) {
      InsertPosition = MBB->end();
    } else {
      auto *MI = cast<MachineInstr *>(Pos);
      InsertPosition = Replacements.lookup_or(MI, MI)->getIterator();
    }

    Remater.recreateReg(Reg.Idx, InsertPosition, Reg.DefReg);

    const Rematerializer::Reg &RecreateReg = Remater.getReg(Reg.Idx);
    if (!Replacements.insert({Reg.DefMI, RecreateReg.DefMI}).second)
      llvm_unreachable("duplicate deleted MI");
  }

  // Rollback rematerializations.
  for (const auto &[RegIdx, RematsOf] : Rematerializations) {
    for (RegisterIdx RematRegIdx : RematsOf) {
      // It is possible that rematerializations were deleted. Their users would
      // have been transfered to some other rematerialization so we can safely
      // ignore them. Original registers that were deleted were just re-created
      // so we do not need to check for that.
      if (Remater.getReg(RematRegIdx).isAlive())
        Remater.transferAllUsers(RematRegIdx, RegIdx);
    }
  }

  Remater.updateLiveIntervals();
  DeadRegs.clear();
  Positions.clear();
  PosToIdx.clear();
  Rematerializations.clear();
  RollingBack = false;
}

bool Rollbacker::isRollbackableMI(const MachineInstr &MI,
                                  const Rematerializer &Remater) const {
  RegisterIdx RegIdx = Remater.getDefRegIdx(MI);
  if (RegIdx == Rematerializer::NoReg ||
      !Remater.isRematerializedRegister(RegIdx))
    return false;
  // It is possible that the MI defines a rematerializable register that was not
  // recorded if the rollbacker was attached to the rematerializer after the
  // rematerialization happened. In such cases the MI won't be rolled back.
  auto RematsOf = Rematerializations.find(Remater.getOriginOf(RegIdx));
  if (RematsOf == Rematerializations.end())
    return false;
  return RematsOf->getSecond().contains(RegIdx);
}

void Rollbacker::invalidatePosition(MachineInstr *MI,
                                    MachineBasicBlock::iterator It) {
  const InsertBeforePos MIPos = InsertBeforePos(MI),
                        NewPos = makePos(It, MI->getParent());
  auto MIIndices = PosToIdx.find(MIPos);
  if (MIIndices == PosToIdx.end())
    return;
  assert(!MIIndices->getSecond().empty() && "no index hold position");
  for (unsigned I : MIIndices->getSecond())
    Positions[I] = NewPos;
  PosToIdx.try_emplace(NewPos).first->getSecond().insert_range(
      MIIndices->getSecond());
  PosToIdx.erase(MIPos);
}