[Evaluator] Walk initial elements when handling load through bitcast

When evaluating a store through a bitcast, the evaluator tries to move the
bitcast from the pointer onto the stored value. If the cast is invalid, it
tries to "introspect" the type to get a valid cast by obtaining a pointer to
the initial element (if the type is nested, this may require walking several
initial elements).

In some situations it is possible to get a bitcast on a load (e.g. with
unions, where the bitcast may not be the same type as the store). However,
equivalent logic to the store to introspect the type is missing. This patch
add this logic.

Note, when developing the patch I was unhappy with adding similar logic
directly to the load case as it could get out of step. Instead, I have
abstracted the "introspection" into a helper function, with the specifics
being handled by a passed-in lambda function.

Differential Revision: https://reviews.llvm.org/D60793

llvm-svn: 359205

llvm/lib/Transforms/Utils/Evaluator.cpp

@@ -174,6 +174,34 @@ static bool isSimpleEnoughPointerToCommit(Constant *C) {
   return false;
 }
 
+/// Apply 'Func' to Ptr. If this returns nullptr, introspect the pointer's
+/// type and walk down through the initial elements to obtain additional
+/// pointers to try. Returns the first non-null return value from Func, or
+/// nullptr if the type can't be introspected further.
+static Constant *
+evaluateBitcastFromPtr(Constant *Ptr, const DataLayout &DL,
+                       const TargetLibraryInfo *TLI,
+                       std::function<Constant *(Constant *)> Func) {
+  Constant *Val;
+  while (!(Val = Func(Ptr))) {
+    // If Ty is a struct, we can convert the pointer to the struct
+    // into a pointer to its first member.
+    // FIXME: This could be extended to support arrays as well.
+    Type *Ty = cast<PointerType>(Ptr->getType())->getElementType();
+    if (!isa<StructType>(Ty))
+      break;
+
+    IntegerType *IdxTy = IntegerType::get(Ty->getContext(), 32);
+    Constant *IdxZero = ConstantInt::get(IdxTy, 0, false);
+    Constant *const IdxList[] = {IdxZero, IdxZero};
+
+    Ptr = ConstantExpr::getGetElementPtr(Ty, Ptr, IdxList);
+    if (auto *FoldedPtr = ConstantFoldConstant(Ptr, DL, TLI))
+      Ptr = FoldedPtr;
+  }
+  return Val;
+}
+
 static Constant *getInitializer(Constant *C) {
   auto *GV = dyn_cast<GlobalVariable>(C);
   return GV && GV->hasDefinitiveInitializer() ? GV->getInitializer() : nullptr;
@@ -184,8 +212,14 @@ static Constant *getInitializer(Constant *C) {
 Constant *Evaluator::ComputeLoadResult(Constant *P) {
   // If this memory location has been recently stored, use the stored value: it
   // is the most up-to-date.
-  DenseMap<Constant*, Constant*>::const_iterator I = MutatedMemory.find(P);
-  if (I != MutatedMemory.end()) return I->second;
+  auto findMemLoc = [this](Constant *Ptr) {
+    DenseMap<Constant *, Constant *>::const_iterator I =
+        MutatedMemory.find(Ptr);
+    return I != MutatedMemory.end() ? I->second : nullptr;
+  };
+
+  if (Constant *Val = findMemLoc(P))
+    return Val;
 
   // Access it.
   if (GlobalVariable *GV = dyn_cast<GlobalVariable>(P)) {
@@ -203,13 +237,17 @@ Constant *Evaluator::ComputeLoadResult(Constant *P) {
       break;
     // Handle a constantexpr bitcast.
     case Instruction::BitCast:
-      Constant *Val = getVal(CE->getOperand(0));
-      auto MM = MutatedMemory.find(Val);
-      auto *I = (MM != MutatedMemory.end()) ? MM->second
-                                            : getInitializer(CE->getOperand(0));
-      if (I)
+      // We're evaluating a load through a pointer that was bitcast to a
+      // different type. See if the "from" pointer has recently been stored.
+      // If it hasn't, we may still be able to find a stored pointer by
+      // introspecting the type.
+      Constant *Val =
+          evaluateBitcastFromPtr(CE->getOperand(0), DL, TLI, findMemLoc);
+      if (!Val)
+        Val = getInitializer(CE->getOperand(0));
+      if (Val)
         return ConstantFoldLoadThroughBitcast(
-            I, P->getType()->getPointerElementType(), DL);
+            Val, P->getType()->getPointerElementType(), DL);
       break;
     }
   }
@@ -329,37 +367,26 @@ bool Evaluator::EvaluateBlock(BasicBlock::iterator CurInst,
                      << "Attempting to resolve bitcast on constant ptr.\n");
           // If we're evaluating a store through a bitcast, then we need
           // to pull the bitcast off the pointer type and push it onto the
-          // stored value.
-          Ptr = CE->getOperand(0);
+          // stored value. In order to push the bitcast onto the stored value,
+          // a bitcast from the pointer's element type to Val's type must be
+          // legal. If it's not, we can try introspecting the type to find a
+          // legal conversion.
 
-          Type *NewTy = cast<PointerType>(Ptr->getType())->getElementType();
-
-          // In order to push the bitcast onto the stored value, a bitcast
-          // from NewTy to Val's type must be legal.  If it's not, we can try
-          // introspecting NewTy to find a legal conversion.
-          Constant *NewVal;
-          while (!(NewVal = ConstantFoldLoadThroughBitcast(Val, NewTy, DL))) {
-            // If NewTy is a struct, we can convert the pointer to the struct
-            // into a pointer to its first member.
-            // FIXME: This could be extended to support arrays as well.
-            if (StructType *STy = dyn_cast<StructType>(NewTy)) {
-
-              IntegerType *IdxTy = IntegerType::get(NewTy->getContext(), 32);
-              Constant *IdxZero = ConstantInt::get(IdxTy, 0, false);
-              Constant * const IdxList[] = {IdxZero, IdxZero};
-
-              Ptr = ConstantExpr::getGetElementPtr(NewTy, Ptr, IdxList);
-              if (auto *FoldedPtr = ConstantFoldConstant(Ptr, DL, TLI))
-                Ptr = FoldedPtr;
-              NewTy = STy->getTypeAtIndex(0U);
-
-              // If we can't improve the situation by introspecting NewTy,
-              // we have to give up.
-            } else {
-              LLVM_DEBUG(dbgs() << "Failed to bitcast constant ptr, can not "
-                                   "evaluate.\n");
-              return false;
+          auto castValTy = [&](Constant *P) -> Constant * {
+            Type *Ty = cast<PointerType>(P->getType())->getElementType();
+            if (Constant *FV = ConstantFoldLoadThroughBitcast(Val, Ty, DL)) {
+              Ptr = P;
+              return FV;
             }
+            return nullptr;
+          };
+
+          Constant *NewVal =
+              evaluateBitcastFromPtr(CE->getOperand(0), DL, TLI, castValTy);
+          if (!NewVal) {
+            LLVM_DEBUG(dbgs() << "Failed to bitcast constant ptr, can not "
+                                 "evaluate.\n");
+            return false;
           }
 
           Val = NewVal;

llvm/test/Transforms/GlobalOpt/evaluate-bitcast-2.ll

@@ -0,0 +1,51 @@
+; RUN: opt < %s -globalopt -S | FileCheck %s
+
+; Test the evaluation of a store and a load via a bitcast, and check
+; that globals are constant folded to the correct value.
+
+; CHECK: @u = dso_local local_unnamed_addr global %union.A { i8* inttoptr (i64 12345 to i8*) }, align 8
+; CHECK: @l = dso_local local_unnamed_addr global i64 12345, align 8
+
+; Test derived from:
+;
+; union A {
+;   A(long long ll) : l(ll) {}
+;   void *p;
+;   long long l;
+; } u(12345);
+;
+; long long l = u.l;
+
+target datalayout = "e-m:e-i64:64-f80:128-n8:16:32:64-S128"
+target triple = "x86_64-unknown-linux-gnu"
+
+%union.A = type { i8* }
+
+$_ZN1AC2Ex = comdat any
+
+@u = dso_local global %union.A zeroinitializer, align 8
+@l = dso_local global i64 0, align 8
+@llvm.global_ctors = appending global [1 x { i32, void ()*, i8* }] [{ i32, void ()*, i8* } { i32 65535, void ()* @_GLOBAL__sub_I_test.cpp, i8* null }]
+
+define internal void @__cxx_global_var_init() section ".text.startup" {
+  call void @_ZN1AC2Ex(%union.A* @u, i64 12345)
+  ret void
+}
+
+define linkonce_odr dso_local void @_ZN1AC2Ex(%union.A* %this, i64 %ll) unnamed_addr comdat align 2 {
+  %l = bitcast %union.A* %this to i64*
+  store i64 %ll, i64* %l, align 8
+  ret void
+}
+
+define internal void @__cxx_global_var_init.1() section ".text.startup" {
+  %1 = load i64, i64* bitcast (%union.A* @u to i64*), align 8
+  store i64 %1, i64* @l, align 8
+  ret void
+}
+
+define internal void @_GLOBAL__sub_I_test.cpp() section ".text.startup" {
+  call void @__cxx_global_var_init()
+  call void @__cxx_global_var_init.1()
+  ret void
+}

llvm/test/Transforms/GlobalOpt/evaluate-bitcast-3.ll

@@ -0,0 +1,42 @@
+; RUN: opt < %s -globalopt -S | FileCheck %s
+
+; Test the evaluation of a load via a bitcast and a store via a GEP.
+; Check that globals are constant folded to the correct value.
+
+; CHECK: @u = dso_local local_unnamed_addr global %union.A { i8* inttoptr (i64 12345 to i8*) }, align 8
+; CHECK: @l = dso_local local_unnamed_addr global i64 12345, align 8
+
+target datalayout = "e-m:e-i64:64-f80:128-n8:16:32:64-S128"
+target triple = "x86_64-unknown-linux-gnu"
+
+%union.A = type { i8* }
+
+$_ZN1AC2Ex = comdat any
+
+@u = dso_local global %union.A zeroinitializer, align 8
+@l = dso_local global i64 0, align 8
+@llvm.global_ctors = appending global [1 x { i32, void ()*, i8* }] [{ i32, void ()*, i8* } { i32 65535, void ()* @_GLOBAL__sub_I_test.cpp, i8* null }]
+
+define internal void @__cxx_global_var_init() section ".text.startup" {
+  call void @_ZN1AC2Ex(%union.A* @u, i64 12345)
+  ret void
+}
+
+define linkonce_odr dso_local void @_ZN1AC2Ex(%union.A* %this, i64 %ll) unnamed_addr comdat align 2 {
+  %l = inttoptr i64 %ll to i8*
+  %p = getelementptr inbounds %union.A, %union.A* %this, i64 0, i32 0
+  store i8* %l, i8** %p
+  ret void
+}
+
+define internal void @__cxx_global_var_init.1() section ".text.startup" {
+  %1 = load i64, i64* bitcast (%union.A* @u to i64*), align 8
+  store i64 %1, i64* @l, align 8
+  ret void
+}
+
+define internal void @_GLOBAL__sub_I_test.cpp() section ".text.startup" {
+  call void @__cxx_global_var_init()
+  call void @__cxx_global_var_init.1()
+  ret void
+}