[WinEH] Hoist state stores from successors

[majnemer]

Summary

If we know that all of our successors want to be in the exact same
state, it makes sense to hoist the state transition into their common
predecessor.

[rnk]

lgtm

lib/Target/X86/X86WinEHState.cpp

if (auto *II = dyn_cast<InvokeInst>(CS.getInstruction())) {
     return FuncInfo.InvokeStateMap[II];
   }
   // Possibly throwing call instructions have no actions to take after
   // an unwind. Ensure they are in the -1 state.
   return getBaseStateForBB(BlockColors, FuncInfo, CS.getParent());
 }
 
 // Calculate the intersection of all the FinalStates for a BasicBlock's
-// predecessor.
+// predecessors.
 static int getPredState(DenseMap<BasicBlock *, int> &FinalStates, Function &F,
                         int ParentBaseState, BasicBlock *BB) {
   // The entry block has no predecessors but we know that the prologue always
   // sets us up with a fixed state.
   if (&F.getEntryBlock() == BB)
     return ParentBaseState;
 
   // This is an EH Pad, conservatively report this basic block as overdefined.
for (BasicBlock *PredBB : predecessors(BB)) {
     // conservatively report this basic block as overdefined.
     if (CommonState != PredState)
       return OverdefinedState;
   }
 
   return CommonState;
 }
 
+// Calculate the intersection of all the InitialStates for a BasicBlock's
+// successors.
+static int getSuccState(DenseMap<BasicBlock *, int> &InitialStates, Function &F,
+                        int ParentBaseState, BasicBlock *BB) {
+  // This block rejoins normal control flow,
+  // conservatively report this basic block as overdefined.
+  if (isa<CatchReturnInst>(BB->getTerminator()))
+    return OverdefinedState;
+
+  int CommonState = OverdefinedState;
+  for (BasicBlock *SuccBB : successors(BB)) {
+    // We didn't manage to get a state for one of these predecessors,
+    // conservatively report this basic block as overdefined.
+    auto SuccStartState = InitialStates.find(SuccBB);
+    if (SuccStartState == InitialStates.end())
+      return OverdefinedState;
+
+    // This is an EH Pad, conservatively report this basic block as overdefined.
+    if (SuccBB->isEHPad())
+      return OverdefinedState;
+
+    int SuccState = SuccStartState->second;
+    assert(SuccState != OverdefinedState &&
+           "overdefined BBs shouldn't be in FinalStates");
+    if (CommonState == OverdefinedState)
+      CommonState = SuccState;
+
+    // At least two successors have different InitialStates,
+    // conservatively report this basic block as overdefined.
+    if (CommonState != SuccState)
+      return OverdefinedState;
+  }
+
+  return CommonState;
+}
+
 static bool isStateStoreNeeded(EHPersonality Personality, CallSite CS) {
   if (!CS)
     return false;
 
   if (isAsynchronousEHPersonality(Personality))
     return !CS.doesNotAccessMemory();
 
   return !CS.doesNotThrow();
while (!Worklist.empty()) {
     // We successfully inferred this BasicBlock's state via it's predecessors;
     // enqueue it's successors to see if we can infer their states.
     InitialStates.insert({BB, PredState});
     FinalStates.insert({BB, PredState});
     for (BasicBlock *SuccBB : successors(BB))
       Worklist.push_back(SuccBB);
   }
 
+  // Try to hoist stores from successors.
+  for (BasicBlock *BB : RPOT) {
+    int SuccState = getSuccState(InitialStates, F, ParentBaseState, BB);
+    if (SuccState == OverdefinedState)
+      continue;
+
+    // Update our FinalState to reflect the common InitialState of our
+    // successors.
+    FinalStates.insert({BB, SuccState});
+  }
+
   // Finally, insert state stores before call-sites which transition us to a new
   // state.
   for (BasicBlock *BB : RPOT) {
     auto &BBColors = BlockColors[BB];
     BasicBlock *FuncletEntryBB = BBColors.front();
     if (isa<CleanupPadInst>(FuncletEntryBB->getFirstNonPHI()))
       continue;
 
     int PrevState = getPredState(FinalStates, F, ParentBaseState, BB);
     DEBUG(dbgs() << "X86WinEHState: " << BB->getName()
                  << " PrevState=" << PrevState << '\n');
 
     for (Instruction &I : *BB) {
       CallSite CS(&I);
       if (!isStateStoreNeeded(Personality, CS))
         continue;
 
       int State = getStateForCallSite(BlockColors, FuncInfo, CS);
       if (State != PrevState)
         insertStateNumberStore(&I, State);
       PrevState = State;
     }
+
+    // We might have hoisted a state store into this block, emit it now.
+    auto EndState = FinalStates.find(BB);
+    if (EndState != FinalStates.end())
+      if (EndState->second != PrevState)
+        insertStateNumberStore(BB->getTerminator(), EndState->second);
   }
 }
 
 void WinEHStatePass::insertStateNumberStore(Instruction *IP, int State) {
   IRBuilder<> Builder(IP);
   Value *StateField =
       Builder.CreateStructGEP(nullptr, RegNode, StateFieldIndex);
   Builder.CreateStore(Builder.getInt32(State), StateField);
 }

test/CodeGen/X86/win-catchpad.ll

 ; X86-NEXT: addl $12, %ebp
 ; X86: jmp [[contbb]]
 
 ; X86: "?catch$[[catchbb:[0-9]+]]@?0?branch_to_normal_dest@4HA":
 ; X86: LBB1_[[catchbb]]: # %catch{{$}}
 ; X86: pushl %ebp
 ; X86: subl $8, %esp
 ; X86: addl $12, %ebp
-; X86: LBB1_[[loopbb:[0-9]+]]: # %loop
 ; X86: movl    $1, -16(%ebp)
+; X86: LBB1_[[loopbb:[0-9]+]]: # %loop
 ; X86: calll   _getbool
 ; X86: testb   $1, %al
 ; X86: jne LBB1_[[loopbb]]
 ; X86: # %catch.done
 ; X86-NEXT: movl $[[restorebb]], %eax
 ; X86-NEXT: addl $8, %esp
 ; X86-NEXT: popl %ebp
 ; X86-NEXT: retl