[GlobalOpt] Fix a miscompile when evaluating struct initializers.

The bug was that evaluateBitcastFromPtr attempts a narrowing to a struct's 0th
element of a store that covers other elements. While this is okay on the load
side, applying it to stores causes us to miss the writes to the additionally
covered elements.

rdar://79503568

Differential revision: https://reviews.llvm.org/D105838

Author: jroelofs

llvm/lib/Transforms/Utils/Evaluator.cpp

@@ -174,16 +174,16 @@ static bool isSimpleEnoughPointerToCommit(Constant *C, const DataLayout &DL) {
   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.
+/// Apply \p TryLoad to Ptr. If this returns \p 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
+/// \p TryLoad, or \p 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) {
+                       std::function<Constant *(Constant *)> TryLoad) {
   Constant *Val;
-  while (!(Val = Func(Ptr))) {
+  while (!(Val = TryLoad(Ptr))) {
     // If Ty is a non-opaque 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.
@@ -211,9 +211,11 @@ static Constant *getInitializer(Constant *C) {
 Constant *Evaluator::ComputeLoadResult(Constant *P, Type *Ty) {
   // If this memory location has been recently stored, use the stored value: it
   // is the most up-to-date.
-  auto findMemLoc = [this](Constant *Ptr) { return MutatedMemory.lookup(Ptr); };
+  auto TryFindMemLoc = [this](Constant *Ptr) {
+    return MutatedMemory.lookup(Ptr);
+  };
 
-  if (Constant *Val = findMemLoc(P))
+  if (Constant *Val = TryFindMemLoc(P))
     return Val;
 
   // Access it.
@@ -237,7 +239,7 @@ Constant *Evaluator::ComputeLoadResult(Constant *P, Type *Ty) {
       // 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);
+          evaluateBitcastFromPtr(CE->getOperand(0), DL, TLI, TryFindMemLoc);
       if (!Val)
         Val = getInitializer(CE->getOperand(0));
       if (Val)
@@ -369,17 +371,25 @@ bool Evaluator::EvaluateBlock(BasicBlock::iterator CurInst, BasicBlock *&NextBB,
           // legal. If it's not, we can try introspecting the type to find a
           // legal conversion.
 
-          auto castValTy = [&](Constant *P) -> Constant * {
-            Type *Ty = cast<PointerType>(P->getType())->getElementType();
-            if (Constant *FV = ConstantFoldLoadThroughBitcast(Val, Ty, DL)) {
+          auto TryCastValTy = [&](Constant *P) -> Constant * {
+            // The conversion is illegal if the store is wider than the
+            // pointee proposed by `evaluateBitcastFromPtr`, since that would
+            // drop stores to other struct elements when the caller attempts to
+            // look through a struct's 0th element.
+            Type *NewTy = cast<PointerType>(P->getType())->getElementType();
+            Type *STy = Val->getType();
+            if (DL.getTypeSizeInBits(NewTy) < DL.getTypeSizeInBits(STy))
+              return nullptr;
+
+            if (Constant *FV = ConstantFoldLoadThroughBitcast(Val, NewTy, DL)) {
               Ptr = P;
               return FV;
             }
             return nullptr;
           };
 
           Constant *NewVal =
-              evaluateBitcastFromPtr(CE->getOperand(0), DL, TLI, castValTy);
+              evaluateBitcastFromPtr(CE->getOperand(0), DL, TLI, TryCastValTy);
           if (!NewVal) {
             LLVM_DEBUG(dbgs() << "Failed to bitcast constant ptr, can not "
                                  "evaluate.\n");

llvm/test/Transforms/GlobalOpt/store-struct-element.ll

@@ -0,0 +1,36 @@
+; RUN: opt < %s -globalopt -S -o - | FileCheck %s
+
+%class.Class = type { i8, i8, i8, i8 }
+@A = local_unnamed_addr global %class.Class undef, align 4
+@B = local_unnamed_addr global %class.Class undef, align 4
+
+@llvm.global_ctors = appending global [2 x { i32, void ()*, i8* }] [
+  { i32, void ()*, i8* } { i32 65535, void ()* @initA, i8* null },
+  { i32, void ()*, i8* } { i32 65535, void ()* @initB, i8* null }
+]
+
+define internal void @initA() section "__TEXT,__StaticInit,regular,pure_instructions" {
+entry:
+  store i32 -1, i32* bitcast (%class.Class* @A to i32*), align 4
+  ret void
+}
+
+define internal void @initB() section "__TEXT,__StaticInit,regular,pure_instructions" {
+entry:
+  store i8 -1, i8* bitcast (%class.Class* @B to i8*), align 4
+  ret void
+}
+
+; rdar://79503568
+; Check that we don't miscompile when the store covers the whole struct.
+; CHECK-NOT: @A = local_unnamed_addr global %class.Class { i8 -1, i8 undef, i8 undef, i8 undef }, align 4
+
+; FIXME: We could optimzie this as { i8 -1, i8 -1, i8 -1, i8 -1 } if constant folding were a little smarter.
+; CHECK: @A = local_unnamed_addr global %class.Class undef, align 4
+
+; Check that we still perform the transform when store is smaller than the width of the 0th element.
+; CHECK: @B = local_unnamed_addr global %class.Class { i8 -1, i8 undef, i8 undef, i8 undef }, align 4
+
+; CHECK: define internal void @initA()
+; CHECK-NOT: define internal void @initB()
+