[msan] More precise instrumentation of select IR.

Some bits of select result may be initialized even if select condition
is not.

https://code.google.com/p/memory-sanitizer/issues/detail?id=50


git-svn-id: https://llvm.org/svn/llvm-project/llvm/trunk@204716 91177308-0d34-0410-b5e6-96231b3b80d8

2 files changed, 74 insertions, 34 deletions

diff --git a/lib/Transforms/Instrumentation/MemorySanitizer.cpp b/lib/Transforms/Instrumentation/MemorySanitizer.cpp
index 4ac6eab552..ec1a195c95 100644
--- a/lib/Transforms/Instrumentation/MemorySanitizer.cpp
+++ b/lib/Transforms/Instrumentation/MemorySanitizer.cpp
@@ -1324,6 +1324,17 @@ struct MemorySanitizerVisitor : public InstVisitor<MemorySanitizerVisitor> {
     // TODO: handle struct types.
   }
 
+  /// \brief Cast an application value to the type of its own shadow.
+  Value *CreateAppToShadowCast(IRBuilder<> &IRB, Value *V) {
+    Type *ShadowTy = getShadowTy(V);
+    if (V->getType() == ShadowTy)
+      return V;
+    if (V->getType()->isPtrOrPtrVectorTy())
+      return IRB.CreatePtrToInt(V, ShadowTy);
+    else
+      return IRB.CreateBitCast(V, ShadowTy);
+  }
+
   /// \brief Propagate shadow for arbitrary operation.
   void handleShadowOr(Instruction &I) {
     IRBuilder<> IRB(&I);
@@ -2180,40 +2191,51 @@ struct MemorySanitizerVisitor : public InstVisitor<MemorySanitizerVisitor> {
   void visitSelectInst(SelectInst& I) {
     IRBuilder<> IRB(&I);
     // a = select b, c, d
-    Value *S = IRB.CreateSelect(I.getCondition(), getShadow(I.getTrueValue()),
-                                getShadow(I.getFalseValue()));
+    Value *B = I.getCondition();
+    Value *C = I.getTrueValue();
+    Value *D = I.getFalseValue();
+    Value *Sb = getShadow(B);
+    Value *Sc = getShadow(C);
+    Value *Sd = getShadow(D);
+
+    // Result shadow if condition shadow is 0.
+    Value *Sa0 = IRB.CreateSelect(B, Sc, Sd);
+    Value *Sa1;
     if (I.getType()->isAggregateType()) {
       // To avoid "sign extending" i1 to an arbitrary aggregate type, we just do
       // an extra "select". This results in much more compact IR.
       // Sa = select Sb, poisoned, (select b, Sc, Sd)
-      S = IRB.CreateSelect(getShadow(I.getCondition()),
-                           getPoisonedShadow(getShadowTy(I.getType())), S,
-                           "_msprop_select_agg");
+      Sa1 = getPoisonedShadow(getShadowTy(I.getType()));
     } else {
-      // Sa = (sext Sb) | (select b, Sc, Sd)
-      S = IRB.CreateOr(S, CreateShadowCast(IRB, getShadow(I.getCondition()),
-                                           S->getType(), true),
-                       "_msprop_select");
+      // Sa = select Sb, [ (c^d) | Sc | Sd ], [ b ? Sc : Sd ]
+      // If Sb (condition is poisoned), look for bits in c and d that are equal
+      // and both unpoisoned.
+      // If !Sb (condition is unpoisoned), simply pick one of Sc and Sd.
+
+      // Cast arguments to shadow-compatible type.
+      C = CreateAppToShadowCast(IRB, C);
+      D = CreateAppToShadowCast(IRB, D);
+
+      // Result shadow if condition shadow is 1.
+      Sa1 = IRB.CreateOr(IRB.CreateXor(C, D), IRB.CreateOr(Sc, Sd));
     }
-    setShadow(&I, S);
+    Value *Sa = IRB.CreateSelect(Sb, Sa1, Sa0, "_msprop_select");
+    setShadow(&I, Sa);
     if (MS.TrackOrigins) {
       // Origins are always i32, so any vector conditions must be flattened.
       // FIXME: consider tracking vector origins for app vectors?
-      Value *Cond = I.getCondition();
-      Value *CondShadow = getShadow(Cond);
-      if (Cond->getType()->isVectorTy()) {
-        Type *FlatTy = getShadowTyNoVec(Cond->getType());
-        Cond = IRB.CreateICmpNE(IRB.CreateBitCast(Cond, FlatTy),
+      if (B->getType()->isVectorTy()) {
+        Type *FlatTy = getShadowTyNoVec(B->getType());
+        B = IRB.CreateICmpNE(IRB.CreateBitCast(B, FlatTy),
                                 ConstantInt::getNullValue(FlatTy));
-        CondShadow = IRB.CreateICmpNE(IRB.CreateBitCast(CondShadow, FlatTy),
+        Sb = IRB.CreateICmpNE(IRB.CreateBitCast(Sb, FlatTy),
                                       ConstantInt::getNullValue(FlatTy));
       }
       // a = select b, c, d
       // Oa = Sb ? Ob : (b ? Oc : Od)
       setOrigin(&I, IRB.CreateSelect(
-                        CondShadow, getOrigin(I.getCondition()),
-                        IRB.CreateSelect(Cond, getOrigin(I.getTrueValue()),
-                                         getOrigin(I.getFalseValue()))));
+                        Sb, getOrigin(I.getCondition()),
+                        IRB.CreateSelect(B, getOrigin(C), getOrigin(D))));
     }
   }
 
diff --git a/test/Instrumentation/MemorySanitizer/msan_basic.ll b/test/Instrumentation/MemorySanitizer/msan_basic.ll
index 69e5c18465..624c706c6e 100644
--- a/test/Instrumentation/MemorySanitizer/msan_basic.ll
+++ b/test/Instrumentation/MemorySanitizer/msan_basic.ll
@@ -249,13 +249,16 @@ entry:
 }
 
 ; CHECK: @Select
-; CHECK: select
-; CHECK-NEXT: sext i1 {{.*}} to i32
+; CHECK: select i1
+; CHECK-NEXT: or i32
+; CHECK-NEXT: xor i32
 ; CHECK-NEXT: or i32
-; CHECK-NEXT: select
+; CHECK-NEXT: select i1
 ; CHECK-ORIGINS: select
 ; CHECK-ORIGINS: select
-; CHECK: select
+; CHECK-NEXT: select i1
+; CHECK: store i32{{.*}}@__msan_retval_tls
+; CHECK-ORIGINS: store i32{{.*}}@__msan_retval_origin_tls
 ; CHECK: ret i32
 
 
@@ -271,19 +274,18 @@ entry:
 
 ; CHECK: @SelectVector
 ; CHECK: select <8 x i1>
-; CHECK-NEXT: sext <8 x i1> {{.*}} to <8 x i16>
 ; CHECK-NEXT: or <8 x i16>
-; CHECK-ORIGINS: bitcast <8 x i1> {{.*}} to i8
-; CHECK-ORIGINS: icmp ne i8 {{.*}}, 0
-; CHECK-ORIGINS: bitcast <8 x i1> {{.*}} to i8
-; CHECK-ORIGINS: icmp ne i8 {{.*}}, 0
-; CHECK-ORIGINS: select i1
-; CHECK-ORIGINS: select i1
-; CHECK: select <8 x i1>
+; CHECK-NEXT: xor <8 x i16>
+; CHECK-NEXT: or <8 x i16>
+; CHECK-NEXT: select <8 x i1>
+; CHECK-ORIGINS: select
+; CHECK-ORIGINS: select
+; CHECK-NEXT: select <8 x i1>
+; CHECK: store <8 x i16>{{.*}}@__msan_retval_tls
+; CHECK-ORIGINS: store i32{{.*}}@__msan_retval_origin_tls
 ; CHECK: ret <8 x i16>
 
 
-
 ; Check that we propagate origin for "select" with scalar condition and vector
 ; arguments. Select condition shadow is sign-extended to the vector type and
 ; mixed into the result shadow.
@@ -296,9 +298,10 @@ entry:
 
 ; CHECK: @SelectVector2
 ; CHECK: select i1
-; CHECK: sext i1 {{.*}} to i128
-; CHECK: bitcast i128 {{.*}} to <8 x i16>
 ; CHECK: or <8 x i16>
+; CHECK: xor <8 x i16>
+; CHECK: or <8 x i16>
+; CHECK: select i1
 ; CHECK-ORIGINS: select i1
 ; CHECK-ORIGINS: select i1
 ; CHECK: select i1
@@ -320,6 +323,21 @@ entry:
 ; CHECK: ret { i64, i64 }
 
 
+define { i64*, double } @SelectStruct2(i1 zeroext %x, { i64*, double } %a, { i64*, double } %b) readnone sanitize_memory {
+entry:
+  %c = select i1 %x, { i64*, double } %a, { i64*, double } %b
+  ret { i64*, double } %c
+}
+
+; CHECK: @SelectStruct2
+; CHECK: select i1 {{.*}}, { i64, i64 }
+; CHECK-NEXT: select i1 {{.*}}, { i64, i64 } { i64 -1, i64 -1 }, { i64, i64 }
+; CHECK-ORIGINS: select i1
+; CHECK-ORIGINS: select i1
+; CHECK-NEXT: select i1 {{.*}}, { i64*, double }
+; CHECK: ret { i64*, double }
+
+
 define i8* @IntToPtr(i64 %x) nounwind uwtable readnone sanitize_memory {
 entry:
   %0 = inttoptr i64 %x to i8*