[InstCombine] Generalize ptrtoint(gep) fold (NFC)

We're currently handling a special case of
ptrtoint gep -> add ptrtoint. Reframe the code to make it easier
to add more patterns for this transform.

1 files changed, 21 insertions, 4 deletions

diff --git a/llvm/lib/Transforms/InstCombine/InstCombineCasts.cpp b/llvm/lib/Transforms/InstCombine/InstCombineCasts.cpp
index 8f83047..5f51c92 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineCasts.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstCombineCasts.cpp
@@ -2025,6 +2025,25 @@ Instruction *InstCombinerImpl::visitIntToPtr(IntToPtrInst &CI) {
   return nullptr;
 }
 
+// Whether we should convert ptrtoint(gep P, X) to ptrtoint(P) + X
+static bool shouldPushPtrToIntThroughGEP(GEPOperator *GEP, Type *IntTy,
+                                         const DataLayout &DL) {
+  if (!GEP->hasOneUse())
+    return false;
+
+  // Skip cases whether there is a mismatch between the pointer integer type
+  // and the index type, or the GEP performs an implicit splat operation.
+  if (DL.getIndexType(GEP->getType()) != IntTy ||
+      GEP->getType() != GEP->getPointerOperand()->getType())
+    return false;
+
+  // (ptrtoint (gep (inttoptr Base), Offset)) -> Base + Offset
+  if (match(GEP->getPointerOperand(), m_OneUse(m_IntToPtr(m_Value()))))
+    return true;
+
+  return false;
+}
+
 Instruction *InstCombinerImpl::visitPtrToInt(PtrToIntInst &CI) {
   // If the destination integer type is not the intptr_t type for this target,
   // do a ptrtoint to intptr_t then do a trunc or zext.  This allows the cast
@@ -2064,10 +2083,8 @@ Instruction *InstCombinerImpl::visitPtrToInt(PtrToIntInst &CI) {
     }
 
     // (ptrtoint (gep (inttoptr Base), ...)) -> Base + Offset
-    Value *Base;
-    if (GEP->hasOneUse() &&
-        match(GEP->getPointerOperand(), m_OneUse(m_IntToPtr(m_Value(Base)))) &&
-        Base->getType() == Ty) {
+    if (shouldPushPtrToIntThroughGEP(GEP, Ty, DL)) {
+      Value *Base = Builder.CreatePtrToInt(GEP->getPointerOperand(), Ty);
       Value *Offset = EmitGEPOffset(GEP);
       auto *NewOp = BinaryOperator::CreateAdd(Base, Offset);
       if (GEP->hasNoUnsignedWrap() ||