[llvm] Add a pass to convert jump tables to switches. (PR #77709)

[Alexander Shaposhnikov via llvm-commits]

================
@@ -0,0 +1,205 @@
+//===- JumpTableToSwitch.cpp ----------------------------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/Transforms/Scalar/JumpTableToSwitch.h"
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/SmallSet.h"
+#include "llvm/Analysis/DomTreeUpdater.h"
+#include "llvm/Analysis/PostDominators.h"
+#include "llvm/Analysis/TargetLibraryInfo.h"
+#include "llvm/Analysis/TargetTransformInfo.h"
+#include "llvm/Analysis/ValueTracking.h"
+#include "llvm/IR/IRBuilder.h"
+#include "llvm/IR/IntrinsicInst.h"
+#include "llvm/IR/PatternMatch.h"
+#include "llvm/Support/CommandLine.h"
+#include "llvm/Support/Debug.h"
+#include "llvm/Transforms/Utils/BasicBlockUtils.h"
+#include "llvm/Transforms/Utils/Cloning.h"
+#include "llvm/Transforms/Utils/Local.h"
+
+using namespace llvm;
+using namespace PatternMatch;
+
+static cl::opt<unsigned>
+    JumpTableSizeThreshold("jump-table-to-switch-size-threshold", cl::Hidden,
+                           cl::desc("Only split jump tables with size less or "
+                                    "equal than JumpTableSizeThreshold."),
+                           cl::init(10));
+
+#define DEBUG_TYPE "jump-table-to-switch"
+
+static Constant *getElementWithGivenTypeAtOffset(Constant *C, const Type *Ty,
+                                                 uint64_t Offset, Module &M) {
+  if (Offset == 0 && C->getType() == Ty)
+    return C;
+  if (auto *CS = dyn_cast<ConstantStruct>(C)) {
+    const DataLayout &DL = M.getDataLayout();
+    const StructLayout *SL = DL.getStructLayout(CS->getType());
+    if (Offset >= SL->getSizeInBytes())
+      return nullptr;
+    const unsigned Op = SL->getElementContainingOffset(Offset);
+    const uint64_t AdjustedOffset = Offset - SL->getElementOffset(Op);
+    Constant *Element = cast<Constant>(CS->getOperand(Op));
+    return getElementWithGivenTypeAtOffset(Element, Ty, AdjustedOffset, M);
+  }
+  // TODO: add support for arrays.
+  return nullptr;
+}
+
+namespace {
+struct JumpTableTy {
+  Value *Index;
+  SmallVector<Function *, 5> Funcs;
+};
+} // anonymous namespace
+
+static std::optional<JumpTableTy> parseJumpTable(GetElementPtrInst *GEP) {
+  if (!GEP || !GEP->isInBounds())
+    return std::nullopt;
+  ArrayType *ArrayTy = dyn_cast<ArrayType>(GEP->getSourceElementType());
+  if (!ArrayTy || ArrayTy->getArrayNumElements() > JumpTableSizeThreshold)
+    return std::nullopt;
+  Constant *Ptr = dyn_cast<Constant>(GEP->getPointerOperand());
+  if (!Ptr)
+    return std::nullopt;
+
+  Function &F = *GEP->getParent()->getParent();
+  const DataLayout &DL = F.getParent()->getDataLayout();
+  const unsigned BitWidth =
+      DL.getIndexSizeInBits(GEP->getPointerAddressSpace());
+  MapVector<Value *, APInt> VariableOffsets;
+  APInt ConstantOffset(BitWidth, 0);
+  if (!GEP->collectOffset(DL, BitWidth, VariableOffsets, ConstantOffset))
+    return std::nullopt;
+  if (VariableOffsets.empty() || VariableOffsets.size() > 1)
+    return std::nullopt;
+
+  Module &M = *F.getParent();
+  unsigned Offset = ConstantOffset.getZExtValue();
+  // TODO: support more general patterns
+  // (see also TODO in getElementWithGivenTypeAtOffset).
+  if (Offset != 0)
+    return std::nullopt;
+  if (!Ptr->getNumOperands())
+    return std::nullopt;
+  Constant *ConstArray = getElementWithGivenTypeAtOffset(
+      cast<Constant>(Ptr->getOperand(0)), ArrayTy, Offset, M);
+  if (!ConstArray)
+    return std::nullopt;
+
+  JumpTableTy JumpTable;
+  JumpTable.Index = VariableOffsets.front().first;
----------------
alexander-shaposhnikov wrote:

Done

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