[esan|cfrag] Handle complex GEP instr in the cfrag tool

[zhaoqin]

Summary

Iterates all (except the first and the last) operands within the GEP instruction
for instrumentation.

Adds test struct_field_gep.ll.

[bruening]

The commit message has spelling and grammar errors.

llvm/trunk/lib/Transforms/Instrumentation/EfficiencySanitizer.cpp

if (isa<MemSetInst>(MI)) {
     Res = true;
   } else
     llvm_unreachable("Unsupported mem intrinsic type");
   return Res;
 }
 
 bool EfficiencySanitizer::instrumentGetElementPtr(Instruction *I, Module &M) {
   GetElementPtrInst *GepInst = dyn_cast<GetElementPtrInst>(I);
-  if (GepInst == nullptr || !isa<StructType>(GepInst->getSourceElementType()) ||
-      StructTyMap.count(GepInst->getSourceElementType()) == 0 ||
-      !GepInst->hasAllConstantIndices() ||
-      // Only handle simple struct field GEP.
-      GepInst->getNumIndices() != 2) {
-    ++NumIgnoredGEPs;
-    return false;
-  }
-  StructType *StructTy = dyn_cast<StructType>(GepInst->getSourceElementType());
-  if (shouldIgnoreStructType(StructTy)) {
+  bool Res = false;
+  if (GepInst == nullptr || GepInst->getNumIndices() == 1) {
     ++NumIgnoredGEPs;
     return false;
   }
-  ++NumInstrumentedGEPs;
-  // Use the last index as the index within the struct.
-  ConstantInt *Idx = dyn_cast<ConstantInt>(GepInst->getOperand(2));
+  Type *SourceTy = GepInst->getSourceElementType();
+  // Iterate all (except the first and the last) idx within each GEP instruction
+  // for possible nested struct field address calculation.
+  for (unsigned i = 1; i < GepInst->getNumIndices(); ++i) {
+    SmallVector<Value *, 8> IdxVec(GepInst->idx_begin(),
+                                   GepInst->idx_begin() + i);
+    StructType *StructTy = dyn_cast<StructType>(
+        GetElementPtrInst::getIndexedType(SourceTy, IdxVec));
+    if (StructTy == nullptr || shouldIgnoreStructType(StructTy) ||
+        StructTyMap.count(StructTy) == 0)
+      continue;
+    // Get the StructTy's subfield index.
+    ConstantInt *Idx = dyn_cast<ConstantInt>(GepInst->getOperand(i+1));
     if (Idx == nullptr || Idx->getZExtValue() > StructTy->getNumElements())
-    return false;
-
+      continue;
     GlobalVariable *CounterArray = StructTyMap[StructTy];
     if (CounterArray == nullptr)
       return false;
     IRBuilder<> IRB(I);
     Constant *Indices[2];
     // Xref http://llvm.org/docs/LangRef.html#i-getelementptr and
     // http://llvm.org/docs/GetElementPtr.html.
     // The first index of the GEP instruction steps through the first operand,
     // i.e., the array itself.
     Indices[0] = ConstantInt::get(IRB.getInt32Ty(), 0);
     // The second index is the index within the array.
     Indices[1] = ConstantInt::get(IRB.getInt32Ty(), Idx->getZExtValue());
     Constant *Counter =
-      ConstantExpr::getGetElementPtr(ArrayType::get(IRB.getInt64Ty(),
-                                                    StructTy->getNumElements()),
-                                     CounterArray, Indices);
+        ConstantExpr::getGetElementPtr(
+            ArrayType::get(IRB.getInt64Ty(), StructTy->getNumElements()),
     Value *Load = IRB.CreateLoad(Counter);
     IRB.CreateStore(IRB.CreateAdd(Load, ConstantInt::get(IRB.getInt64Ty(), 1)),
                     Counter);
-  return true;
+    Res = true;
+  }
+  if (Res)
+    ++NumInstrumentedGEPs;
+  else
+    ++NumIgnoredGEPs;
+  return Res;
 }
 
 int EfficiencySanitizer::getMemoryAccessFuncIndex(Value *Addr,
                                                   const DataLayout &DL) {
   Type *OrigPtrTy = Addr->getType();
   Type *OrigTy = cast<PointerType>(OrigPtrTy)->getElementType();
   assert(OrigTy->isSized());
   // The size is always a multiple of 8.

llvm/trunk/test/Instrumentation/EfficiencySanitizer/struct_field_gep.ll

+; Test the complex GetElementPtr instruction handling in the EfficiencySanitizer
+; cache fragmentation tool.
+;
+; RUN: opt < %s -esan -esan-cache-frag -S | FileCheck %s
+
+; Code from http://llvm.org/docs/LangRef.html#getelementptr-instruction
+; struct RT {
+;   char A;
+;   int B[10][20];
+;   char C;
+; };
+; struct ST {
+;   int X;
+;   double Y;
+;   struct RT Z;
+; };
+; 
+; int *foo(struct ST *s) {
+;   return &s[1].Z.B[5][13];
+; }
+
+%struct.RT = type { i8, [10 x [20 x i32]], i8 }
+%struct.ST = type { i32, double, %struct.RT }
+
+define i32* @foo(%struct.ST* %s) nounwind uwtable readnone optsize ssp {
+entry:
+  %arrayidx = getelementptr inbounds %struct.ST, %struct.ST* %s, i64 1, i32 2, i32 1, i64 5, i64 13
+  ret i32* %arrayidx
+}
+
+; CHECK:        %0 = load i64, i64* getelementptr inbounds ([3 x i64], [3 x i64]* @"struct.ST#3#13#3#11", i32 0, i32 2)
+; CHECK-NEXT:   %1 = add i64 %0, 1
+; CHECK-NEXT:   store i64 %1, i64* getelementptr inbounds ([3 x i64], [3 x i64]* @"struct.ST#3#13#3#11", i32 0, i32 2)
+; CHECK-NEXT:   %2 = load i64, i64* getelementptr inbounds ([3 x i64], [3 x i64]* @"struct.RT#3#11#14#11", i32 0, i32 1)
+; CHECK-NEXT:   %3 = add i64 %2, 1
+; CHECK-NEXT:   store i64 %3, i64* getelementptr inbounds ([3 x i64], [3 x i64]* @"struct.RT#3#11#14#11", i32 0, i32 1)
+; CHECK-NEXT:   %arrayidx = getelementptr inbounds %struct.ST, %struct.ST* %s, i64 1, i32 2, i32 1, i64 5, i64 13
+; CHECK-NEXT:   ret i32* %arrayidx