From 36fdfaba191ce9e7f951d994379a0827c7b08ffe Mon Sep 17 00:00:00 2001
From: Nikita Popov <npopov@redhat.com>
Date: Thu, 17 Feb 2022 11:59:04 +0100
Subject: [PATCH] [RelLookupTableConverter] Ensure that GV, GEP and load types
 match

This code could be generalized to be type-independent, but for now
just ensure that the same type constraints are enforced with opaque
pointers as with typed pointers.
---
 .../Utils/RelLookupTableConverter.cpp         |  6 ++-
 .../RelLookupTableConverter/X86/opaque-ptr.ll | 53 +++++++++++++++++++
 2 files changed, 57 insertions(+), 2 deletions(-)
 create mode 100644 llvm/test/Transforms/RelLookupTableConverter/X86/opaque-ptr.ll

diff --git a/llvm/lib/Transforms/Utils/RelLookupTableConverter.cpp b/llvm/lib/Transforms/Utils/RelLookupTableConverter.cpp
index 65207056a3f4..ad1101540755 100644
--- a/llvm/lib/Transforms/Utils/RelLookupTableConverter.cpp
+++ b/llvm/lib/Transforms/Utils/RelLookupTableConverter.cpp
@@ -38,11 +38,13 @@ static bool shouldConvertToRelLookupTable(Module &M, GlobalVariable &GV) {
 
   GetElementPtrInst *GEP =
       dyn_cast<GetElementPtrInst>(GV.use_begin()->getUser());
-  if (!GEP || !GEP->hasOneUse())
+  if (!GEP || !GEP->hasOneUse() ||
+      GV.getValueType() != GEP->getSourceElementType())
     return false;
 
   LoadInst *Load = dyn_cast<LoadInst>(GEP->use_begin()->getUser());
-  if (!Load || !Load->hasOneUse())
+  if (!Load || !Load->hasOneUse() ||
+      Load->getType() != GEP->getResultElementType())
     return false;
 
   // If the original lookup table does not have local linkage and is
diff --git a/llvm/test/Transforms/RelLookupTableConverter/X86/opaque-ptr.ll b/llvm/test/Transforms/RelLookupTableConverter/X86/opaque-ptr.ll
new file mode 100644
index 000000000000..bed4fc6f5ba7
--- /dev/null
+++ b/llvm/test/Transforms/RelLookupTableConverter/X86/opaque-ptr.ll
@@ -0,0 +1,53 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; REQUIRES: x86-registered-target
+; RUN: opt < %s -passes=rel-lookup-table-converter -relocation-model=pic -opaque-pointers -S | FileCheck %s
+target datalayout = "e-p:64:64:64-i1:8:8-i8:8:8-i16:16:16-i32:32:32-i64:64:64-f32:32:32-f64:64:64-v64:64:64-v128:128:128-a0:0:64-s0:64:64-f80:128:128-n8:16:32:64-S128"
+target triple = "x86_64-unknown-linux-gnu"
+
+@a = internal constant i32 0, align 4
+@b = internal constant i32 0, align 4
+@c = internal constant i32 0, align 4
+
+@table1 = private unnamed_addr constant [3 x ptr] [ptr @a, ptr @b, ptr @c], align 8
+@table2 = private unnamed_addr constant [3 x ptr] [ptr @a, ptr @b, ptr @c], align 8
+@table3 = private unnamed_addr constant [3 x ptr] [ptr @a, ptr @b, ptr @c], align 8
+
+define ptr @test(i32 %cond) {
+; CHECK-LABEL: @test(
+; CHECK-NEXT:    [[RELTABLE_SHIFT:%.*]] = shl i32 [[COND:%.*]], 2
+; CHECK-NEXT:    [[RELTABLE_INTRINSIC:%.*]] = call ptr @llvm.load.relative.i32(ptr @reltable.test, i32 [[RELTABLE_SHIFT]])
+; CHECK-NEXT:    ret ptr [[RELTABLE_INTRINSIC]]
+;
+  %switch.gep = getelementptr inbounds [3 x ptr], ptr @table1, i32 0, i32 %cond
+  %switch.load = load ptr, ptr %switch.gep, align 8
+  ret ptr %switch.load
+}
+
+define i32 @test_different_load_type(i32 %cond) {
+; CHECK-LABEL: @test_different_load_type(
+; CHECK-NEXT:    [[SWITCH_GEP:%.*]] = getelementptr inbounds [3 x ptr], ptr @table2, i32 0, i32 [[COND:%.*]]
+; CHECK-NEXT:    [[SWITCH_LOAD:%.*]] = load i32, ptr [[SWITCH_GEP]], align 8
+; CHECK-NEXT:    ret i32 [[SWITCH_LOAD]]
+;
+  %switch.gep = getelementptr inbounds [3 x ptr], ptr @table2, i32 0, i32 %cond
+  %switch.load = load i32, ptr %switch.gep, align 8
+  ret i32 %switch.load
+}
+
+define i8 @test_different_gep_type(i32 %cond) {
+; CHECK-LABEL: @test_different_gep_type(
+; CHECK-NEXT:    [[SWITCH_GEP:%.*]] = getelementptr inbounds [3 x i8], ptr @table3, i32 0, i32 [[COND:%.*]]
+; CHECK-NEXT:    [[SWITCH_LOAD:%.*]] = load i8, ptr [[SWITCH_GEP]], align 8
+; CHECK-NEXT:    ret i8 [[SWITCH_LOAD]]
+;
+  %switch.gep = getelementptr inbounds [3 x i8], ptr @table3, i32 0, i32 %cond
+  %switch.load = load i8, ptr %switch.gep, align 8
+  ret i8 %switch.load
+}
+
+!llvm.module.flags = !{!0, !1}
+!0 = !{i32 7, !"PIC Level", i32 2}
+!1 = !{i32 1, !"Code Model", i32 1}
+!4 = !{!"any pointer", !5, i64 0}
+!5 = !{!"omnipotent char", !6, i64 0}
+!6 = !{!"Simple C/C++ TBAA"}