[IR] Extract method to get single GEP index from offset (NFC)

This exposes the core logic of getGEPIndicesForOffset() as a
getGEPIndexForOffset() method that only returns a single offset,
instead of following the whole chain.

Author: nikic

llvm/include/llvm/IR/DataLayout.h

@@ -591,6 +591,12 @@ class DataLayout {
   /// the result element type and Offset to be the residual offset.
   SmallVector<APInt> getGEPIndicesForOffset(Type *&ElemTy, APInt &Offset) const;
 
+  /// Get single GEP index to access Offset inside ElemTy. Returns None if
+  /// index cannot be computed, e.g. because the type is not an aggregate.
+  /// ElemTy is updated to be the result element type and Offset to be the
+  /// residual offset.
+  Optional<APInt> getGEPIndexForOffset(Type *&ElemTy, APInt &Offset) const;
+
   /// Returns a StructLayout object, indicating the alignment of the
   /// struct, its size, and the offsets of its fields.
   ///

llvm/lib/IR/DataLayout.cpp

@@ -903,16 +903,14 @@ int64_t DataLayout::getIndexedOffsetInType(Type *ElemTy,
   return Result;
 }
 
-static void addElementIndex(SmallVectorImpl<APInt> &Indices, TypeSize ElemSize,
-                            APInt &Offset) {
+static APInt getElementIndex(TypeSize ElemSize, APInt &Offset) {
   // Skip over scalable or zero size elements. Also skip element sizes larger
   // than the positive index space, because the arithmetic below may not be
   // correct in that case.
   unsigned BitWidth = Offset.getBitWidth();
   if (ElemSize.isScalable() || ElemSize == 0 ||
       !isUIntN(BitWidth - 1, ElemSize)) {
-    Indices.push_back(APInt::getZero(BitWidth));
-    return;
+    return APInt::getZero(BitWidth);
   }
 
   APInt Index = Offset.sdiv(ElemSize);
@@ -923,47 +921,52 @@ static void addElementIndex(SmallVectorImpl<APInt> &Indices, TypeSize ElemSize,
     Offset += ElemSize;
     assert(Offset.isNonNegative() && "Remaining offset shouldn't be negative");
   }
-  Indices.push_back(Index);
+  return Index;
 }
 
-SmallVector<APInt> DataLayout::getGEPIndicesForOffset(Type *&ElemTy,
-                                                      APInt &Offset) const {
-  assert(ElemTy->isSized() && "Element type must be sized");
-  SmallVector<APInt> Indices;
-  addElementIndex(Indices, getTypeAllocSize(ElemTy), Offset);
-  while (Offset != 0) {
-    if (auto *ArrTy = dyn_cast<ArrayType>(ElemTy)) {
-      ElemTy = ArrTy->getElementType();
-      addElementIndex(Indices, getTypeAllocSize(ElemTy), Offset);
-      continue;
-    }
+Optional<APInt> DataLayout::getGEPIndexForOffset(Type *&ElemTy,
+                                                 APInt &Offset) const {
+  if (auto *ArrTy = dyn_cast<ArrayType>(ElemTy)) {
+    ElemTy = ArrTy->getElementType();
+    return getElementIndex(getTypeAllocSize(ElemTy), Offset);
+  }
 
-    if (auto *VecTy = dyn_cast<VectorType>(ElemTy)) {
-      ElemTy = VecTy->getElementType();
-      unsigned ElemSizeInBits = getTypeSizeInBits(ElemTy).getFixedSize();
-      // GEPs over non-multiple of 8 size vector elements are invalid.
-      if (ElemSizeInBits % 8 != 0)
-        break;
+  if (auto *VecTy = dyn_cast<VectorType>(ElemTy)) {
+    ElemTy = VecTy->getElementType();
+    unsigned ElemSizeInBits = getTypeSizeInBits(ElemTy).getFixedSize();
+    // GEPs over non-multiple of 8 size vector elements are invalid.
+    if (ElemSizeInBits % 8 != 0)
+      return None;
 
-      addElementIndex(Indices, TypeSize::Fixed(ElemSizeInBits / 8), Offset);
-      continue;
-    }
+    return getElementIndex(TypeSize::Fixed(ElemSizeInBits / 8), Offset);
+  }
 
-    if (auto *STy = dyn_cast<StructType>(ElemTy)) {
-      const StructLayout *SL = getStructLayout(STy);
-      uint64_t IntOffset = Offset.getZExtValue();
-      if (IntOffset >= SL->getSizeInBytes())
-        break;
+  if (auto *STy = dyn_cast<StructType>(ElemTy)) {
+    const StructLayout *SL = getStructLayout(STy);
+    uint64_t IntOffset = Offset.getZExtValue();
+    if (IntOffset >= SL->getSizeInBytes())
+      return None;
 
-      unsigned Index = SL->getElementContainingOffset(IntOffset);
-      Offset -= SL->getElementOffset(Index);
-      ElemTy = STy->getElementType(Index);
-      Indices.push_back(APInt(32, Index));
-      continue;
-    }
+    unsigned Index = SL->getElementContainingOffset(IntOffset);
+    Offset -= SL->getElementOffset(Index);
+    ElemTy = STy->getElementType(Index);
+    return APInt(32, Index);
+  }
+
+  // Non-aggregate type.
+  return None;
+}
 
-    // Can't index into non-aggregate type.
-    break;
+SmallVector<APInt> DataLayout::getGEPIndicesForOffset(Type *&ElemTy,
+                                                      APInt &Offset) const {
+  assert(ElemTy->isSized() && "Element type must be sized");
+  SmallVector<APInt> Indices;
+  Indices.push_back(getElementIndex(getTypeAllocSize(ElemTy), Offset));
+  while (Offset != 0) {
+    Optional<APInt> Index = getGEPIndexForOffset(ElemTy, Offset);
+    if (!Index)
+      break;
+    Indices.push_back(*Index);
   }
 
   return Indices;