Add visitGEPOfAlloc to restore the visitGEPOfBitcast part of the behavior of non-opaque pointers.

llvm/lib/Transforms/InstCombine/InstCombineInternal.h

@@ -155,6 +155,7 @@ class LLVM_LIBRARY_VISIBILITY InstCombinerImpl final
   Instruction *visitPHINode(PHINode &PN);
   Instruction *visitGetElementPtrInst(GetElementPtrInst &GEP);
   Instruction *visitGEPOfGEP(GetElementPtrInst &GEP, GEPOperator *Src);
+  Instruction *visitGEPOfAlloc(GetElementPtrInst &GEP, AllocaInst *AI);
   Instruction *visitAllocaInst(AllocaInst &AI);
   Instruction *visitAllocSite(Instruction &FI);
   Instruction *visitFree(CallInst &FI, Value *FreedOp);

llvm/lib/Transforms/InstCombine/InstructionCombining.cpp

@@ -1975,6 +1975,48 @@ static Instruction *foldSelectGEP(GetElementPtrInst &GEP,
   return SelectInst::Create(Cond, NewTrueC, NewFalseC, "", nullptr, Sel);
 }
 
+static Type *findElementAtOffset(Type *PtrTy, Type *Ty, int64_t IntOffset,
+                                 SmallVectorImpl<Value *> &NewIndices,
+                                 const DataLayout &DL) {
+  if (!Ty->isSized())
+    return nullptr;
+  APInt Offset(DL.getIndexTypeSizeInBits(PtrTy), IntOffset);
+  SmallVector<APInt> Indices = DL.getGEPIndicesForOffset(Ty, Offset);
+  if (!Offset.isZero())
+    return nullptr;
+
+  for (const APInt &Index : Indices)
+    NewIndices.push_back(ConstantInt::get(Ty->getContext(), Index));
+  return Ty;
+}
+
+// See if we can simplify:
+//   X = alloc %Type
+//   Y = gep X, <...constant indices...>
+// into a gep of the original struct. This is important for SROA and alias
+// analysis of unions.
+Instruction *InstCombinerImpl::visitGEPOfAlloc(GetElementPtrInst &GEP,
+                                               AllocaInst *AI) {
+  Type *SrcType = AI->getAllocatedType();
+  if (SrcType->isArrayTy() || GEP.getSourceElementType() == SrcType)
+    return nullptr;
+  unsigned OffsetBits = DL.getIndexTypeSizeInBits(GEP.getType());
+  APInt Offset(OffsetBits, 0);
+  if (GEP.accumulateConstantOffset(DL, Offset)) {
+    if (!Offset)
+      return nullptr;
+    // we need to find out if there is a field at Offset in 'A's type.
+    SmallVector<Value *, 8> NewIndices;
+    if (findElementAtOffset(GEP.getType(), SrcType, Offset.getSExtValue(),
+                            NewIndices, DL)) {
+      Value *NGEP = Builder.CreateGEP(SrcType, GEP.getOperand(0), NewIndices,
+                                      "", GEP.isInBounds());
+      return replaceInstUsesWith(GEP, NGEP);
+    }
+  }
+  return nullptr;
+}
+
 Instruction *InstCombinerImpl::visitGEPOfGEP(GetElementPtrInst &GEP,
                                              GEPOperator *Src) {
   // Combine Indices - If the source pointer to this getelementptr instruction
@@ -2334,6 +2376,10 @@ Instruction *InstCombinerImpl::visitGetElementPtrInst(GetElementPtrInst &GEP) {
   if (GEPType->isVectorTy())
     return nullptr;
 
+  if (auto *AI = dyn_cast<AllocaInst>(PtrOp))
+    if (Instruction *I = visitGEPOfAlloc(GEP, AI))
+      return I;
+
   if (!GEP.isInBounds()) {
     unsigned IdxWidth =
         DL.getIndexSizeInBits(PtrOp->getType()->getPointerAddressSpace());

llvm/test/Transforms/InstCombine/2009-01-08-AlignAlloca.ll

@@ -13,7 +13,7 @@ define i32 @bar(i64 %key_token2) nounwind {
 ; CHECK-NEXT:    [[IOSPEC:%.*]] = alloca [[STRUCT_KEY:%.*]], align 8
 ; CHECK-NEXT:    [[RET:%.*]] = alloca i32, align 4
 ; CHECK-NEXT:    store i32 0, ptr [[IOSPEC]], align 8
-; CHECK-NEXT:    [[TMP0:%.*]] = getelementptr inbounds { i32, i32 }, ptr [[IOSPEC]], i32 0, i32 1
+; CHECK-NEXT:    [[TMP0:%.*]] = getelementptr inbounds [[STRUCT_KEY]], ptr [[IOSPEC]], i32 0, i32 0, i32 1
 ; CHECK-NEXT:    store i32 0, ptr [[TMP0]], align 4
 ; CHECK-NEXT:    store i64 [[KEY_TOKEN2:%.*]], ptr [[IOSPEC]], align 8
 ; CHECK-NEXT:    [[TMP1:%.*]] = call i32 (...) @foo(ptr nonnull byval([[STRUCT_KEY]]) align 4 [[IOSPEC]], ptr nonnull [[RET]]) #[[ATTR0:[0-9]+]]

llvm/test/Transforms/InstCombine/restore-gep-original-struct.ll

@@ -0,0 +1,26 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 3
+; RUN: opt < %s -passes=instcombine -S | FileCheck %s
+%Zip = type { i64, i64, [1 x i8] }
+
+define i64 @foo(i64 %idx, ptr %v) {
+; CHECK-LABEL: define i64 @foo(
+; CHECK-SAME: i64 [[IDX:%.*]], ptr [[V:%.*]]) {
+; CHECK-NEXT:    [[Z1:%.*]] = alloca [[ZIP:%.*]], align 8
+; CHECK-NEXT:    [[SROA_IDX:%.*]] = getelementptr inbounds [[ZIP]], ptr [[Z1]], i64 0, i32 1
+; CHECK-NEXT:    store i64 32, ptr [[SROA_IDX]], align 8
+; CHECK-NEXT:    [[AF:%.*]] = getelementptr inbounds [[ZIP]], ptr [[Z1]], i64 0, i32 2, i64 [[IDX]]
+; CHECK-NEXT:    [[A:%.*]] = load i8, ptr [[AF]], align 8
+; CHECK-NEXT:    store i8 [[A]], ptr [[V]], align 8
+; CHECK-NEXT:    ret i64 32
+;
+  %z1 = alloca %Zip, align 8
+  %sroa_idx = getelementptr inbounds i8, ptr %z1, i64 8
+  store i64 32, ptr %sroa_idx, align 8
+  %offset = getelementptr inbounds %Zip, ptr %z1, i64 0, i32 1
+  %af = getelementptr inbounds %Zip, ptr %z1, i64 0, i32 2, i64 %idx
+  %a = load i8, ptr %af, align 8
+  store i8 %a, ptr %v, align 8
+  %len = load i64, ptr %offset, align 8
+  ret i64 %len
+}
+