[NFCI][InstCombine] Move store merging from `visitStoreInst()` into `visitUnconditionalBranchInst()`

[lebedev.ri]

Summary

As @nikic is pointing out in https://bugs.llvm.org/show_bug.cgi?id=46680#c5,
InstCombine should not have forward instruction scans,
so let's move this transform into the proper place.

This is pretty much NFCI.

[nikic]

llvm/lib/Transforms/InstCombine/InstructionCombining.cpp

2760

Can we just return &BI; here, and let the usual reprocessing deal with it?

[lebedev.ri]

If changed, just return BI, let the usual handling deal with it.

llvm/lib/Transforms/InstCombine/InstructionCombining.cpp

2760

.. or did you mean *always* return &BI; ?

[nikic]

llvm/lib/Transforms/InstCombine/InstructionCombining.cpp

2760

Yeah, I meant always returning. That will try rerunning this code and will either find the next store to merge and do so, or not.

[lebedev.ri]

Addressed comments.

[nikic]

LG

llvm/lib/Transforms/InstCombine/InstCombineInternal.h

public:
   Instruction *visitPHINode(PHINode &PN);
   Instruction *visitGetElementPtrInst(GetElementPtrInst &GEP);
   Instruction *visitAllocaInst(AllocaInst &AI);
   Instruction *visitAllocSite(Instruction &FI);
   Instruction *visitFree(CallInst &FI);
   Instruction *visitLoadInst(LoadInst &LI);
   Instruction *visitStoreInst(StoreInst &SI);
   Instruction *visitAtomicRMWInst(AtomicRMWInst &SI);
+  Instruction *visitUnconditionalBranchInst(BranchInst &BI);
   Instruction *visitBranchInst(BranchInst &BI);
   Instruction *visitFenceInst(FenceInst &FI);
   Instruction *visitSwitchInst(SwitchInst &SI);
   Instruction *visitReturnInst(ReturnInst &RI);
   Instruction *visitInsertValueInst(InsertValueInst &IV);
   Instruction *visitInsertElementInst(InsertElementInst &IE);
   Instruction *visitExtractElementInst(ExtractElementInst &EI);
   Instruction *visitShuffleVectorInst(ShuffleVectorInst &SVI);

llvm/lib/Transforms/InstCombine/InstCombineLoadStoreAlloca.cpp

if (!isa<UndefValue>(Val))
       return replaceOperand(SI, 0, UndefValue::get(Val->getType()));
     return nullptr;  // Do not modify these!
   }
 
   // store undef, Ptr -> noop
   if (isa<UndefValue>(Val))
     return eraseInstFromFunction(SI);
 
-  auto IsNoopInstrForStoreMerging = [](BasicBlock::iterator BBI) {
-    return isa<DbgInfoIntrinsic>(BBI) ||
-           (isa<BitCastInst>(BBI) && BBI->getType()->isPointerTy());
-  };
-
-  // If this store is the second-to-last instruction in the basic block
-  // (excluding debug info and bitcasts of pointers) and if the block ends with
-  // an unconditional branch, try to move the store to the successor block.
-  BBI = SI.getIterator();
-  do {
-    ++BBI;
-  } while (IsNoopInstrForStoreMerging(BBI));
-
-  if (BranchInst *BI = dyn_cast<BranchInst>(BBI))
-    if (BI->isUnconditional())
-      if (mergeStoreIntoSuccessor(SI)) {
-        // Okay, we've managed to do that. Now, let's see if now-second-to-last
-        // instruction is also a store that we can also sink.
-        BasicBlock::iterator FirstInstr = BBI->getParent()->begin();
-        do {
-          if (BBI != FirstInstr)
-            --BBI;
-        } while (BBI != FirstInstr && IsNoopInstrForStoreMerging(BBI));
-        if (StoreInst *PrevStore = dyn_cast<StoreInst>(BBI))
-          Worklist.add(PrevStore);
-        return nullptr;
-      }
-
   return nullptr;
 }
 
 /// Try to transform:
 ///   if () { *P = v1; } else { *P = v2 }
 /// or:
 ///   *P = v1; if () { *P = v2; }
 /// into a phi node with a store in the successor.
 bool InstCombiner::mergeStoreIntoSuccessor(StoreInst &SI) {
-  assert(SI.isUnordered() &&
-         "This code has not been audited for volatile or ordered store case.");
+  if (!SI.isUnordered())
+    return false; // This code has not been audited for volatile/ordered case.
 
   // Check if the successor block has exactly 2 incoming edges.
   BasicBlock *StoreBB = SI.getParent();
   BasicBlock *DestBB = StoreBB->getTerminator()->getSuccessor(0);
   if (!DestBB->hasNPredecessors(2))
     return false;
 
   // Capture the other block (the block that doesn't contain our store).

llvm/lib/Transforms/InstCombine/InstructionCombining.cpp

Instruction *InstCombiner::visitReturnInst(ReturnInst &RI) {
   KnownBits Known = computeKnownBits(ResultOp, 0, &RI);
   if (Known.isConstant())
     return replaceOperand(RI, 0,
         Constant::getIntegerValue(VTy, Known.getConstant()));
 
   return nullptr;
 }
 
+Instruction *InstCombiner::visitUnconditionalBranchInst(BranchInst &BI) {
+  assert(BI.isUnconditional() && "Only for unconditional branches.");
+
+  // If this store is the second-to-last instruction in the basic block
+  // (excluding debug info and bitcasts of pointers) and if the block ends with
+  // an unconditional branch, try to move the store to the successor block.
+
+  auto GetLastSinkableStore = [](BasicBlock::iterator BBI) {
+    auto IsNoopInstrForStoreMerging = [](BasicBlock::iterator BBI) {
+      return isa<DbgInfoIntrinsic>(BBI) ||
+             (isa<BitCastInst>(BBI) && BBI->getType()->isPointerTy());
+    };
+
+    BasicBlock::iterator FirstInstr = BBI->getParent()->begin();
+    do {
+      if (BBI != FirstInstr)
+        --BBI;
+    } while (BBI != FirstInstr && IsNoopInstrForStoreMerging(BBI));
+
+    return dyn_cast<StoreInst>(BBI);
+  };
+
+  if (StoreInst *SI = GetLastSinkableStore(BasicBlock::iterator(BI)))
+    if (mergeStoreIntoSuccessor(*SI)) {
+      // Okay, we've managed to do that. Now, let's see if now-second-to-last
+      // instruction is also a store that we can also sink.
+      if (StoreInst *SI = GetLastSinkableStore(BasicBlock::iterator(BI)))
+        return &BI; // Let's revisit this `br` once more then.
+    }
+
+  return nullptr;
+}
+
 Instruction *InstCombiner::visitBranchInst(BranchInst &BI) {
-  // Nothing to do about unconditional branches.
   if (BI.isUnconditional())
-    return nullptr;
+    return visitUnconditionalBranchInst(BI);
 
   // Change br (not X), label True, label False to: br X, label False, True
   Value *X = nullptr;
   if (match(&BI, m_Br(m_Not(m_Value(X)), m_BasicBlock(), m_BasicBlock())) &&
       !isa<Constant>(X)) {
     // Swap Destinations and condition...
     BI.swapSuccessors();
     return replaceOperand(BI, 0, X);