From 9f779195d311c983031271d0243d6e6af988ce55 Mon Sep 17 00:00:00 2001
From: Nikita Popov <nikita.ppv@gmail.com>
Date: Mon, 21 Jun 2021 17:43:06 +0200
Subject: [PATCH] [OpaquePtr] Return opaque pointer from opaque pointer GEP

For a GEP on an opaque pointer, also return an opaque pointer (or
vector of opaque pointer) result.

This requires explicitly enumerating the GEP source element type,
because it is now no longer implicitly enumerated as part of either
the source or result pointer types.

Differential Revision: https://reviews.llvm.org/D104652
---
 llvm/include/llvm/IR/Instructions.h         |  8 ++++--
 llvm/lib/Bitcode/Writer/ValueEnumerator.cpp |  2 ++
 llvm/test/Assembler/opaque-ptr.ll           | 28 ++++++++++++++++-----
 3 files changed, 30 insertions(+), 8 deletions(-)

diff --git a/llvm/include/llvm/IR/Instructions.h b/llvm/include/llvm/IR/Instructions.h
index a133267ffe9c..5b334493eb41 100644
--- a/llvm/include/llvm/IR/Instructions.h
+++ b/llvm/include/llvm/IR/Instructions.h
@@ -1098,8 +1098,12 @@ public:
   /// instruction, which may be a vector of pointers.
   static Type *getGEPReturnType(Type *ElTy, Value *Ptr,
                                 ArrayRef<Value *> IdxList) {
-    Type *PtrTy = PointerType::get(checkGEPType(getIndexedType(ElTy, IdxList)),
-                                   Ptr->getType()->getPointerAddressSpace());
+    PointerType *OrigPtrTy = cast<PointerType>(Ptr->getType()->getScalarType());
+    unsigned AddrSpace = OrigPtrTy->getAddressSpace();
+    Type *ResultElemTy = checkGEPType(getIndexedType(ElTy, IdxList));
+    Type *PtrTy = OrigPtrTy->isOpaque()
+      ? PointerType::get(OrigPtrTy->getContext(), AddrSpace)
+      : PointerType::get(ResultElemTy, AddrSpace);
     // Vector GEP
     if (auto *PtrVTy = dyn_cast<VectorType>(Ptr->getType())) {
       ElementCount EltCount = PtrVTy->getElementCount();
diff --git a/llvm/lib/Bitcode/Writer/ValueEnumerator.cpp b/llvm/lib/Bitcode/Writer/ValueEnumerator.cpp
index 80c230da6a4b..abfdcd67524e 100644
--- a/llvm/lib/Bitcode/Writer/ValueEnumerator.cpp
+++ b/llvm/lib/Bitcode/Writer/ValueEnumerator.cpp
@@ -463,6 +463,8 @@ ValueEnumerator::ValueEnumerator(const Module &M,
         }
         if (auto *SVI = dyn_cast<ShuffleVectorInst>(&I))
           EnumerateType(SVI->getShuffleMaskForBitcode()->getType());
+        if (auto *GEP = dyn_cast<GetElementPtrInst>(&I))
+          EnumerateType(GEP->getSourceElementType());
         EnumerateType(I.getType());
         if (const auto *Call = dyn_cast<CallBase>(&I))
           EnumerateAttributes(Call->getAttributes());
diff --git a/llvm/test/Assembler/opaque-ptr.ll b/llvm/test/Assembler/opaque-ptr.ll
index b3b35631188e..3a71da2f3537 100644
--- a/llvm/test/Assembler/opaque-ptr.ll
+++ b/llvm/test/Assembler/opaque-ptr.ll
@@ -41,12 +41,28 @@ define void @store(ptr %a, i32 %i) {
     ret void
 }
 
-; CHECK: define void @gep(ptr %a)
-; CHECK:     %b = getelementptr i8, ptr %a, i32 2
-; CHECK:     ret void
-define void @gep(ptr %a) {
-    %b = getelementptr i8, ptr %a, i32 2
-    ret void
+; CHECK: define ptr @gep(ptr %a)
+; CHECK:     %res = getelementptr i8, ptr %a, i32 2
+; CHECK:     ret ptr %res
+define ptr @gep(ptr %a) {
+  %res = getelementptr i8, ptr %a, i32 2
+  ret ptr %res
+}
+
+; CHECK: define <2 x ptr> @gep_vec1(ptr %a)
+; CHECK:     %res = getelementptr i8, ptr %a, <2 x i32> <i32 1, i32 2>
+; CHECK:     ret <2 x ptr> %res
+define <2 x ptr> @gep_vec1(ptr %a) {
+  %res = getelementptr i8, ptr %a, <2 x i32> <i32 1, i32 2>
+  ret <2 x ptr> %res
+}
+
+; CHECK: define <2 x ptr> @gep_vec2(<2 x ptr> %a)
+; CHECK:     %res = getelementptr i8, <2 x ptr> %a, i32 2
+; CHECK:     ret <2 x ptr> %res
+define <2 x ptr> @gep_vec2(<2 x ptr> %a) {
+  %res = getelementptr i8, <2 x ptr> %a, i32 2
+  ret <2 x ptr> %res
 }
 
 ; CHECK: define void @cmpxchg(ptr %p, i32 %a, i32 %b)