[llvm] [SimplifyCFG] Convert switch to cmp/select sequence (PR #82795)

[Thomas Symalla via llvm-commits]

================
@@ -6166,6 +6166,279 @@ static bool trySwitchToSelect(SwitchInst *SI, IRBuilder<> &Builder,
   return true;
 }
 
+// The first field contains the value that the switch produces when a certain
+// case group is selected, and the second field is a vector containing the
+// cases composing the case group.
+using SwitchCaseResultVectorTy2 =
+    SmallVector<std::pair<Value *, SmallVector<Value *, 4>>, 2>;
+
+using PHINodeToCaseEntryValueMapTy = std::map<
+    llvm::PHINode *,
+    llvm::SmallVector<std::pair<llvm::ConstantInt *, llvm::Value *>, 4>>;
+
+// The first field contains the phi node that generates a result of the switch
+// and the second field contains the value generated for a certain case in the
+// switch for that PHI.
+using SwitchCaseResultsTy2 = SmallVector<std::pair<PHINode *, Value *>, 4>;
+
+static void mapCaseToResultWithMap(SwitchCaseResultsTy2 &Results,
+                                   PHINodeToCaseEntryValueMapTy &Map,
+                                   ConstantInt *CaseVal,
+                                   SwitchCaseResultVectorTy2 &UniqueResults,
+                                   Value *Result) {
+  // fill in PHINodeToCaseEntryValueMapTy &Map
+  for (const auto &Pair : Results) {
+    PHINode *PHI = Pair.first;
+    Value *ReturnValue = Pair.second;
+    Map[PHI].emplace_back(std::make_pair(CaseVal, ReturnValue));
+  }
+
+  // fill in SwitchCaseResultVectorTy2 &UniqueResults
+  for (auto &I : UniqueResults) {
+    if (I.first == Result) {
+      I.second.push_back(CaseVal);
+      return;
+    }
+  }
+  UniqueResults.push_back(
+      std::make_pair(Result, SmallVector<Value *, 4>(1, CaseVal)));
+}
+
+static bool getCaseResultsWithoutConstants(
+    SwitchInst *SI, BasicBlock *CaseDest, BasicBlock **CommonDest,
+    SmallVectorImpl<std::pair<PHINode *, Value *>> &Res,
+    bool IsDefault = false) {
+  BasicBlock *Pred = SI->getParent();
+  int NumOfInsts = 0;
+  // Check if there is only one instruction per block, excepet from default
+  // block
+  if (!IsDefault) {
+    for (Instruction &I : CaseDest->instructionsWithoutDebug(false)) {
+
+      if (I.isTerminator()) {
+        // If the terminator is a simple branch, continue to the next block.
+        if (I.getNumSuccessors() != 1 || I.isSpecialTerminator())
+          return false;
+        Pred = CaseDest;
+        CaseDest = I.getSuccessor(0);
+      } else {
+        // samo jedna instrukcija se hendluje po bloku
+        if (++NumOfInsts > 1)
+          return false;
+        for (auto &Use : I.uses()) {
+          User *User = Use.getUser();
+          if (Instruction *I = dyn_cast<Instruction>(User))
+            if (I->getParent() == CaseDest)
+              continue;
+          if (PHINode *Phi = dyn_cast<PHINode>(User))
+            if (Phi->getIncomingBlock(Use) == CaseDest)
+              continue;
+          return false;
+        }
+      }
+    }
+  }
+
+  // If we did not have a CommonDest before, use the current one.
+  if (!*CommonDest)
+    *CommonDest = CaseDest;
+  // If the destination isn't the common one, abort.
+  if (CaseDest != *CommonDest)
+    return false;
+
+  // Get the values for this case from phi nodes in the destination block.
+  for (PHINode &PHI : (*CommonDest)->phis()) {
+    int Idx = PHI.getBasicBlockIndex(Pred);
+    if (Idx == -1)
+      continue;
+
+    Value *Val = PHI.getIncomingValue(Idx);
+
+    if (!Val)
+      return false;
+
+    Res.push_back(std::make_pair(&PHI, Val));
+  }
+
+  return Res.size() > 0;
+}
+
+static bool initializeuniqueCasesWithoutConstants(
+    SwitchInst *SI, PHINode *&PHI, BasicBlock *&CommonDest,
+    PHINodeToCaseEntryValueMapTy &Map, SwitchCaseResultVectorTy2 &UniqueResults,
+    Value *&DefaultResult, const DataLayout &DL, const TargetTransformInfo &TTI,
+    uintptr_t MaxUniqueResults) {
+  for (const auto &Case : SI->cases()) {
+    SwitchCaseResultsTy2 Results;
+    ConstantInt *CaseVal = Case.getCaseValue();
+    if (!getCaseResultsWithoutConstants(SI, Case.getCaseSuccessor(),
+                                        &CommonDest, Results))
+      return false;
+
+    mapCaseToResultWithMap(Results, Map, CaseVal, UniqueResults,
+                           Results.begin()->second);
+
+    // Check the PHI consistency.
+    if (!PHI)
+      PHI = Results[0].first;
+    else if (PHI != Results[0].first)
+      return false;
+  }
+
+  // Find the default result value.
+  SmallVector<std::pair<PHINode *, Value *>, 1> DefaultResults;
+  BasicBlock *DefaultDest = SI->getDefaultDest();
+  getCaseResultsWithoutConstants(SI, SI->getDefaultDest(), &CommonDest,
+                                 DefaultResults, true);
+  DefaultResult = DefaultResults.begin()->second;
+
+  if ((!DefaultResult &&
+       !isa<UnreachableInst>(DefaultDest->getFirstNonPHIOrDbg())))
+    return false;
+
+  return true;
+}
+
+Value *createSelectChain(Value *Condition, Value *DefaultResult,
+                         const SwitchCaseResultVectorTy2 &ResultVector,
+                         unsigned StartIndex, PHINode *PHI, SwitchInst *SI,
+                         IRBuilder<> &Builder, DomTreeUpdater *DTU) {
+
+  if (StartIndex >= ResultVector.size() && DefaultResult)
+    return DefaultResult;
+
+  if (StartIndex >= ResultVector.size() && !DefaultResult)
+    return nullptr;
+
+  Value *CurrentCase = ResultVector[StartIndex].second[0];
+  // check if the ReturnValue is a constant
+  if (Constant *C = dyn_cast<Constant>(ResultVector[StartIndex].first)) {
+    Value *ValueCompare =
+        Builder.CreateICmpEQ(Condition, CurrentCase, "switch.selectcmp");
+    Value *NextSelect =
+        createSelectChain(Condition, DefaultResult, ResultVector,
+                          StartIndex + 1, PHI, SI, Builder, DTU);
+    Value *SelectedValue =
+        Builder.CreateSelect(ValueCompare, ResultVector[StartIndex].first,
+                             NextSelect, "switch.select");
+
+    if (DTU) {
+      DTU->applyUpdates(
+          {{DominatorTree::Delete, SI->getParent(), SI->getDefaultDest()}});
+    }
+
+    return SelectedValue;
+
+  } else {
+    // Move to the next case if CurrentCase is not a constant
+    return createSelectChain(Condition, DefaultResult, ResultVector,
+                             StartIndex + 1, PHI, SI, Builder, DTU);
+  }
+}
+
+static Value *foldSwitchToSelect2(const SwitchCaseResultVectorTy2 &ResultVector,
+                                  Value *DefaultResult, Value *Condition,
+                                  PHINode *&PHI, SwitchInst *SI,
+                                  IRBuilder<> &Builder, DomTreeUpdater *DTU) {
+  if (ResultVector.size() > 2) {
+    Value *FinalSelect = createSelectChain(
+        Condition, DefaultResult, ResultVector, 0, PHI, SI, Builder, DTU);
+
+    if (FinalSelect)
+      return FinalSelect;
+  }
+
+  return nullptr;
+}
+
+/// Helper function that will try to convert switch with cmp/select sequence
+static bool convertSwitchToCmpAndSelect(SwitchInst *SI, IRBuilder<> &Builder,
+                                        DomTreeUpdater *DTU,
+                                        const DataLayout &DL,
+                                        const TargetTransformInfo &TTI) {
+  Value *const Cond = SI->getCondition();
+  PHINode *PHI = nullptr;
+  BasicBlock *CommonDest = nullptr;
+  Value *DefaultResult;
+  SwitchCaseResultVectorTy2 UniqueResults;
+  PHINodeToCaseEntryValueMapTy Map;
+
+  if (!initializeuniqueCasesWithoutConstants(
----------------
tsymalla wrote:

```suggestion
  if (!initializeUniqueCasesWithoutConstants(
```

https://github.com/llvm/llvm-project/pull/82795