[llvm] [JumpTableToSwitch] Bail out on function type mismatch (PR #207700)

[via llvm-commits]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-llvm-transforms

Author: Nikita Popov (nikic)

<details>
<summary>Changes</summary>

If the function types of the functions in the jump table don't match the function type of the call, we currently get an assertion failure. We could avoid that by using a different API, but the usefulness of the transform in this case is questionable (e.g. we wouldn't be able to inline the call). For simplicity just bail out.

---
Full diff: https://github.com/llvm/llvm-project/pull/207700.diff


2 Files Affected:

- (modified) llvm/lib/Transforms/Scalar/JumpTableToSwitch.cpp (+5-3) 
- (added) llvm/test/Transforms/JumpTableToSwitch/callee-type-mismatch.ll (+48) 


``````````diff
diff --git a/llvm/lib/Transforms/Scalar/JumpTableToSwitch.cpp b/llvm/lib/Transforms/Scalar/JumpTableToSwitch.cpp
index d5696bc9ea956..bd4946ac5e69f 100644
--- a/llvm/lib/Transforms/Scalar/JumpTableToSwitch.cpp
+++ b/llvm/lib/Transforms/Scalar/JumpTableToSwitch.cpp
@@ -61,7 +61,8 @@ struct JumpTableTy {
 } // anonymous namespace
 
 static std::optional<JumpTableTy> parseJumpTable(GetElementPtrInst *GEP,
-                                                 PointerType *PtrTy) {
+                                                 PointerType *PtrTy,
+                                                 FunctionType *CallFTy) {
   Constant *Ptr = dyn_cast<Constant>(GEP->getPointerOperand());
   if (!Ptr)
     return std::nullopt;
@@ -101,7 +102,7 @@ static std::optional<JumpTableTy> parseJumpTable(GetElementPtrInst *GEP,
     Constant *C =
         ConstantFoldLoadFromConst(GV->getInitializer(), PtrTy, Offset, DL);
     auto *Func = dyn_cast_or_null<Function>(C);
-    if (!Func || Func->isDeclaration() ||
+    if (!Func || Func->isDeclaration() || Func->getFunctionType() != CallFTy ||
         Func->getInstructionCount() > FunctionSizeThreshold)
       return std::nullopt;
     JumpTable.Funcs.push_back(Func);
@@ -240,7 +241,8 @@ PreservedAnalyses JumpTableToSwitchPass::run(Function &F,
           continue;
         auto *PtrTy = dyn_cast<PointerType>(L->getType());
         assert(PtrTy && "call operand must be a pointer");
-        std::optional<JumpTableTy> JumpTable = parseJumpTable(GEP, PtrTy);
+        std::optional<JumpTableTy> JumpTable =
+            parseJumpTable(GEP, PtrTy, Call->getFunctionType());
         if (!JumpTable)
           continue;
         SplittedOutTail =
diff --git a/llvm/test/Transforms/JumpTableToSwitch/callee-type-mismatch.ll b/llvm/test/Transforms/JumpTableToSwitch/callee-type-mismatch.ll
new file mode 100644
index 0000000000000..04d3206268f38
--- /dev/null
+++ b/llvm/test/Transforms/JumpTableToSwitch/callee-type-mismatch.ll
@@ -0,0 +1,48 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 6
+; RUN: opt < %s -passes=jump-table-to-switch -S | FileCheck %s
+
+ at arr_argcount = constant [1 x ptr] [ptr @f_argcount]
+ at arr_rettype = constant [1 x ptr] [ptr @f_rettype]
+
+define i32 @f_argcount(i32 %x) {
+; CHECK-LABEL: define i32 @f_argcount(
+; CHECK-SAME: i32 [[X:%.*]]) {
+; CHECK-NEXT:    ret i32 [[X]]
+;
+  ret i32 %x
+}
+
+define i64 @f_rettype() {
+; CHECK-LABEL: define i64 @f_rettype() {
+; CHECK-NEXT:    ret i64 123
+;
+  ret i64 123
+}
+
+define i32 @argcount(i32 %idx) {
+; CHECK-LABEL: define i32 @argcount(
+; CHECK-SAME: i32 [[IDX:%.*]]) {
+; CHECK-NEXT:    [[GEP:%.*]] = getelementptr [1 x ptr], ptr @arr_argcount, i32 0, i32 [[IDX]]
+; CHECK-NEXT:    [[FP:%.*]] = load ptr, ptr [[GEP]], align 8
+; CHECK-NEXT:    [[R:%.*]] = call i32 [[FP]]()
+; CHECK-NEXT:    ret i32 [[R]]
+;
+  %gep = getelementptr [1 x ptr], ptr @arr_argcount, i32 0, i32 %idx
+  %fp = load ptr, ptr %gep
+  %r = call i32 %fp()
+  ret i32 %r
+}
+
+define i32 @rettype(i32 %idx) {
+; CHECK-LABEL: define i32 @rettype(
+; CHECK-SAME: i32 [[IDX:%.*]]) {
+; CHECK-NEXT:    [[GEP:%.*]] = getelementptr [1 x ptr], ptr @arr_rettype, i32 0, i32 [[IDX]]
+; CHECK-NEXT:    [[FP:%.*]] = load ptr, ptr [[GEP]], align 8
+; CHECK-NEXT:    [[R:%.*]] = call i32 [[FP]]()
+; CHECK-NEXT:    ret i32 [[R]]
+;
+  %gep = getelementptr [1 x ptr], ptr @arr_rettype, i32 0, i32 %idx
+  %fp = load ptr, ptr %gep
+  %r = call i32 %fp()
+  ret i32 %r
+}

``````````

</details>


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