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

[Nikita Popov via llvm-commits]

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

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.

>From 40ffd0df6c033b5255cae3ce8644ccb7e15f80c2 Mon Sep 17 00:00:00 2001
From: Nikita Popov <npopov at redhat.com>
Date: Mon, 6 Jul 2026 12:52:38 +0200
Subject: [PATCH] [JumpTableToSwitch] Bail out on function type mismatch

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 if the
type doesn't match we won't be able to inline anyway, so the
transform is pointless. Instead just bail out in this case.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply at anthropic.com>
---
 .../Transforms/Scalar/JumpTableToSwitch.cpp   |  8 ++--
 .../JumpTableToSwitch/callee-type-mismatch.ll | 48 +++++++++++++++++++
 2 files changed, 53 insertions(+), 3 deletions(-)
 create mode 100644 llvm/test/Transforms/JumpTableToSwitch/callee-type-mismatch.ll

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
+}