[msan] Strict checking of loads.

[eugenis]

Summary

Insert shadow checks after each load.

DO NOT COMMIT. Experimental.

[eugenis]

Btw, this makes building msan unit tests insanely slow. Must be all the extra basic blocks.

I'm guessing this change does not exactly improve code size, too.

[eugenis]

A reasonable compromise could be to check function call arguments. I.e. allow loading and storing (basically, copying) uninit values, but not passing them to other functions as arguments or return values.

llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp

struct MemorySanitizerVisitor : public InstVisitor<MemorySanitizerVisitor> {
   bool PoisonStack;
   bool PoisonUndef;
   bool CheckReturnValue;
 
   struct ShadowOriginAndInsertPoint {
     Value *Shadow;
     Value *Origin;
     Instruction *OrigIns;
+    bool After;
 
-    ShadowOriginAndInsertPoint(Value *S, Value *O, Instruction *I)
-      : Shadow(S), Origin(O), OrigIns(I) {}
+    ShadowOriginAndInsertPoint(Value *S, Value *O, Instruction *I, bool A)
+        : Shadow(S), Origin(O), OrigIns(I), After(A) {}
   };
   SmallVector<ShadowOriginAndInsertPoint, 16> InstrumentationList;
   SmallVector<StoreInst *, 16> StoreList;
 
   MemorySanitizerVisitor(Function &F, MemorySanitizer &MS)
       : F(F), MS(MS), VAHelper(CreateVarArgHelper(F, MS, *this)) {
     bool SanitizeFunction = F.hasFnAttribute(Attribute::SanitizeMemory);
     InsertChecks = SanitizeFunction;
void materializeOneCheck(Instruction *OrigIns, Value *Shadow, Value *Origin,
     }
   }
 
   void materializeChecks(bool InstrumentWithCalls) {
     for (const auto &ShadowData : InstrumentationList) {
       Instruction *OrigIns = ShadowData.OrigIns;
       Value *Shadow = ShadowData.Shadow;
       Value *Origin = ShadowData.Origin;
+      if (ShadowData.After)
+        OrigIns = OrigIns->getNextNode();
       materializeOneCheck(OrigIns, Shadow, Origin, InstrumentWithCalls);
     }
     DEBUG(dbgs() << "DONE:\n" << F);
   }
 
   /// \brief Add MemorySanitizer instrumentation to a function.
   bool runOnFunction() {
     MS.initializeCallbacks(*F.getParent());
struct MemorySanitizerVisitor : public InstVisitor<MemorySanitizerVisitor> {
   Value *getOrigin(Instruction *I, int i) {
     return getOrigin(I->getOperand(i));
   }
 
   /// \brief Remember the place where a shadow check should be inserted.
   ///
   /// This location will be later instrumented with a check that will print a
   /// UMR warning in runtime if the shadow value is not 0.
-  void insertShadowCheck(Value *Shadow, Value *Origin, Instruction *OrigIns) {
+  void insertShadowCheck(Value *Shadow, Value *Origin, Instruction *OrigIns,
+                         bool After = false) {
     assert(Shadow);
     if (!InsertChecks) return;
 #ifndef NDEBUG
     Type *ShadowTy = Shadow->getType();
     assert((isa<IntegerType>(ShadowTy) || isa<VectorType>(ShadowTy)) &&
            "Can only insert checks for integer and vector shadow types");
 #endif
     InstrumentationList.push_back(
-        ShadowOriginAndInsertPoint(Shadow, Origin, OrigIns));
+        ShadowOriginAndInsertPoint(Shadow, Origin, OrigIns, After));
   }
 
   /// \brief Remember the place where a shadow check should be inserted.
   ///
   /// This location will be later instrumented with a check that will print a
   /// UMR warning in runtime if the value is not fully defined.
-  void insertShadowCheck(Value *Val, Instruction *OrigIns) {
+  void insertShadowCheck(Value *Val, Instruction *OrigIns, bool After = false) {
     assert(Val);
     Value *Shadow, *Origin;
     if (ClCheckConstantShadow) {
       Shadow = getShadow(Val);
       if (!Shadow) return;
       Origin = getOrigin(Val);
     } else {
       Shadow = dyn_cast_or_null<Instruction>(getShadow(Val));
       if (!Shadow) return;
       Origin = dyn_cast_or_null<Instruction>(getOrigin(Val));
     }
-    insertShadowCheck(Shadow, Origin, OrigIns);
+    insertShadowCheck(Shadow, Origin, OrigIns, After);
   }
 
   AtomicOrdering addReleaseOrdering(AtomicOrdering a) {
     switch (a) {
       case AtomicOrdering::NotAtomic:
         return AtomicOrdering::NotAtomic;
       case AtomicOrdering::Unordered:
       case AtomicOrdering::Monotonic:
if (MS.TrackOrigins) {
         unsigned Alignment = I.getAlignment();
         unsigned OriginAlignment = std::max(kMinOriginAlignment, Alignment);
         setOrigin(&I, IRB.CreateAlignedLoad(getOriginPtr(Addr, IRB, Alignment),
                                             OriginAlignment));
       } else {
         setOrigin(&I, getCleanOrigin());
       }
     }
+
+    // Shadow check of the loaded value must follow loads of shadow and origin.
+    // The first instruction that satisfies this requirement is
+    // IRB.GetInsertPoint(), but it has not been instrumented yet => may get
+    // stuff inserted before it, or even disappear without trace. Record the
+    // location as "after the last instruction of this load's instrumentation".
+
+    // FIXME: implement shadow check for aggregate types and re-enable here and
+    // in storeOrigin().
+    if (PropagateShadow && !I.getType()->isAggregateType())
+      insertShadowCheck(&I, &*IRB.GetInsertPoint()->getPrevNode(),
+                        /*After=*/true);
   }
 
   /// \brief Instrument StoreInst
   ///
   /// Stores the corresponding shadow and (optionally) origin.
   /// Optionally, checks that the store address is fully defined.
   void visitStoreInst(StoreInst &I) {
     StoreList.push_back(&I);