[ConstantFolding] Fix vector GEPs harder

[mkuper]

Summary

CastGEPIndices still handles Vector GEPs incorrectly - in some cases we can fall into an infinite recursion, because we compare the vector type to the scalar pointer type, and then try to cast a type to itself.

[chandlerc]

LGTM, nice find.

[davide]

LGTM.

lib/Analysis/ConstantFolding.cpp

Constant *CastGEPIndices(Type *SrcElemTy, ArrayRef<Constant *> Ops,
   Type *IntPtrScalarTy = IntPtrTy->getScalarType();
 
   bool Any = false;
   SmallVector<Constant*, 32> NewIdxs;
   for (unsigned i = 1, e = Ops.size(); i != e; ++i) {
     if ((i == 1 ||
          !isa<StructType>(GetElementPtrInst::getIndexedType(
              SrcElemTy, Ops.slice(1, i - 1)))) &&
-        Ops[i]->getType() != (i == 1 ? IntPtrTy : IntPtrScalarTy)) {
+        Ops[i]->getType()->getScalarType() != IntPtrScalarTy) {
       Any = true;
+      Type *NewType = Ops[i]->getType()->isVectorTy()
+                          ? IntPtrTy
+                          : IntPtrTy->getScalarType();
       NewIdxs.push_back(ConstantExpr::getCast(CastInst::getCastOpcode(Ops[i],
                                                                       true,
-                                                                      IntPtrTy,
+                                                                      NewType,
                                                                       true),
-                                              Ops[i], IntPtrTy));
+                                              Ops[i], NewType));
     } else
       NewIdxs.push_back(Ops[i]);
   }
 
   if (!Any)
     return nullptr;
 
   Constant *C = ConstantExpr::getGetElementPtr(

test/Analysis/ConstantFolding/vectorgep-crash.ll

 ; Function Attrs: sspreq
 define <8 x i64*> @"julia_axpy!_65480"(%Dual* %arg1, <8 x i64> %arg2) {
 top:
 ; CHECK: %VectorGep14 = getelementptr inbounds %Dual, %Dual* %arg1, <8 x i64> %arg2, i32 1, i32 0, i64 0, i32 1, i32 0, i64 0
   %VectorGep14 = getelementptr inbounds %Dual, %Dual* %arg1, <8 x i64> %arg2, i32 1, i32 0, i64 0, i32 1, i32 0, i64 0
   %0 = bitcast <8 x double*> %VectorGep14 to <8 x i64*>
   ret <8 x i64*> %0
 }
+
+%struct.A = type { i32, %struct.B* }
+%struct.B = type { i64, %struct.C* }
+%struct.C = type { i64 }
+
+@G = internal global [65 x %struct.A] zeroinitializer, align 16
+; CHECK-LABEL: @test
+; CHECK: ret <16 x i32*> getelementptr ([65 x %struct.A], [65 x %struct.A]* @G, <16 x i64> zeroinitializer, <16 x i64> <i64 1, i64 2, i64 3, i64 4, i64 5, i64 6, i64 7, i64 8, i64 9, i64 10, i64 11, i64 12, i64 13, i64 14, i64 15, i64 16>, <16 x i32> zeroinitializer)
+define <16 x i32*> @test() {
+vector.body:
+  %VectorGep = getelementptr [65 x %struct.A], [65 x %struct.A]* @G, <16 x i64> zeroinitializer, <16 x i64> <i64 1, i64 2, i64 3, i64 4, i64 5, i64 6, i64 7, i64 8, i64 9, i64 10, i64 11, i64 12, i64 13, i64 14, i64 15, i64 16>, <16 x i32> zeroinitializer
+  ret <16 x i32*> %VectorGep
+}
+
+; CHECK-LABEL: @test2
+; CHECK: ret <16 x i32*> getelementptr ([65 x %struct.A], [65 x %struct.A]* @G, <16 x i64> zeroinitializer, <16 x i64> <i64 1, i64 2, i64 3, i64 4, i64 5, i64 6, i64 7, i64 8, i64 9, 
+define <16 x i32*> @test2() {
+vector.body:
+  %VectorGep = getelementptr [65 x %struct.A], [65 x %struct.A]* @G, <16 x i32> zeroinitializer, <16 x i64> <i64 1, i64 2, i64 3, i64 4, i64 5, i64 6, i64 7, i64 8, i64 9, i64 10, i64 11, i64 12, i64 13, i64 14, i64 15, i64 16>, <16 x i32> zeroinitializer
+  ret <16 x i32*> %VectorGep
+}