text: //===--- Passes/ShrinkWrapping.h ------------------------------------------===//
  //
  //                     The LLVM Compiler Infrastructure
  //
  // This file is distributed under the University of Illinois Open Source
  // License. See LICENSE.TXT for details.
  //
  //===----------------------------------------------------------------------===//
  //
  //===----------------------------------------------------------------------===//
  
  #ifndef LLVM_TOOLS_LLVM_BOLT_PASSES_SHRINKWRAPPING_H
  #define LLVM_TOOLS_LLVM_BOLT_PASSES_SHRINKWRAPPING_H
  
  #include "BinaryPasses.h"
  #include "FrameAnalysis.h"
  #include "DataflowInfoManager.h"
  
  namespace llvm {
  namespace bolt {
  
  /// Encapsulates logic required to analyze a binary function and detect which
  /// registers are being saved as callee-saved, where are these saves and where
  /// are the points where their original value are being restored.
  class CalleeSavedAnalysis {
    const FrameAnalysis &FA;
    const BinaryContext &BC;
    BinaryFunction &BF;
    DataflowInfoManager &Info;
    Optional<unsigned> SaveTagIndex;
    Optional<unsigned> RestoreTagIndex;
  
    /// Compute all stores of callee-saved regs. Those are the ones that stores a
    /// register whose definition is not local.
    void analyzeSaves();
  
    /// Similar to analyzeSaves, tries to determine all instructions that recover
    /// the original value of the callee-saved register before exiting the
    /// function.
    void analyzeRestores();
  
    /// Returns the identifying string used to annotate instructions with metadata
    /// for this analysis. These are deleted in the destructor.
    static StringRef getSaveTagName() {
      return StringRef("CSA-SavedReg");
    }
  
    unsigned getSaveTag() {
      if (SaveTagIndex)
        return *SaveTagIndex;
      SaveTagIndex = BC.MIB->getOrCreateAnnotationIndex(getSaveTagName());
      return *SaveTagIndex;
    }
  
    static StringRef getRestoreTagName() {
      return StringRef("CSA-RestoredReg");
    }
  
    unsigned getRestoreTag() {
      if (RestoreTagIndex)
        return *RestoreTagIndex;
      RestoreTagIndex = BC.MIB->getOrCreateAnnotationIndex(getRestoreTagName());
      return *RestoreTagIndex;
    }
  
  public:
    BitVector CalleeSaved;
    std::vector<int64_t> OffsetsByReg;
    BitVector HasRestores;
    std::vector<uint64_t> SavingCost;
    std::vector<const FrameIndexEntry*> SaveFIEByReg;
    std::vector<const FrameIndexEntry*> LoadFIEByReg;
  
    CalleeSavedAnalysis(const FrameAnalysis &FA, const BinaryContext &BC,
                        BinaryFunction &BF, DataflowInfoManager &Info)
        : FA(FA), BC(BC), BF(BF), Info(Info),
          CalleeSaved(BC.MRI->getNumRegs(), false),
          OffsetsByReg(BC.MRI->getNumRegs(), 0LL),
          HasRestores(BC.MRI->getNumRegs(), false),
          SavingCost(BC.MRI->getNumRegs(), 0ULL),
          SaveFIEByReg(BC.MRI->getNumRegs(), nullptr),
          LoadFIEByReg(BC.MRI->getNumRegs(), nullptr) {}
  
    ~CalleeSavedAnalysis();
  
    void compute() {
      analyzeSaves();
      analyzeRestores();
    }
  
    /// Retrieves the value of the callee-saved register that is saved by this
    /// instruction or 0 if this is not a CSR save instruction.
    uint16_t getSavedReg(const MCInst &Inst) {
      auto Val = BC.MIB->tryGetAnnotationAs<decltype(FrameIndexEntry::RegOrImm)>(
          Inst, getSaveTag());
      if (Val)
        return *Val;
      return 0;
    }
  
    /// Retrieves the value of the callee-saved register that is restored by this
    /// instruction or 0 if this is not a CSR restore instruction.
    uint16_t getRestoredReg(const MCInst &Inst) {
      auto Val = BC.MIB->tryGetAnnotationAs<decltype(FrameIndexEntry::RegOrImm)>(
          Inst, getRestoreTag());
      if (Val)
        return *Val;
      return 0;
    }
  
    /// Routines to compute all saves/restores for a Reg (needs to traverse all
    /// instructions).
    std::vector<MCInst *> getSavesByReg(uint16_t Reg);
    std::vector<MCInst *> getRestoresByReg(uint16_t Reg);
  };
  
  /// Identifies in a given binary function all stack regions being used and allow
  /// us to edit the layout, removing or inserting new regions. When the layout is
  /// modified, all affected stack-accessing instructions are updated.
  class StackLayoutModifier {
    const FrameAnalysis &FA;
    const BinaryContext &BC;
    BinaryFunction &BF;
    DataflowInfoManager &Info;
  
    // Keep track of stack slots we know how to safely move
    std::map<int64_t, int64_t> AvailableRegions;
  
    DenseSet<int64_t> CollapsedRegions;
    DenseSet<int64_t> InsertedRegions;
  
    // A map of chunks of stack memory we don't really know what's happening there
    // and we need to leave it untouched.
    std::map<int64_t, int64_t> BlacklistedRegions;
  
    // Maps stack slots to the regs that are saved to them
    DenseMap<int64_t, std::set<MCPhysReg>> RegionToRegMap;
    DenseMap<int, std::set<int64_t>> RegToRegionMap;
  
    // If we can't understand how to move stack slots, IsSimple will be false
    bool IsSimple{true};
  
    bool IsInitialized{false};
  
    Optional<unsigned> TodoTagIndex;
    Optional<unsigned> SlotTagIndex;
    Optional<unsigned> OffsetCFIRegTagIndex;
  
  public:
    // Keep a worklist of operations to perform on the function to perform
    // the requested layout modifications via collapseRegion()/insertRegion().
    struct WorklistItem {
      enum ActionType : uint8_t {
        None = 0,
        AdjustLoadStoreOffset,
        AdjustCFI,
      } Action;
  
      int64_t OffsetUpdate{0};
      WorklistItem() : Action(None) {}
      WorklistItem(ActionType Action) : Action(Action) {}
      WorklistItem(ActionType Action, int OffsetUpdate)
          : Action(Action), OffsetUpdate(OffsetUpdate) {}
    };
  private:
  
    /// Mark the stack region identified by \p Offset and \p Size to be a
    /// no-touch zone, whose accesses cannot be relocated to another region.
    void blacklistRegion(int64_t Offset, int64_t Size);
  
    /// Check if this region overlaps with blacklisted addresses
    bool isRegionBlacklisted(int64_t Offset, int64_t Size);
  
    /// Check if the region identified by \p Offset and \p Size has any conflicts
    /// with available regions so far. If it has, blacklist all involved regions
    /// and return true.
    bool blacklistAllInConflictWith(int64_t Offset, int64_t Size);
  
    /// If \p Point is identified as frame pointer initialization (defining the
    /// value of FP with SP), check for non-standard initialization that precludes
    /// us from changing the stack layout. If positive, update blacklisted
    /// regions.
    void checkFramePointerInitialization(MCInst &Point);
  
    /// If \p Point is restoring the value with SP with FP plus offset,
    /// add a slottag to this instruction as it needs to be updated when we
    /// change the stack layout.
    void checkStackPointerRestore(MCInst &Point);
  
    /// Make sense of each stack offsets we can freely change
    void classifyStackAccesses();
    void classifyCFIs();
  
    /// Used to keep track of modifications to the function that will later be
    /// performed by performChanges();
    void scheduleChange(MCInst &Inst, WorklistItem Item);
  
    unsigned getTodoTag() {
      if (TodoTagIndex)
        return *TodoTagIndex;
      TodoTagIndex = BC.MIB->getOrCreateAnnotationIndex(getTodoTagName());
      return *TodoTagIndex;
    }
  
    unsigned getSlotTag() {
      if (SlotTagIndex)
        return *SlotTagIndex;
      SlotTagIndex = BC.MIB->getOrCreateAnnotationIndex(getSlotTagName());
      return *SlotTagIndex;
    }
  
    unsigned getOffsetCFIRegTag() {
      if (OffsetCFIRegTagIndex)
        return *OffsetCFIRegTagIndex;
      OffsetCFIRegTagIndex =
        BC.MIB->getOrCreateAnnotationIndex(getOffsetCFIRegTagName());
      return *OffsetCFIRegTagIndex;
    }
  
    static StringRef getTodoTagName() {
      return StringRef("SLM-TodoTag");
    }
    static StringRef getSlotTagName() {
      return StringRef("SLM-SlotTag");
    }
    static StringRef getOffsetCFIRegTagName() {
      return StringRef("SLM-OffsetCFIReg");
    }
  
  public:
    StackLayoutModifier(const FrameAnalysis &FA, const BinaryContext &BC,
                        BinaryFunction &BF, DataflowInfoManager &Info)
        : FA(FA), BC(BC), BF(BF), Info(Info) {}
  
    ~StackLayoutModifier() {
      for (auto &BB : BF) {
        for (auto &Inst : BB) {
          BC.MIB->removeAnnotation(Inst, getTodoTag());
          BC.MIB->removeAnnotation(Inst, getSlotTag());
          BC.MIB->removeAnnotation(Inst, getOffsetCFIRegTag());
        }
      }
    }
  
    /// Retrieves the value of the callee-saved register that is restored by this
    /// instruction or 0 if this is not a CSR restore instruction.
    uint16_t getOffsetCFIReg(const MCInst &Inst) {
      auto Val =
          BC.MIB->tryGetAnnotationAs<uint16_t>(Inst, getOffsetCFIRegTag());
      if (Val)
        return *Val;
      return 0;
    }
  
    /// Check if it is possible to delete the push instruction \p DeletedPush.
    /// This involves collapsing the region accessed by this push and updating all
    /// other instructions that access affected memory regions. Return true if we
    /// can update this.
    bool canCollapseRegion(int64_t RegionAddr);
    bool canCollapseRegion(MCInst *DeletedPush);
  
    /// Notify the layout manager that \p DeletedPush was deleted and that it
    /// needs to update other affected stack-accessing instructions.
    bool collapseRegion(MCInst *Alloc, int64_t RegionAddr, int64_t RegionSize);
    bool collapseRegion(MCInst *DeletedPush);
  
    /// Set the new stack address difference for load/store instructions that
    /// referenced a stack location that was deleted via collapseRegion.
    void setOffsetForCollapsedAccesses(int64_t NewOffset);
  
    /// Check if it is possible to insert a push instruction at point \p P.
    /// This involves inserting a new region in the stack, possibly affecting
    /// instructions that access the frame. Return true if we can update them all.
    bool canInsertRegion(ProgramPoint P);
  
    /// Notify the layout manager that a new push instruction has been inserted
    /// at point \p P and that it will need to update relevant instructions.
    bool insertRegion(ProgramPoint P, int64_t RegionSz);
  
    /// Perform all changes scheduled by collapseRegion()/insertRegion()
    void performChanges();
  
    /// Perform initial assessment of the function trying to understand its stack
    /// accesses.
    void initialize();
  };
  
  /// Implements a pass to optimize callee-saved register spills. These spills
  /// typically happen at function prologue/epilogue. When these are hot basic
  /// blocks, this pass will try to move these spills to cold blocks whenever
  /// possible.
  class ShrinkWrapping {
    const FrameAnalysis &FA;
    const BinaryContext &BC;
    BinaryFunction &BF;
    DataflowInfoManager &Info;
    StackLayoutModifier SLM;
    /// For each CSR, store a vector of all CFI indexes deleted as a consequence
    /// of moving this Callee-Saved Reg
    DenseMap<unsigned, std::vector<uint32_t>> DeletedPushCFIs;
    DenseMap<unsigned, std::vector<uint32_t>> DeletedPopCFIs;
    std::vector<bool> HasDeletedOffsetCFIs;
    SmallPtrSet<const MCCFIInstruction *, 16> UpdatedCFIs;
    std::vector<BitVector> UsesByReg;
    std::vector<int64_t> PushOffsetByReg;
    std::vector<int64_t> PopOffsetByReg;
    std::vector<MCPhysReg> DomOrder;
    CalleeSavedAnalysis CSA;
    std::vector<SmallSetVector<MCInst *, 4>> SavePos;
    std::vector<uint64_t> BestSaveCount;
    std::vector<MCInst *> BestSavePos;
  
    /// Pass stats
    static uint64_t SpillsMovedRegularMode;
    static uint64_t SpillsMovedPushPopMode;
  
    Optional<unsigned> AnnotationIndex;
  
    /// Allow our custom worklist-sensitive analysis
    /// PredictiveStackPointerTracking to access WorklistItem
  public:
    struct WorklistItem {
      enum ActionType : uint8_t {
        Erase = 0,
        ChangeToAdjustment,
        InsertLoadOrStore,
        InsertPushOrPop
      } Action;
      FrameIndexEntry FIEToInsert;
      unsigned AffectedReg;
      int Adjustment{0};
      WorklistItem(ActionType Action, unsigned AffectedReg)
          : Action(Action), FIEToInsert(), AffectedReg(AffectedReg) {}
      WorklistItem(ActionType Action, unsigned AffectedReg, int Adjustment)
          : Action(Action), FIEToInsert(), AffectedReg(AffectedReg),
            Adjustment(Adjustment) {}
      WorklistItem(ActionType Action, const FrameIndexEntry &FIE,
                   unsigned AffectedReg)
          : Action(Action), FIEToInsert(FIE), AffectedReg(AffectedReg) {}
    };
  
    /// Insertion todo items scheduled to happen at the end of BBs. Since we
    /// can't annotate BBs we maintain this bookkeeping here.
    DenseMap<BinaryBasicBlock*, std::vector<WorklistItem>> Todo;
  
    /// Annotation name used to tag instructions with removal or insertion actions
    static StringRef getAnnotationName() {
      return StringRef("ShrinkWrap-Todo");
    }
  
    unsigned getAnnotationIndex() {
      if (AnnotationIndex)
        return *AnnotationIndex;
      AnnotationIndex = BC.MIB->getOrCreateAnnotationIndex(getAnnotationName());
      return *AnnotationIndex;
    }
  
  private:
    using BBIterTy = BinaryBasicBlock::iterator;
  
    /// Calculate all possible uses/defs of these callee-saved regs
    void classifyCSRUses();
  
    // Ensure we don't work on cases where there are no uses of the callee-saved
    // register. These unnecessary spills should have been removed by previous
    // passes.
    void pruneUnwantedCSRs();
  
    // Map regs to their possible save possibilities (at start of these BBs)
    void computeSaveLocations();
  
    /// Look into the best save location found for saving callee-saved reg
    /// \p CSR and evaluates whether we would benefit by moving the spill to this
    /// new save location. Returns true in case it is profitable to perform the
    /// move.
    bool validateBestSavePos(unsigned CSR, MCInst *&BestPosSave,
                             uint64_t &TotalEstimatedWin);
  
    /// Populate the Todo map with worklistitems to change the function
    template <typename ...T>
    void scheduleChange(ProgramPoint PP, T&& ...Item) {
      if (PP.isInst()) {
        auto &WList = BC.MIB->getOrCreateAnnotationAs<std::vector<WorklistItem>>(
            *PP.getInst(), getAnnotationIndex());
        WList.emplace_back(std::forward<T>(Item)...);
        return;
      }
      // Avoid inserting on BBs with no instructions because we have a dataflow
      // analysis that depends on insertions happening before real instructions
      // (PredictiveStackPointerTracking)
      BinaryBasicBlock *BB = PP.getBB();
      if (BB->size() != 0) {
        Todo[BB].emplace_back(std::forward<T>(Item)...);
        return;
      }
      while (BB->size() == 0) {
        assert (BB->succ_size() == 1);
        BB = *BB->succ_begin();
      }
      auto &WList = BC.MIB->getOrCreateAnnotationAs<std::vector<WorklistItem>>(
        *BB->begin(), getAnnotationIndex());
      WList.emplace_back(std::forward<T>(Item)...);
    }
  
    /// Determine the POP ordering according to which CSR save is the dominator.
    void computeDomOrder();
  
    /// Check that the best possible location for a spill save (as determined by
    /// computeSaveLocations) is cold enough to be worth moving the save to it.
    /// \p CSR is the callee-saved register number, \p BestPosSave returns the
    /// pointer to the cold location in case the function returns true, while
    /// \p TotalEstimatedWin contains the ins dyn count reduction after moving.
    bool isBestSavePosCold(unsigned CSR, MCInst *&BestPosSave,
                           uint64_t &TotalEstimatedWin);
  
    /// Auxiliary function used to create basic blocks for critical edges and
    /// update the dominance frontier with these new locations
    void splitFrontierCritEdges(
        BinaryFunction *Func, SmallVector<ProgramPoint, 4> &Frontier,
        const SmallVector<bool, 4> &IsCritEdge,
        const SmallVector<BinaryBasicBlock *, 4> &From,
        const SmallVector<SmallVector<BinaryBasicBlock *, 4>, 4> &To);
  
    /// After the best save location for a spill has been established in
    /// \p BestPosSave for reg \p CSR, compute adequate locations to restore
    /// the spilled value. This will be at the dominance frontier.
    /// Returns an empty vector if we failed. In case of success, set
    /// \p UsePushPops to true if we can operate in the push/pops mode.
    SmallVector<ProgramPoint, 4> doRestorePlacement(MCInst *BestPosSave,
                                                    unsigned CSR,
                                                    uint64_t TotalEstimatedWin);
  
    /// Checks whether using push and pops (instead of the longer load-store
    /// counterparts) is correct for reg \p CSR
    bool validatePushPopsMode(unsigned CSR, MCInst *BestPosSave,
                              int64_t SaveOffset);
  
    /// Adjust restore locations to the correct SP offset if we are using POPs
    /// instead of random-access load instructions.
    SmallVector<ProgramPoint, 4>
    fixPopsPlacements(const SmallVector<ProgramPoint, 4> &RestorePoints,
                      int64_t SaveOffset, unsigned CSR);
  
    /// When moving spills, mark all old spill locations to be deleted
    void scheduleOldSaveRestoresRemoval(unsigned CSR, bool UsePushPops);
    /// Return true if \p Inst uses reg \p CSR
    bool doesInstUsesCSR(const MCInst &Inst, uint16_t CSR);
    /// When moving spills, mark all new spill locations for insertion
    void
    scheduleSaveRestoreInsertions(unsigned CSR, MCInst *BestPosSave,
                                  SmallVector<ProgramPoint, 4> &RestorePoints,
                                  bool UsePushPops);
  
    /// Coordinate the replacement of callee-saved spills from their original
    /// place (at prologue and epilogues) to colder basic blocks as determined
    /// by computeSaveLocations().
    void moveSaveRestores();
  
    /// Compare multiple basic blocks created by splitting critical edges. If they
    /// have the same contents and successor, fold them into one.
    bool foldIdenticalSplitEdges();
  
    /// After the spill locations for reg \p CSR has been moved and all affected
    /// CFI has been removed, insert new updated CFI information for these
    /// locations.
    void insertUpdatedCFI(unsigned CSR, int SPValPush, int SPValPop);
  
    /// In case the function anchors the CFA reg as SP and we inserted pushes/pops
    /// insert def_cfa_offsets at appropriate places (and delete old
    /// def_cfa_offsets)
    void rebuildCFIForSP();
  
    /// Rebuild all CFI for affected Callee-Saved Registers.
    void rebuildCFI();
  
    /// Create a load-store instruction (depending on the contents of \p FIE).
    /// If \p CreatePushOrPop is true, create a push/pop instead. Current SP/FP
    /// values, as determined by StackPointerTracking, should be informed via
    /// \p SPVal and \p FPVal in order to emit the correct offset form SP/FP.
    MCInst createStackAccess(int SPVal, int FPVal, const FrameIndexEntry &FIE,
                             bool CreatePushOrPop);
  
    /// Update the CFI referenced by \p Inst with \p NewOffset, if the CFI has
    /// an offset.
    void updateCFIInstOffset(MCInst &Inst, int64_t NewOffset);
  
    /// Insert any CFI that should be attached to a register spill save/restore.
    BBIterTy insertCFIsForPushOrPop(BinaryBasicBlock &BB, BBIterTy Pos,
                                    unsigned Reg, bool isPush, int Sz,
                                    int64_t NewOffset);
  
    /// Auxiliary function to processInsertionsList, adding a new instruction
    /// before \p InsertionPoint as requested by \p Item. Return an updated
    /// InsertionPoint for other instructions that need to be inserted at the same
    /// original location, since this insertion may have invalidated the previous
    /// location.
    BBIterTy processInsertion(BBIterTy InsertionPoint, BinaryBasicBlock *CurBB,
                              const WorklistItem &Item, int64_t SPVal,
                              int64_t FPVal);
  
    /// Auxiliary function to processInsertions(), helping perform all the
    /// insertion tasks in the todo list associated with a single insertion point.
    /// Return true if at least one insertion was performed.
    BBIterTy processInsertionsList(BBIterTy InsertionPoint,
                                   BinaryBasicBlock *CurBB,
                                   std::vector<WorklistItem> &TodoList,
                                   int64_t SPVal, int64_t FPVal);
  
    /// Apply all insertion todo tasks regarding insertion of new stores/loads or
    /// push/pops at annotated points. Return false if the entire function had
    /// no todo tasks annotation and this pass has nothing to do.
    bool processInsertions();
  
    /// Apply all deletion todo tasks (or tasks to change a push/pop to a memory
    /// access no-op)
    void processDeletions();
  
  public:
    ShrinkWrapping(const FrameAnalysis &FA, const BinaryContext &BC,
                    BinaryFunction &BF, DataflowInfoManager &Info)
        : FA(FA), BC(BC), BF(BF), Info(Info), SLM(FA, BC, BF, Info),
          CSA(FA, BC, BF, Info) {}
  
    ~ShrinkWrapping() {
      for (auto &BB : BF) {
        for (auto &Inst : BB) {
          BC.MIB->removeAnnotation(Inst, getAnnotationIndex());
        }
      }
    }
  
    bool perform();
  
    static void printStats();
  };
  
  } // end namespace bolt
  } // end namespace llvm
  
  #endif
  