Propagate flags properly in indirectCallTransformer (#107790)

dotnet · Sep 13, 2024 · 2267228 · 2267228
1 parent edf1c66
commit 2267228
Showing 1 changed file with 13 additions and 10 deletions.
diff --git a/src/coreclr/jit/indirectcalltransformer.cpp b/src/coreclr/jit/indirectcalltransformer.cpp
@@ -223,15 +223,20 @@ class IndirectCallTransformer
  // and insert in into the basic block list.
  //
  // Arguments:
- // jumpKind - jump kind for the new basic block
+ // jumpKind  - jump kind for the new basic block
  // insertAfter - basic block, after which compiler has to insert the new one.
+ // flagsSource - basic block to copy BBF_SPLIT_GAINED flags from
  //
  // Return Value:
  // new basic block.
- BasicBlock* CreateAndInsertBasicBlock(BBKinds jumpKind, BasicBlock* insertAfter)
+ BasicBlock* CreateAndInsertBasicBlock(BBKinds jumpKind, BasicBlock* insertAfter, BasicBlock* flagsSource)
  {
  BasicBlock* block = compiler->fgNewBBafter(jumpKind, insertAfter, true);
  block->SetFlags(BBF_IMPORTED);
+ if (flagsSource != nullptr)
+ {
+ block->CopyFlags(flagsSource, BBF_SPLIT_GAINED);
+ }
  return block;
  }
 
@@ -380,7 +385,7 @@ class IndirectCallTransformer
  {
  assert(checkIdx == 0);
 
- checkBlock = CreateAndInsertBasicBlock(BBJ_ALWAYS, currBlock);
+ checkBlock = CreateAndInsertBasicBlock(BBJ_ALWAYS, currBlock, currBlock);
  GenTree* fatPointerMask = new (compiler, GT_CNS_INT) GenTreeIntCon(TYP_I_IMPL, FAT_POINTER_MASK);
  GenTree* fptrAddressCopy = compiler->gtCloneExpr(fptrAddress);
  GenTree* fatPointerAnd = compiler->gtNewOperNode(GT_AND, TYP_I_IMPL, fptrAddressCopy, fatPointerMask);
@@ -398,7 +403,7 @@ class IndirectCallTransformer
  virtual void CreateThen(uint8_t checkIdx)
  {
  assert(remainderBlock != nullptr);
- thenBlock = CreateAndInsertBasicBlock(BBJ_ALWAYS, checkBlock);
+ thenBlock = CreateAndInsertBasicBlock(BBJ_ALWAYS, checkBlock, currBlock);
  Statement* copyOfOriginalStmt = compiler->gtCloneStmt(stmt);
  compiler->fgInsertStmtAtEnd(thenBlock, copyOfOriginalStmt);
  }
@@ -408,7 +413,7 @@ class IndirectCallTransformer
  //
  virtual void CreateElse()
  {
- elseBlock = CreateAndInsertBasicBlock(BBJ_ALWAYS, thenBlock);
+ elseBlock = CreateAndInsertBasicBlock(BBJ_ALWAYS, thenBlock, currBlock);
 
  GenTree* fixedFptrAddress = GetFixedFptrAddress();
  GenTree* actualCallAddress = compiler->gtNewIndir(pointerType, fixedFptrAddress);
@@ -605,7 +610,7 @@ class IndirectCallTransformer
  // In case of multiple checks, append to the previous thenBlock block
  // (Set jump target of new checkBlock in CreateThen())
  BasicBlock* prevCheckBlock = checkBlock;
- checkBlock = CreateAndInsertBasicBlock(BBJ_ALWAYS, thenBlock);
+ checkBlock = CreateAndInsertBasicBlock(BBJ_ALWAYS, thenBlock, currBlock);
  checkFallsThrough = false;
 
  // We computed the "then" likelihood in CreateThen, so we
@@ -1059,8 +1064,7 @@ class IndirectCallTransformer
 
  // thenBlock always jumps to remainderBlock
  //
- thenBlock = CreateAndInsertBasicBlock(BBJ_ALWAYS, checkBlock);
- thenBlock->CopyFlags(currBlock, BBF_SPLIT_GAINED);
+ thenBlock = CreateAndInsertBasicBlock(BBJ_ALWAYS, checkBlock, currBlock);
  thenBlock->inheritWeight(checkBlock);
  thenBlock->scaleBBWeight(adjustedThenLikelihood);
  FlowEdge* const thenRemainderEdge = compiler->fgAddRefPred(remainderBlock, thenBlock);
@@ -1091,8 +1095,7 @@ class IndirectCallTransformer
  //
  virtual void CreateElse()
  {
- elseBlock = CreateAndInsertBasicBlock(BBJ_ALWAYS, thenBlock);
- elseBlock->CopyFlags(currBlock, BBF_SPLIT_GAINED);
+ elseBlock = CreateAndInsertBasicBlock(BBJ_ALWAYS, thenBlock, currBlock);
 
  // We computed the "then" likelihood in CreateThen, so we
  // just use that to figure out the "else" likelihood.