lib/CodeGen/TaintRelaxedAtomicsUtils.cpp

   1 //===- TaintRelaxedAtomicsUtil.cpp - Utils for tainting relaxed atomics --------===//
   2 //
   3 //                     The LLVM Compiler Infrastructure
   4 //
   5 // This file is distributed under the University of Illinois Open Source
   6 // License. See LICENSE.TXT for details.
   7 //
   8 //===----------------------------------------------------------------------===//
   9
  10 #include "TaintRelaxedAtomicsUtils.h"
  11 #include "llvm/CodeGen/Passes.h"
  12 #include "llvm/ADT/DenseMap.h"
  13 #include "llvm/ADT/SetVector.h"
  14 #include "llvm/ADT/SmallPtrSet.h"
  15 #include "llvm/ADT/SmallSet.h"
  16 #include "llvm/ADT/SmallVector.h"
  17 #include "llvm/ADT/Statistic.h"
  18 #include "llvm/Analysis/InstructionSimplify.h"
  19 #include "llvm/Analysis/MemoryLocation.h"
  20 #include "llvm/Analysis/TargetLibraryInfo.h"
  21 #include "llvm/Analysis/TargetTransformInfo.h"
  22 #include "llvm/Analysis/ValueTracking.h"
  23 #include "llvm/IR/CallSite.h"
  24 #include "llvm/IR/Constants.h"
  25 #include "llvm/IR/DataLayout.h"
  26 #include "llvm/IR/DerivedTypes.h"
  27 #include "llvm/IR/Dominators.h"
  28 #include "llvm/IR/Function.h"
  29 #include "llvm/IR/GetElementPtrTypeIterator.h"
  30 #include "llvm/IR/IRBuilder.h"
  31 #include "llvm/IR/InlineAsm.h"
  32 #include "llvm/IR/InstIterator.h"
  33 #include "llvm/IR/InstrTypes.h"
  34 #include "llvm/IR/Instructions.h"
  35 #include "llvm/IR/IntrinsicInst.h"
  36 #include "llvm/IR/MDBuilder.h"
  37 #include "llvm/IR/NoFolder.h"
  38 #include "llvm/IR/PatternMatch.h"
  39 #include "llvm/IR/Statepoint.h"
  40 #include "llvm/IR/ValueHandle.h"
  41 #include "llvm/IR/ValueMap.h"
  42 #include "llvm/Pass.h"
  43 #include "llvm/Support/CommandLine.h"
  44 #include "llvm/Support/Debug.h"
  45 #include "llvm/Support/raw_ostream.h"
  46 #include "llvm/Target/TargetLowering.h"
  47 #include "llvm/Target/TargetSubtargetInfo.h"
  48 #include "llvm/Transforms/Utils/BasicBlockUtils.h"
  49 #include "llvm/Transforms/Utils/BuildLibCalls.h"
  50 #include "llvm/Transforms/Utils/BypassSlowDivision.h"
  51 #include "llvm/Transforms/Utils/Local.h"
  52 #include "llvm/Transforms/Utils/SimplifyLibCalls.h"
  53 using namespace llvm;
  54 using namespace llvm::PatternMatch;
  55
  56 #define DEBUG_TYPE "taintrelaxedatomics"
  57
  58 namespace llvm {
  59
  60 // The depth we trace down a variable to look for its dependence set.
  61 const unsigned kDependenceDepth = 4;
  62
  63 // Recursively looks for variables that 'Val' depends on at the given depth
  64 // 'Depth', and adds them in 'DepSet'. If 'InsertOnlyLeafNodes' is true, only
  65 // inserts the leaf node values; otherwise, all visited nodes are included in
  66 // 'DepSet'. Note that constants will be ignored.
  67 template <typename SetType>
  68 void recursivelyFindDependence(SetType* DepSet, Value* Val,
  69                                bool InsertOnlyLeafNodes = false,
  70                                unsigned Depth = kDependenceDepth) {
  71   if (Val == nullptr) {
  72     return;
  73   }
  74   if (!InsertOnlyLeafNodes && !isa<Constant>(Val)) {
  75     DepSet->insert(Val);
  76   }
  77   if (Depth == 0) {
  78     // Cannot go deeper. Insert the leaf nodes.
  79     if (InsertOnlyLeafNodes && !isa<Constant>(Val)) {
  80       DepSet->insert(Val);
  81     }
  82     return;
  83   }
  84
  85   // Go one step further to explore the dependence of the operands.
  86   Instruction* I = nullptr;
  87   if ((I = dyn_cast<Instruction>(Val))) {
  88     if (isa<LoadInst>(I)) {
  89       // A load is considerd the leaf load of the dependence tree. Done.
  90       DepSet->insert(Val);
  91       return;
  92     } else if (I->isBinaryOp()) {
  93       BinaryOperator* I = dyn_cast<BinaryOperator>(Val);
  94       Value *Op0 = I->getOperand(0), *Op1 = I->getOperand(1);
  95       recursivelyFindDependence(DepSet, Op0, InsertOnlyLeafNodes, Depth - 1);
  96       recursivelyFindDependence(DepSet, Op1, InsertOnlyLeafNodes, Depth - 1);
  97     } else if (I->isCast()) {
  98       Value* Op0 = I->getOperand(0);
  99       recursivelyFindDependence(DepSet, Op0, InsertOnlyLeafNodes, Depth - 1);
 100     } else if (I->getOpcode() == Instruction::Select) {
 101       Value* Op0 = I->getOperand(0);
 102       Value* Op1 = I->getOperand(1);
 103       Value* Op2 = I->getOperand(2);
 104       recursivelyFindDependence(DepSet, Op0, InsertOnlyLeafNodes, Depth - 1);
 105       recursivelyFindDependence(DepSet, Op1, InsertOnlyLeafNodes, Depth - 1);
 106       recursivelyFindDependence(DepSet, Op2, InsertOnlyLeafNodes, Depth - 1);
 107     } else if (I->getOpcode() == Instruction::GetElementPtr) {
 108       for (unsigned i = 0; i < I->getNumOperands(); i++) {
 109         recursivelyFindDependence(DepSet, I->getOperand(i), InsertOnlyLeafNodes,
 110                                   Depth - 1);
 111       }
 112     } else if (I->getOpcode() == Instruction::Store) {
 113       auto* SI = dyn_cast<StoreInst>(Val);
 114       recursivelyFindDependence(DepSet, SI->getPointerOperand(),
 115                                 InsertOnlyLeafNodes, Depth - 1);
 116       recursivelyFindDependence(DepSet, SI->getValueOperand(),
 117                                 InsertOnlyLeafNodes, Depth - 1);
 118     } else {
 119       Value* Op0 = nullptr;
 120       Value* Op1 = nullptr;
 121       switch (I->getOpcode()) {
 122         case Instruction::ICmp:
 123         case Instruction::FCmp: {
 124           Op0 = I->getOperand(0);
 125           Op1 = I->getOperand(1);
 126           recursivelyFindDependence(DepSet, Op0, InsertOnlyLeafNodes,
 127                                     Depth - 1);
 128           recursivelyFindDependence(DepSet, Op1, InsertOnlyLeafNodes,
 129                                     Depth - 1);
 130           break;
 131         }
 132         case Instruction::PHI: {
 133           for (unsigned i = 0; i < I->getNumOperands(); i++) {
 134             auto* op = I->getOperand(i);
 135             if (DepSet->count(op) == 0) {
 136               recursivelyFindDependence(DepSet, I->getOperand(i),
 137                                         InsertOnlyLeafNodes, Depth - 1);
 138             }
 139           }
 140           break;
 141         }
 142         default: {
 143           // Be conservative. Add it and be done with it.
 144           DepSet->insert(Val);
 145           return;
 146         }
 147       }
 148     }
 149   } else if (isa<Constant>(Val)) {
 150     // Not interested in constant values. Done.
 151     return;
 152   } else {
 153     // Be conservative. Add it and be done with it.
 154     DepSet->insert(Val);
 155     return;
 156   }
 157 }
 158
 159 // Helper function to create a Cast instruction.
 160 template <typename BuilderTy>
 161 Value* createCast(BuilderTy& Builder, Value* DepVal,
 162                   Type* TargetIntegerType) {
 163   Instruction::CastOps CastOp = Instruction::BitCast;
 164   switch (DepVal->getType()->getTypeID()) {
 165     case Type::IntegerTyID: {
 166       assert(TargetIntegerType->getTypeID() == Type::IntegerTyID);
 167       auto* FromType = dyn_cast<IntegerType>(DepVal->getType());
 168       auto* ToType = dyn_cast<IntegerType>(TargetIntegerType);
 169       assert(FromType && ToType);
 170       if (FromType->getBitWidth() <= ToType->getBitWidth()) {
 171         CastOp = Instruction::ZExt;
 172       } else {
 173         CastOp = Instruction::Trunc;
 174       }
 175       break;
 176     }
 177     case Type::FloatTyID:
 178     case Type::DoubleTyID: {
 179       CastOp = Instruction::FPToSI;
 180       break;
 181     }
 182     case Type::PointerTyID: {
 183       CastOp = Instruction::PtrToInt;
 184       break;
 185     }
 186     default: { break; }
 187   }
 188
 189   return Builder.CreateCast(CastOp, DepVal, TargetIntegerType);
 190 }
 191
 192 // Given a value, if it's a tainted address, this function returns the
 193 // instruction that ORs the "dependence value" with the "original address".
 194 // Otherwise, returns nullptr.  This instruction is the first OR instruction
 195 // where one of its operand is an AND instruction with an operand being 0.
 196 //
 197 // E.g., it returns '%4 = or i32 %3, %2' given 'CurrentAddress' is '%5'.
 198 // %0 = load i32, i32* @y, align 4, !tbaa !1
 199 // %cmp = icmp ne i32 %0, 42  // <== this is like the condition
 200 // %1 = sext i1 %cmp to i32
 201 // %2 = ptrtoint i32* @x to i32
 202 // %3 = and i32 %1, 0
 203 // %4 = or i32 %3, %2
 204 // %5 = inttoptr i32 %4 to i32*
 205 // store i32 1, i32* %5, align 4
 206 Instruction* getOrAddress(Value* CurrentAddress) {
 207   // Is it a cast from integer to pointer type.
 208   Instruction* OrAddress = nullptr;
 209   Instruction* AndDep = nullptr;
 210   Constant* ZeroConst = nullptr;
 211
 212   const Instruction* CastToPtr = dyn_cast<Instruction>(CurrentAddress);
 213   if (CastToPtr && CastToPtr->getOpcode() == Instruction::IntToPtr) {
 214     // Is it an OR instruction: %1 = or %and, %actualAddress.
 215     if ((OrAddress = dyn_cast<Instruction>(CastToPtr->getOperand(0))) &&
 216         OrAddress->getOpcode() == Instruction::Or) {
 217       // The first operand should be and AND instruction.
 218       AndDep = dyn_cast<Instruction>(OrAddress->getOperand(0));
 219       if (AndDep && AndDep->getOpcode() == Instruction::And) {
 220         // Also make sure its first operand of the "AND" is 0, or the "AND" is
 221         // marked explicitly by "NoInstCombine".
 222         if ((ZeroConst = dyn_cast<Constant>(AndDep->getOperand(1))) &&
 223             ZeroConst->isNullValue()) {
 224           return OrAddress;
 225         }
 226       }
 227     }
 228   }
 229   // Looks like it's not been tainted.
 230   return nullptr;
 231 }
 232
 233 // Given a value, if it's a tainted address, this function returns the
 234 // instruction that taints the "dependence value". Otherwise, returns nullptr.
 235 // This instruction is the last AND instruction where one of its operand is 0.
 236 // E.g., it returns '%3' given 'CurrentAddress' is '%5'.
 237 // %0 = load i32, i32* @y, align 4, !tbaa !1
 238 // %cmp = icmp ne i32 %0, 42  // <== this is like the condition
 239 // %1 = sext i1 %cmp to i32
 240 // %2 = ptrtoint i32* @x to i32
 241 // %3 = and i32 %1, 0
 242 // %4 = or i32 %3, %2
 243 // %5 = inttoptr i32 %4 to i32*
 244 // store i32 1, i32* %5, align 4
 245 Instruction* getAndDependence(Value* CurrentAddress) {
 246   // If 'CurrentAddress' is tainted, get the OR instruction.
 247   auto* OrAddress = getOrAddress(CurrentAddress);
 248   if (OrAddress == nullptr) {
 249     return nullptr;
 250   }
 251
 252   // No need to check the operands.
 253   auto* AndDepInst = dyn_cast<Instruction>(OrAddress->getOperand(0));
 254   assert(AndDepInst);
 255   return AndDepInst;
 256 }
 257
 258 // Given a value, if it's a tainted address, this function returns
 259 // the "dependence value", which is the first operand in the AND instruction.
 260 // E.g., it returns '%1' given 'CurrentAddress' is '%5'.
 261 // %0 = load i32, i32* @y, align 4, !tbaa !1
 262 // %cmp = icmp ne i32 %0, 42  // <== this is like the condition
 263 // %1 = sext i1 %cmp to i32
 264 // %2 = ptrtoint i32* @x to i32
 265 // %3 = and i32 %1, 0
 266 // %4 = or i32 %3, %2
 267 // %5 = inttoptr i32 %4 to i32*
 268 // store i32 1, i32* %5, align 4
 269 Value* getDependence(Value* CurrentAddress) {
 270   auto* AndInst = getAndDependence(CurrentAddress);
 271   if (AndInst == nullptr) {
 272     return nullptr;
 273   }
 274   return AndInst->getOperand(0);
 275 }
 276
 277 // Given an address that has been tainted, returns the only condition it depends
 278 // on, if any; otherwise, returns nullptr.
 279 Value* getConditionDependence(Value* Address) {
 280   auto* Dep = getDependence(Address);
 281   if (Dep == nullptr) {
 282     // 'Address' has not been dependence-tainted.
 283     return nullptr;
 284   }
 285
 286   Value* Operand = Dep;
 287   while (true) {
 288     auto* Inst = dyn_cast<Instruction>(Operand);
 289     if (Inst == nullptr) {
 290       // Non-instruction type does not have condition dependence.
 291       return nullptr;
 292     }
 293     if (Inst->getOpcode() == Instruction::ICmp) {
 294       return Inst;
 295     } else {
 296       if (Inst->getNumOperands() != 1) {
 297         return nullptr;
 298       } else {
 299         Operand = Inst->getOperand(0);
 300       }
 301     }
 302   }
 303 }
 304
 305 // Conservatively decides whether the dependence set of 'Val1' includes the
 306 // dependence set of 'Val2'. If 'ExpandSecondValue' is false, we do not expand
 307 // 'Val2' and use that single value as its dependence set.
 308 // If it returns true, it means the dependence set of 'Val1' includes that of
 309 // 'Val2'; otherwise, it only means we cannot conclusively decide it.
 310 bool dependenceSetInclusion(Value* Val1, Value* Val2,
 311                             int Val1ExpandLevel = 2 * kDependenceDepth,
 312                             int Val2ExpandLevel = kDependenceDepth) {
 313   typedef SmallSet<Value*, 8> IncludingSet;
 314   typedef SmallSet<Value*, 4> IncludedSet;
 315
 316   IncludingSet DepSet1;
 317   IncludedSet DepSet2;
 318   // Look for more depths for the including set.
 319   recursivelyFindDependence(&DepSet1, Val1, false /*Insert all visited nodes*/,
 320                             Val1ExpandLevel);
 321   recursivelyFindDependence(&DepSet2, Val2, true /*Only insert leaf nodes*/,
 322                             Val2ExpandLevel);
 323
 324   auto set_inclusion = [](IncludingSet FullSet, IncludedSet Subset) {
 325     for (auto* Dep : Subset) {
 326       if (0 == FullSet.count(Dep)) {
 327         return false;
 328       }
 329     }
 330     return true;
 331   };
 332   bool inclusion = set_inclusion(DepSet1, DepSet2);
 333   DEBUG(dbgs() << "[dependenceSetInclusion]: " << inclusion << "\n");
 334   DEBUG(dbgs() << "Including set for: " << *Val1 << "\n");
 335   DEBUG(for (const auto* Dep : DepSet1) { dbgs() << "\t\t" << *Dep << "\n"; });
 336   DEBUG(dbgs() << "Included set for: " << *Val2 << "\n");
 337   DEBUG(for (const auto* Dep : DepSet2) { dbgs() << "\t\t" << *Dep << "\n"; });
 338
 339   return inclusion;
 340 }
 341
 342 // Recursively iterates through the operands spawned from 'DepVal'. If there
 343 // exists a single value that 'DepVal' only depends on, we call that value the
 344 // root dependence of 'DepVal' and return it. Otherwise, return 'DepVal'.
 345 Value* getRootDependence(Value* DepVal) {
 346   SmallSet<Value*, 8> DepSet;
 347   for (unsigned depth = kDependenceDepth; depth > 0; --depth) {
 348     recursivelyFindDependence(&DepSet, DepVal, true /*Only insert leaf nodes*/,
 349                               depth);
 350     if (DepSet.size() == 1) {
 351       return *DepSet.begin();
 352     }
 353     DepSet.clear();
 354   }
 355   return DepVal;
 356 }
 357
 358 // This function actually taints 'DepVal' to the address to 'SI'. If the
 359 // address
 360 // of 'SI' already depends on whatever 'DepVal' depends on, this function
 361 // doesn't do anything and returns false. Otherwise, returns true.
 362 //
 363 // This effect forces the store and any stores that comes later to depend on
 364 // 'DepVal'. For example, we have a condition "cond", and a store instruction
 365 // "s: STORE addr, val". If we want "s" (and any later store) to depend on
 366 // "cond", we do the following:
 367 // %conv = sext i1 %cond to i32
 368 // %addrVal = ptrtoint i32* %addr to i32
 369 // %andCond = and i32 conv, 0;
 370 // %orAddr = or i32 %andCond, %addrVal;
 371 // %NewAddr = inttoptr i32 %orAddr to i32*;
 372 //
 373 // This is a more concrete example:
 374 // ------
 375 // %0 = load i32, i32* @y, align 4, !tbaa !1
 376 // %cmp = icmp ne i32 %0, 42  // <== this is like the condition
 377 // %1 = sext i1 %cmp to i32
 378 // %2 = ptrtoint i32* @x to i32
 379 // %3 = and i32 %1, 0
 380 // %4 = or i32 %3, %2
 381 // %5 = inttoptr i32 %4 to i32*
 382 // store i32 1, i32* %5, align 4
 383 bool taintStoreAddress(StoreInst* SI, Value* DepVal) {
 384   // Set the insertion point right after the 'DepVal'.
 385   IRBuilder<true, NoFolder> Builder(SI);
 386   BasicBlock* BB = SI->getParent();
 387   Value* Address = SI->getPointerOperand();
 388   Type* TargetIntegerType =
 389       IntegerType::get(Address->getContext(),
 390                        BB->getModule()->getDataLayout().getPointerSizeInBits());
 391
 392   // Does SI's address already depends on whatever 'DepVal' depends on?
 393   if (StoreAddressDependOnValue(SI, DepVal)) {
 394     return false;
 395   }
 396
 397   // Figure out if there's a root variable 'DepVal' depends on. For example, we
 398   // can extract "getelementptr inbounds %struct, %struct* %0, i64 0, i32 123"
 399   // to be "%struct* %0" since all other operands are constant.
 400   auto* RootVal = getRootDependence(DepVal);
 401   auto* RootInst = dyn_cast<Instruction>(RootVal);
 402   auto* DepValInst = dyn_cast<Instruction>(DepVal);
 403   if (RootInst && DepValInst &&
 404       RootInst->getParent() == DepValInst->getParent()) {
 405     DepVal = RootVal;
 406   }
 407
 408   // Is this already a dependence-tainted store?
 409   Value* OldDep = getDependence(Address);
 410   if (OldDep) {
 411     // The address of 'SI' has already been tainted.  Just need to absorb the
 412     // DepVal to the existing dependence in the address of SI.
 413     Instruction* AndDep = getAndDependence(Address);
 414     IRBuilder<true, NoFolder> Builder(AndDep);
 415     Value* NewDep = nullptr;
 416     if (DepVal->getType() == AndDep->getType()) {
 417       NewDep = Builder.CreateAnd(OldDep, DepVal);
 418     } else {
 419       NewDep = Builder.CreateAnd(
 420           OldDep, createCast(Builder, DepVal, TargetIntegerType));
 421     }
 422
 423     // Use the new AND instruction as the dependence
 424     AndDep->setOperand(0, NewDep);
 425     return true;
 426   }
 427
 428   // SI's address has not been tainted. Now taint it with 'DepVal'.
 429   Value* CastDepToInt = createCast(Builder, DepVal, TargetIntegerType);
 430   Value* PtrToIntCast = Builder.CreatePtrToInt(Address, TargetIntegerType);
 431   Value* AndDepVal =
 432       Builder.CreateAnd(CastDepToInt, ConstantInt::get(TargetIntegerType, 0));
 433   // XXX-comment: The original IR InstCombiner would change our and instruction
 434   // to a select and then the back end optimize the condition out.  We attach a
 435   // flag to instructions and set it here to inform the InstCombiner to not to
 436   // touch this and instruction at all.
 437   Value* OrAddr = Builder.CreateOr(AndDepVal, PtrToIntCast);
 438   Value* NewAddr = Builder.CreateIntToPtr(OrAddr, Address->getType());
 439
 440   DEBUG(dbgs() << "[taintStoreAddress]\n"
 441                << "Original store: " << *SI << '\n');
 442   SI->setOperand(1, NewAddr);
 443
 444   // Debug output.
 445   DEBUG(dbgs() << "\tTargetIntegerType: " << *TargetIntegerType << '\n'
 446                << "\tCast dependence value to integer: " << *CastDepToInt
 447                << '\n'
 448                << "\tCast address to integer: " << *PtrToIntCast << '\n'
 449                << "\tAnd dependence value: " << *AndDepVal << '\n'
 450                << "\tOr address: " << *OrAddr << '\n'
 451                << "\tCast or instruction to address: " << *NewAddr << "\n\n");
 452
 453   return true;
 454 }
 455
 456 // Given the load part result of a RMW 'LoadPart', taints the address of the store
 457 // exclusive 'Addr' and returns it.
 458 Value* taintRMWStoreAddressWithLoadPart(IRBuilder<>& Builder, Value* Address, Instruction* LoadPart) {
 459   auto* BB = LoadPart->getParent();
 460
 461   Type* TargetIntegerType =
 462       IntegerType::get(Address->getContext(),
 463                        BB->getModule()->getDataLayout().getPointerSizeInBits());
 464
 465   // SI's address has not been tainted. Now taint it with 'DepVal'.
 466   Value* CastDepToInt = createCast(Builder, LoadPart, TargetIntegerType);
 467   Value* PtrToIntCast = Builder.CreatePtrToInt(Address, TargetIntegerType);
 468   Value* AndDepVal =
 469       Builder.CreateAnd(CastDepToInt, ConstantInt::get(TargetIntegerType, 0));
 470   Value* OrAddr = Builder.CreateOr(AndDepVal, PtrToIntCast);
 471   Value* NewAddr = Builder.CreateIntToPtr(OrAddr, Address->getType());
 472
 473   DEBUG(dbgs() << "[taintStoreAddressOfRMWStorePart]\n"
 474                << "Original address: " << *Address << '\n');
 475
 476   // Debug output.
 477   DEBUG(dbgs() << "\tTargetIntegerType: " << *TargetIntegerType << '\n'
 478                << "\tCast dependence value to integer: " << *CastDepToInt
 479                << '\n'
 480                << "\tCast address to integer: " << *PtrToIntCast << '\n'
 481                << "\tAnd dependence value: " << *AndDepVal << '\n'
 482                << "\tOr address: " << *OrAddr << '\n'
 483                << "\tCast or instruction to address: " << *NewAddr << "\n\n");
 484
 485   return NewAddr;
 486 }
 487
 488 // Looks for the previous store in the if block --- 'BrBB', which makes the
 489 // speculative store 'StoreToHoist' safe.
 490 Value* getSpeculativeStoreInPrevBB(StoreInst* StoreToHoist, BasicBlock* BrBB) {
 491   assert(StoreToHoist && "StoreToHoist must be a real store");
 492
 493   Value* StorePtr = StoreToHoist->getPointerOperand();
 494
 495   // Look for a store to the same pointer in BrBB.
 496   for (BasicBlock::reverse_iterator RI = BrBB->rbegin(), RE = BrBB->rend();
 497        RI != RE; ++RI) {
 498     Instruction* CurI = &*RI;
 499
 500     StoreInst* SI = dyn_cast<StoreInst>(CurI);
 501     // Found the previous store make sure it stores to the same location.
 502     // XXX-update: If the previous store's original untainted address are the
 503     // same as 'StorePtr', we are also good to hoist the store.
 504     if (SI && (SI->getPointerOperand() == StorePtr ||
 505                GetUntaintedAddress(SI->getPointerOperand()) == StorePtr)) {
 506       // Found the previous store, return its value operand.
 507       return SI;
 508     }
 509   }
 510
 511   assert(false &&
 512          "We should not reach here since this store is safe to speculate");
 513 }
 514
 515 // XXX-comment: Returns true if it changes the code, false otherwise (the branch
 516 // condition already depends on 'DepVal'.
 517 bool taintConditionalBranch(BranchInst* BI, Value* DepVal) {
 518   assert(BI->isConditional());
 519   auto* Cond = BI->getOperand(0);
 520   if (dependenceSetInclusion(Cond, DepVal)) {
 521     // The dependence/ordering is self-evident.
 522     return false;
 523   }
 524
 525   IRBuilder<true, NoFolder> Builder(BI);
 526   auto* AndDep =
 527       Builder.CreateAnd(DepVal, ConstantInt::get(DepVal->getType(), 0));
 528   auto* TruncAndDep =
 529       Builder.CreateTrunc(AndDep, IntegerType::get(DepVal->getContext(), 1));
 530   auto* OrCond = Builder.CreateOr(TruncAndDep, Cond);
 531   BI->setOperand(0, OrCond);
 532
 533   // Debug output.
 534   DEBUG(dbgs() << "\tTainted branch condition:\n" << *BI->getParent());
 535
 536   return true;
 537 }
 538
 539 bool ConditionalBranchDependsOnValue(BranchInst* BI, Value* DepVal) {
 540   assert(BI->isConditional());
 541   auto* Cond = BI->getOperand(0);
 542   return dependenceSetInclusion(Cond, DepVal);
 543 }
 544
 545 // XXX-update: For an instruction (e.g., a relaxed load) 'Inst', find the first
 546 // immediate atomic store or the first conditional branch. Returns nullptr if
 547 // there's no such immediately following store/branch instructions, which we can
 548 // only enforce the load with 'acquire'. 'ChainedBB' contains all the blocks
 549 // chained together with unconditional branches from 'BB' to the block with the
 550 // first store/cond branch.
 551 template <typename Vector>
 552 Instruction* findFirstStoreCondBranchInst(Instruction* Inst, Vector* ChainedBB) {
 553   // In some situations, relaxed loads can be left as is:
 554   // 1. The relaxed load is used to calculate the address of the immediate
 555   // following store;
 556   // 2. The relaxed load is used as a condition in the immediate following
 557   // condition, and there are no stores in between. This is actually quite
 558   // common. E.g.,
 559   // int r1 = x.load(relaxed);
 560   // if (r1 != 0) {
 561   //   y.store(1, relaxed);
 562   // }
 563   // However, in this function, we don't deal with them directly. Instead, we
 564   // just find the immediate following store/condition branch and return it.
 565
 566   assert(ChainedBB != nullptr && "Chained BB should not be nullptr");
 567   auto* BB = Inst->getParent();
 568   ChainedBB->push_back(BB);
 569   auto BE = BB->end();
 570   auto BBI = BasicBlock::iterator(Inst);
 571   BBI++;
 572   while (true) {
 573     for (; BBI != BE; BBI++) {
 574       Instruction* Inst = &*BBI;
 575       IntrinsicInst* II = dyn_cast<IntrinsicInst>(&*BBI);
 576       if (II && II->getIntrinsicID() == Intrinsic::aarch64_stlxr) {
 577         return II;
 578       } else if (Inst->getOpcode() == Instruction::Store) {
 579         return Inst;
 580       } else if (Inst->getOpcode() == Instruction::Br) {
 581         auto* BrInst = dyn_cast<BranchInst>(Inst);
 582         if (BrInst->isConditional()) {
 583           return Inst;
 584         } else {
 585           // Reinitialize iterators with the destination of the unconditional
 586           // branch.
 587           BB = BrInst->getSuccessor(0);
 588           ChainedBB->push_back(BB);
 589           BBI = BB->begin();
 590           BE = BB->end();
 591           break;
 592         }
 593       }
 594     }
 595     if (BBI == BE) {
 596       return nullptr;
 597     }
 598   }
 599 }
 600
 601 // XXX-update: Find the next node of the last relaxed load from 'FromInst' to
 602 // 'ToInst'. If none, return 'ToInst'.
 603 Instruction* findLastLoadNext(Instruction* FromInst, Instruction* ToInst) {
 604   if (FromInst == ToInst) {
 605     return ToInst;
 606   }
 607   Instruction* LastLoad = ToInst;
 608   auto* BB = FromInst->getParent();
 609   auto BE = BB->end();
 610   auto BBI = BasicBlock::iterator(FromInst);
 611   BBI++;
 612   for (; BBI != BE && &*BBI != ToInst; BBI++) {
 613     auto* LI = dyn_cast<LoadInst>(&*BBI);
 614     if (LI == nullptr || !LI->isAtomic() || LI->getOrdering() != Monotonic) {
 615       continue;
 616     }
 617     LastLoad = LI;
 618     LastLoad = LastLoad->getNextNode();
 619   }
 620   return LastLoad;
 621 }
 622
 623 // Inserts a fake conditional branch right after the instruction 'SplitInst',
 624 // and the branch condition is 'Condition'. 'SplitInst' will be placed in the
 625 // newly created block.
 626 void AddFakeConditionalBranch(Instruction* SplitInst, Value* Condition) {
 627   auto* BB = SplitInst->getParent();
 628   TerminatorInst* ThenTerm = nullptr;
 629   TerminatorInst* ElseTerm = nullptr;
 630   SplitBlockAndInsertIfThenElse(Condition, SplitInst, &ThenTerm, &ElseTerm);
 631   assert(ThenTerm && ElseTerm &&
 632          "Then/Else terminators cannot be empty after basic block spliting");
 633   auto* ThenBB = ThenTerm->getParent();
 634   auto* ElseBB = ElseTerm->getParent();
 635   auto* TailBB = ThenBB->getSingleSuccessor();
 636   assert(TailBB && "Tail block cannot be empty after basic block spliting");
 637
 638   ThenBB->disableCanEliminateBlock();
 639   ThenBB->disableCanEliminateBlock();
 640   TailBB->disableCanEliminateBlock();
 641   ThenBB->setName(BB->getName() + "Then.Fake");
 642   ElseBB->setName(BB->getName() + "Else.Fake");
 643   DEBUG(dbgs() << "Add fake conditional branch:\n"
 644                << "Then Block:\n"
 645                << *ThenBB << "Else Block:\n"
 646                << *ElseBB << "\n");
 647 }
 648
 649 // Returns true if the code is changed, and false otherwise.
 650 void TaintRelaxedLoads(Instruction* UsageInst, Instruction* InsertPoint) {
 651   // For better performance, we can add a "AND X 0" instruction before the
 652   // condition.
 653   auto* BB = UsageInst->getParent();
 654   if (InsertPoint == nullptr) {
 655     InsertPoint = UsageInst->getNextNode();
 656   }
 657   // Insert instructions after PHI nodes.
 658   while (dyn_cast<PHINode>(InsertPoint)) {
 659     InsertPoint = InsertPoint->getNextNode();
 660   }
 661   // First thing is to cast 'UsageInst' to an integer type if necessary.
 662   Value* AndTarget = nullptr;
 663   Type* TargetIntegerType =
 664       IntegerType::get(UsageInst->getContext(),
 665                        BB->getModule()->getDataLayout().getPointerSizeInBits());
 666
 667   // Check whether InsertPoint is a added fake conditional branch.
 668   BranchInst* BI = nullptr;
 669   if ((BI = dyn_cast<BranchInst>(InsertPoint)) && BI->isConditional()) {
 670     auto* Cond = dyn_cast<Instruction>(BI->getOperand(0));
 671     if (Cond && Cond->getOpcode() == Instruction::ICmp) {
 672       auto* CmpInst = dyn_cast<ICmpInst>(Cond);
 673       auto* Op0 = dyn_cast<Instruction>(Cond->getOperand(0));
 674       auto* Op1 = dyn_cast<ConstantInt>(Cond->getOperand(1));
 675       // %tmp = And X, 0
 676       // %cmp = ICMP_NE %tmp, 0
 677       // Br %cmp
 678       // =>
 679       // %tmp1 = And X, NewTaintedVal
 680       // %tmp2 = And %tmp1, 0
 681       // %cmp = ICMP_NE %tmp2, 0
 682       // Br %cmp
 683       if (CmpInst && CmpInst->getPredicate() == CmpInst::ICMP_NE && Op0 &&
 684           Op0->getOpcode() == Instruction::And && Op1 && Op1->isZero()) {
 685         auto* Op01 = dyn_cast<ConstantInt>(Op0->getOperand(1));
 686         if (Op01 && Op01->isZero()) {
 687           // Now we have a previously added fake cond branch.
 688           auto* Op00 = Op0->getOperand(0);
 689           IRBuilder<true, NoFolder> Builder(CmpInst);
 690           if (Op00->getType() == UsageInst->getType()) {
 691             AndTarget = UsageInst;
 692           } else {
 693             AndTarget = createCast(Builder, UsageInst, Op00->getType());
 694           }
 695           AndTarget = Builder.CreateAnd(Op00, AndTarget);
 696           auto* AndZero = dyn_cast<Instruction>(Builder.CreateAnd(
 697               AndTarget, Constant::getNullValue(AndTarget->getType())));
 698           CmpInst->setOperand(0, AndZero);
 699           return;
 700         }
 701       }
 702     }
 703   }
 704
 705   IRBuilder<true, NoFolder> Builder(InsertPoint);
 706   if (IntegerType::classof(UsageInst->getType())) {
 707     AndTarget = UsageInst;
 708   } else {
 709     AndTarget = createCast(Builder, UsageInst, TargetIntegerType);
 710   }
 711   auto* AndZero = dyn_cast<Instruction>(
 712       Builder.CreateAnd(AndTarget, Constant::getNullValue(AndTarget->getType())));
 713   auto* FakeCondition = dyn_cast<Instruction>(Builder.CreateICmp(
 714       CmpInst::ICMP_NE, AndZero, Constant::getNullValue(AndTarget->getType())));
 715   AddFakeConditionalBranch(FakeCondition->getNextNode(), FakeCondition);
 716 }
 717
 718 // Taints the 'Inst' (i.e., adds a fake conditional block that uses the 'Inst')
 719 // at the beginning of basic block 'BB'. Note that we need to add an appropriate
 720 // PHI node and taint the PHI node. Returns true if the code is changed, and
 721 // false otherwise.
 722 void TaintAtBlockBeginning(Instruction* Inst, BasicBlock* BB) {
 723   auto* CurBB = Inst->getParent();
 724   auto* FirstInst = BB->getFirstNonPHI();
 725   IRBuilder<true, NoFolder> Builder(FirstInst);
 726   auto* Phi = Builder.CreatePHI(Inst->getType(), 0, Inst->getName() + ".phi");
 727   // Multiple blocks going to BB. We should add a PHI node w.r.t. 'Inst'.
 728   for (auto* Pred : predecessors(BB)) {
 729     Value* Val = nullptr;
 730     if (Pred == CurBB) {
 731       Val = Inst;
 732     } else {
 733       // We don't care what value other paths are.
 734       Val = UndefValue::get(Inst->getType());
 735     }
 736     Phi->addIncoming(Val, Pred);
 737   }
 738   return TaintRelaxedLoads(Phi, Phi);
 739 }
 740
 741 // XXX-comment: Finds the appropriate Value derived from an atomic load.
 742 // 'ChainedBB' contains all the blocks chained together with unconditional
 743 // branches from LI's parent BB to the block with the first store/cond branch.
 744 // If we don't find any, it means 'LI' is not used at all (which should not
 745 // happen in practice). We can simply set 'LI' to be acquire just to be safe.
 746 template <typename Vector>
 747 Instruction* findMostRecentDependenceUsage(LoadInst* LI, Instruction* LaterInst,
 748                                            Vector* ChainedBB,
 749                                            DominatorTree* DT) {
 750   typedef SmallSet<Instruction*, 8> UsageSet;
 751   typedef DenseMap<BasicBlock*, std::unique_ptr<UsageSet>> UsageMap;
 752   assert(ChainedBB->size() >= 1 && "ChainedBB must have >=1 size");
 753   // Mapping from basic block in 'ChainedBB' to the set of dependence usage of
 754   // 'LI' in each block.
 755   UsageMap usage_map;
 756   auto* LoadBB = LI->getParent();
 757   usage_map[LoadBB] = make_unique<UsageSet>();
 758   usage_map[LoadBB]->insert(LI);
 759
 760   for (auto* BB : *ChainedBB) {
 761     if (usage_map[BB] == nullptr) {
 762       usage_map[BB] = make_unique<UsageSet>();
 763     }
 764     auto& usage_set = usage_map[BB];
 765     if (usage_set->size() == 0) {
 766       // The value has not been used.
 767       return nullptr;
 768     }
 769     // Calculate the usage in the current BB first.
 770     std::list<Value*> bb_usage_list;
 771     std::copy(usage_set->begin(), usage_set->end(),
 772               std::back_inserter(bb_usage_list));
 773     for (auto list_iter = bb_usage_list.begin();
 774          list_iter != bb_usage_list.end(); list_iter++) {
 775       auto* val = *list_iter;
 776       for (auto* U : val->users()) {
 777         Instruction* Inst = nullptr;
 778         if (!(Inst = dyn_cast<Instruction>(U))) {
 779           continue;
 780         }
 781         assert(Inst && "Usage value must be an instruction");
 782         auto iter =
 783             std::find(ChainedBB->begin(), ChainedBB->end(), Inst->getParent());
 784         if (iter == ChainedBB->end()) {
 785           // Only care about usage within ChainedBB.
 786           continue;
 787         }
 788         auto* UsageBB = *iter;
 789         if (UsageBB == BB) {
 790           // Current BB.
 791           if (!usage_set->count(Inst)) {
 792             bb_usage_list.push_back(Inst);
 793             usage_set->insert(Inst);
 794           }
 795         } else {
 796           // A following BB.
 797           if (usage_map[UsageBB] == nullptr) {
 798             usage_map[UsageBB] = make_unique<UsageSet>();
 799           }
 800           usage_map[UsageBB]->insert(Inst);
 801         }
 802       }
 803     }
 804   }
 805
 806   // Pick one usage that is in LaterInst's block and that dominates 'LaterInst'.
 807   auto* LaterBB = LaterInst->getParent();
 808   auto& usage_set = usage_map[LaterBB];
 809   Instruction* usage_inst = nullptr;
 810   for (auto* inst : *usage_set) {
 811     if (DT->dominates(inst, LaterInst)) {
 812       usage_inst = inst;
 813       break;
 814     }
 815   }
 816
 817   assert(usage_inst && "The usage instruction in the same block but after the "
 818                        "later instruction");
 819   return usage_inst;
 820 }
 821
 822 // XXX-comment: For an instruction (e.g., a load) 'Inst', and the first upcoming
 823 // store/conditional branch instruction 'FirstInst', returns whether there are
 824 // any intermediate instructions I (including 'FirstInst') that satisfy:
 825 // 1. I is a load/store, and its address depends on 'Inst'.
 826 // 2. I is a conditional branch whose condition depends on 'Inst'.
 827 // Note that 'Inst' and 'FirstInst' can be in different basic blocks, but Inst's
 828 // basic block can unconditionally jumps (by steps) to FirstInst's block.
 829 bool NeedExtraConstraints(Instruction* Inst, Instruction* FirstInst) {
 830   if (!FirstInst) {
 831     return true;
 832   }
 833   auto BBI = Inst->getIterator();
 834   BBI++;
 835   while (true) {
 836     auto* I = &*BBI;
 837     BBI++;
 838     BranchInst *BI = dyn_cast<BranchInst>(I);
 839     if (BI && BI->isUnconditional()) {
 840       BasicBlock *DestBB = BI->getSuccessor(0);
 841       BBI = DestBB->begin();
 842       continue;
 843     }
 844
 845     if (I->getOpcode() == Instruction::Store) {
 846       return !StoreAddressDependOnValue(dyn_cast<StoreInst>(I), Inst);
 847     } else if (I->getOpcode() == Instruction::Load) {
 848       if (I->isAtomic() &&
 849           LoadAddressDependOnValue(dyn_cast<LoadInst>(I), Inst)) {
 850         // Normal loads are subject to be reordered by the backend, so we only
 851         // rely on atomic loads.
 852         return false;
 853       }
 854     } else if (I->getOpcode() == Instruction::Br) {
 855       return !ConditionalBranchDependsOnValue(dyn_cast<BranchInst>(I), Inst);
 856     }
 857     if (I == FirstInst) {
 858       return true;
 859     }
 860   }
 861   return true;
 862 }
 863
 864 // XXX-comment: For an instruction (e.g., a load) 'Inst', returns whether there
 865 // are any intermediate instructions I (including 'FirstInst') that satisfy:
 866 // 1. There are no reachable store/conditional branch before 'I'.
 867 // 2. I is a load/store, and its address depends on 'Inst'.
 868 // 3. I is a conditional branch whose condition depends on 'Inst'.
 869 // Note that 'Inst' and 'FirstInst' can be in different basic blocks, but Inst's
 870 // basic block can unconditionally jumps (by steps) to FirstInst's block.
 871 bool NeedExtraConstraints(Instruction* Inst) {
 872   SmallVector<BasicBlock*, 2> ChainedBB;
 873   auto* FirstInst = findFirstStoreCondBranchInst(Inst, &ChainedBB);
 874   return NeedExtraConstraints(Inst, FirstInst);
 875 }
 876
 877 // XXX-comment: Returns whether the code has been changed.
 878 bool AddFakeConditionalBranchAfterMonotonicLoads(
 879     SmallSet<LoadInst*, 1>& MonotonicLoadInsts, DominatorTree* DT) {
 880   bool Changed = false;
 881   while (!MonotonicLoadInsts.empty()) {
 882     auto* LI = *MonotonicLoadInsts.begin();
 883     MonotonicLoadInsts.erase(LI);
 884     SmallVector<BasicBlock*, 2> ChainedBB;
 885     auto* FirstInst = findFirstStoreCondBranchInst(LI, &ChainedBB);
 886
 887     // First check whether existing load-store ordering constraints exist.
 888     if (FirstInst != nullptr && !NeedExtraConstraints(LI, FirstInst)) {
 889       continue;
 890     }
 891
 892     // We really need to process the relaxed load now. First see if we can delay
 893     // the tainting.
 894     if (FirstInst) {
 895       auto* FirstInstBBTerm = FirstInst->getParent()->getTerminator();
 896       while (FirstInst != FirstInstBBTerm) {
 897         if (!CanDelayTainting(LI, FirstInst)) {
 898           break;
 899         }
 900         FirstInst = FirstInst->getNextNode();
 901       }
 902     }
 903
 904     StoreInst* SI = nullptr;
 905     IntrinsicInst* II = nullptr;
 906     if (FirstInst) {
 907       SI = dyn_cast<StoreInst>(FirstInst);
 908       II = dyn_cast<IntrinsicInst>(FirstInst);
 909     }
 910     if (FirstInst &&
 911         (SI || (II && II->getIntrinsicID() == Intrinsic::aarch64_stlxr))) {
 912       // For immediately coming stores, taint the address of the store.
 913       if (FirstInst->getParent() == LI->getParent() ||
 914           DT->dominates(LI, FirstInst)) {
 915         TaintRelaxedLoads(LI, FirstInst);
 916         Changed = true;
 917       } else {
 918         auto* Inst =
 919             findMostRecentDependenceUsage(LI, FirstInst, &ChainedBB, DT);
 920         if (!Inst) {
 921           LI->setOrdering(Acquire);
 922           Changed = true;
 923         } else {
 924           TaintRelaxedLoads(Inst, FirstInst);
 925           Changed = true;
 926         }
 927       }
 928     } else {
 929       // No upcoming branch
 930       if (!FirstInst) {
 931         TaintRelaxedLoads(LI, nullptr);
 932         Changed = true;
 933       } else {
 934         // For immediately coming branch, directly add a fake branch.
 935         if (FirstInst->getParent() == LI->getParent() ||
 936             DT->dominates(LI, FirstInst)) {
 937           TaintRelaxedLoads(LI, FirstInst);
 938           Changed = true;
 939         } else {
 940           auto* Inst =
 941               findMostRecentDependenceUsage(LI, FirstInst, &ChainedBB, DT);
 942           if (Inst) {
 943             TaintRelaxedLoads(Inst, FirstInst);
 944           } else {
 945             LI->setOrdering(Acquire);
 946           }
 947           Changed = true;
 948         }
 949       }
 950     }
 951   }
 952   return Changed;
 953 }
 954
 955 /**** Implementations of public methods for dependence tainting ****/
 956 Value* GetUntaintedAddress(Value* CurrentAddress) {
 957   auto* OrAddress = getOrAddress(CurrentAddress);
 958   if (OrAddress == nullptr) {
 959     // Is it tainted by a select instruction?
 960     auto* Inst = dyn_cast<Instruction>(CurrentAddress);
 961     if (nullptr != Inst && Inst->getOpcode() == Instruction::Select) {
 962       // A selection instruction.
 963       if (Inst->getOperand(1) == Inst->getOperand(2)) {
 964         return Inst->getOperand(1);
 965       }
 966     }
 967
 968     return CurrentAddress;
 969   }
 970
 971   auto* CastToInt = dyn_cast<Instruction>(OrAddress->getOperand(1));
 972   if (CastToInt && CastToInt->getOpcode() == Instruction::PtrToInt) {
 973     return CastToInt->getOperand(0);
 974   } else {
 975     // This should be a IntToPtr constant expression.
 976     ConstantExpr* PtrToIntExpr =
 977         dyn_cast<ConstantExpr>(OrAddress->getOperand(1));
 978     if (PtrToIntExpr && PtrToIntExpr->getOpcode() == Instruction::PtrToInt) {
 979       return PtrToIntExpr->getOperand(0);
 980     }
 981   }
 982
 983   // Looks like it's not been dependence-tainted. Returns itself.
 984   return CurrentAddress;
 985 }
 986
 987 MemoryLocation GetUntaintedMemoryLocation(StoreInst* SI) {
 988   AAMDNodes AATags;
 989   SI->getAAMetadata(AATags);
 990   const auto& DL = SI->getModule()->getDataLayout();
 991   const auto* OriginalAddr = GetUntaintedAddress(SI->getPointerOperand());
 992   DEBUG(if (OriginalAddr != SI->getPointerOperand()) {
 993     dbgs() << "[GetUntaintedMemoryLocation]\n"
 994            << "Storing address: " << *SI->getPointerOperand()
 995            << "\nUntainted address: " << *OriginalAddr << "\n";
 996   });
 997   return MemoryLocation(OriginalAddr,
 998                         DL.getTypeStoreSize(SI->getValueOperand()->getType()),
 999                         AATags);
1000 }
1001
1002 bool TaintDependenceToStore(StoreInst* SI, Value* DepVal) {
1003   if (dependenceSetInclusion(SI, DepVal)) {
1004     return false;
1005   }
1006
1007   bool tainted = taintStoreAddress(SI, DepVal);
1008   assert(tainted);
1009   return tainted;
1010 }
1011
1012 bool TaintDependenceToStoreAddress(StoreInst* SI, Value* DepVal) {
1013   if (dependenceSetInclusion(SI->getPointerOperand(), DepVal)) {
1014     return false;
1015   }
1016
1017   bool tainted = taintStoreAddress(SI, DepVal);
1018   assert(tainted);
1019   return tainted;
1020 }
1021
1022 bool CompressTaintedStore(BasicBlock* BB) {
1023   // This function looks for windows of adajcent stores in 'BB' that satisfy the
1024   // following condition (and then do optimization):
1025   // *Addr(d1) = v1, d1 is a condition and is the only dependence the store's
1026   //                 address depends on && Dep(v1) includes Dep(d1);
1027   // *Addr(d2) = v2, d2 is a condition and is the only dependnece the store's
1028   //                 address depends on && Dep(v2) includes Dep(d2) &&
1029   //                 Dep(d2) includes Dep(d1);
1030   // ...
1031   // *Addr(dN) = vN, dN is a condition and is the only dependence the store's
1032   //                 address depends on && Dep(dN) includes Dep(d"N-1").
1033   //
1034   // As a result, Dep(dN) includes [Dep(d1) V ... V Dep(d"N-1")], so we can
1035   // safely transform the above to the following. In between these stores, we
1036   // can omit untainted stores to the same address 'Addr' since they internally
1037   // have dependence on the previous stores on the same address.
1038   // =>
1039   // *Addr = v1
1040   // *Addr = v2
1041   // *Addr(d3) = v3
1042   for (auto BI = BB->begin(), BE = BB->end(); BI != BE; BI++) {
1043     // Look for the first store in such a window of adajacent stores.
1044     auto* FirstSI = dyn_cast<StoreInst>(&*BI);
1045     if (!FirstSI) {
1046       continue;
1047     }
1048
1049     // The first store in the window must be tainted.
1050     auto* UntaintedAddress = GetUntaintedAddress(FirstSI->getPointerOperand());
1051     if (UntaintedAddress == FirstSI->getPointerOperand()) {
1052       continue;
1053     }
1054
1055     // The first store's address must directly depend on and only depend on a
1056     // condition.
1057     auto* FirstSIDepCond = getConditionDependence(FirstSI->getPointerOperand());
1058     if (nullptr == FirstSIDepCond) {
1059       continue;
1060     }
1061
1062     // Dep(first store's storing value) includes Dep(tainted dependence).
1063     if (!dependenceSetInclusion(FirstSI->getValueOperand(), FirstSIDepCond)) {
1064       continue;
1065     }
1066
1067     // Look for subsequent stores to the same address that satisfy the condition
1068     // of "compressing the dependence".
1069     SmallVector<StoreInst*, 8> AdajacentStores;
1070     AdajacentStores.push_back(FirstSI);
1071     auto BII = BasicBlock::iterator(FirstSI);
1072     for (BII++; BII != BE; BII++) {
1073       auto* CurrSI = dyn_cast<StoreInst>(&*BII);
1074       if (!CurrSI) {
1075         if (BII->mayHaveSideEffects()) {
1076           // Be conservative. Instructions with side effects are similar to
1077           // stores.
1078           break;
1079         }
1080         continue;
1081       }
1082
1083       auto* OrigAddress = GetUntaintedAddress(CurrSI->getPointerOperand());
1084       auto* CurrSIDepCond = getConditionDependence(CurrSI->getPointerOperand());
1085       // All other stores must satisfy either:
1086       // A. 'CurrSI' is an untainted store to the same address, or
1087       // B. the combination of the following 5 subconditions:
1088       // 1. Tainted;
1089       // 2. Untainted address is the same as the group's address;
1090       // 3. The address is tainted with a sole value which is a condition;
1091       // 4. The storing value depends on the condition in 3.
1092       // 5. The condition in 3 depends on the previous stores dependence
1093       // condition.
1094
1095       // Condition A. Should ignore this store directly.
1096       if (OrigAddress == CurrSI->getPointerOperand() &&
1097           OrigAddress == UntaintedAddress) {
1098         continue;
1099       }
1100       // Check condition B.
1101       if (OrigAddress == CurrSI->getPointerOperand() ||
1102           OrigAddress != UntaintedAddress || CurrSIDepCond == nullptr ||
1103           !dependenceSetInclusion(CurrSI->getValueOperand(), CurrSIDepCond)) {
1104         // Check condition 1, 2, 3 & 4.
1105         break;
1106       }
1107
1108       // Check condition 5.
1109       StoreInst* PrevSI = AdajacentStores[AdajacentStores.size() - 1];
1110       auto* PrevSIDepCond = getConditionDependence(PrevSI->getPointerOperand());
1111       assert(PrevSIDepCond &&
1112              "Store in the group must already depend on a condtion");
1113       if (!dependenceSetInclusion(CurrSIDepCond, PrevSIDepCond)) {
1114         break;
1115       }
1116
1117       AdajacentStores.push_back(CurrSI);
1118     }
1119
1120     if (AdajacentStores.size() == 1) {
1121       // The outer loop should keep looking from the next store.
1122       continue;
1123     }
1124
1125     // Now we have such a group of tainted stores to the same address.
1126     DEBUG(dbgs() << "[CompressTaintedStore]\n");
1127     DEBUG(dbgs() << "Original BB\n");
1128     DEBUG(dbgs() << *BB << '\n');
1129     auto* LastSI = AdajacentStores[AdajacentStores.size() - 1];
1130     for (unsigned i = 0; i < AdajacentStores.size() - 1; ++i) {
1131       auto* SI = AdajacentStores[i];
1132
1133       // Use the original address for stores before the last one.
1134       SI->setOperand(1, UntaintedAddress);
1135
1136       DEBUG(dbgs() << "Store address has been reversed: " << *SI << '\n';);
1137     }
1138     // XXX-comment: Try to make the last store use fewer registers.
1139     // If LastSI's storing value is a select based on the condition with which
1140     // its address is tainted, transform the tainted address to a select
1141     // instruction, as follows:
1142     // r1 = Select Cond ? A : B
1143     // r2 = Cond & 0
1144     // r3 = Addr | r2
1145     // *r3 = r1
1146     // ==>
1147     // r1 = Select Cond ? A : B
1148     // r2 = Select Cond ? Addr : Addr
1149     // *r2 = r1
1150     // The idea is that both Select instructions depend on the same condition,
1151     // so hopefully the backend can generate two cmov instructions for them (and
1152     // this saves the number of registers needed).
1153     auto* LastSIDep = getConditionDependence(LastSI->getPointerOperand());
1154     auto* LastSIValue = dyn_cast<Instruction>(LastSI->getValueOperand());
1155     if (LastSIValue && LastSIValue->getOpcode() == Instruction::Select &&
1156         LastSIValue->getOperand(0) == LastSIDep) {
1157       // XXX-comment: Maybe it's better for us to just leave it as an and/or
1158       // dependence pattern.
1159       //      /*
1160       IRBuilder<true, NoFolder> Builder(LastSI);
1161       auto* Address =
1162           Builder.CreateSelect(LastSIDep, UntaintedAddress, UntaintedAddress);
1163       LastSI->setOperand(1, Address);
1164       DEBUG(dbgs() << "The last store becomes :" << *LastSI << "\n\n";);
1165       //      */
1166     }
1167   }
1168
1169   return true;
1170 }
1171
1172 bool PassDependenceToStore(Value* OldAddress, StoreInst* NewStore) {
1173   Value* OldDep = getDependence(OldAddress);
1174   // Return false when there's no dependence to pass from the OldAddress.
1175   if (!OldDep) {
1176     return false;
1177   }
1178
1179   // No need to pass the dependence to NewStore's address if it already depends
1180   // on whatever 'OldAddress' depends on.
1181   if (StoreAddressDependOnValue(NewStore, OldDep)) {
1182     return false;
1183   }
1184   return taintStoreAddress(NewStore, OldAddress);
1185 }
1186
1187 SmallSet<Value*, 8> FindDependence(Value* Val) {
1188   SmallSet<Value*, 8> DepSet;
1189   recursivelyFindDependence(&DepSet, Val, true /*Only insert leaf nodes*/);
1190   return DepSet;
1191 }
1192
1193 bool StoreAddressDependOnValue(StoreInst* SI, Value* DepVal) {
1194   return dependenceSetInclusion(SI->getPointerOperand(), DepVal);
1195 }
1196
1197 bool LoadAddressDependOnValue(LoadInst* LI, Value* DepVal) {
1198   return dependenceSetInclusion(LI->getPointerOperand(), DepVal);
1199 }
1200
1201 bool StoreDependOnValue(StoreInst* SI, Value* Dep) {
1202   return dependenceSetInclusion(SI, Dep);
1203 }
1204
1205 bool ValueDependOnValue(Value* Inst, Value* Dep) {
1206   return dependenceSetInclusion(Inst, Dep);
1207 }
1208
1209 // XXX-update: Checks whether the relaxed load 'LI' has subsequent instructions
1210 // that naturally prevents it from being reordered across later stores.
1211 bool HasSubsequentOrderingProtection(LoadInst* LI) {
1212   auto* BB = LI->getParent();
1213   auto* Term = BB->getTerminator();
1214   for (auto Iter = BasicBlock::iterator(LI->getNextNode()); Iter != BB->end();
1215        Iter++) {
1216     Instruction* I = &*Iter;
1217
1218     // Reaching the end of the block.
1219     if (I == Term) {
1220       auto* Branch = dyn_cast<BranchInst>(Term);
1221       // The last instruction isn't a branch, end of analysis.
1222       if (!Branch) {
1223         return false;
1224       }
1225       if (Branch->isConditional()) {
1226         if (ValueDependOnValue(Branch, LI)) {
1227           // 'LI' is used in the conditional branch.
1228           return true;
1229         } else {
1230           // Reach the end with a cond branch that doesn't use the result of
1231           // 'LI'.
1232           return false;
1233         }
1234       } else {
1235         // Reach the end with a unconditional branch, keep going to the next
1236         // block.
1237         BB = BB->getSingleSuccessor();
1238         Term = BB->getTerminator();
1239         Iter = BB->begin();
1240         continue;
1241       }
1242     }
1243
1244     // 'I' is a CAS whose old value depends on 'LI'. We don't need to taint 'LI'
1245     // then.
1246     auto* CAS = dyn_cast<AtomicCmpXchgInst>(I);
1247     if (CAS) {
1248       if (ValueDependOnValue(CAS->getCompareOperand(), LI)) {
1249         return true;
1250       }
1251     }
1252
1253     // fetch_* operations that have acquire-release semantics.
1254     auto* RMW = dyn_cast<AtomicRMWInst>(I);
1255     if (RMW) {
1256       auto Order = RMW->getOrdering();
1257       if (Order == AcquireRelease || Order == SequentiallyConsistent) {
1258         return true;
1259       }
1260     }
1261
1262     // A load whose address depends on 'LI' prevents later stores from being
1263     // reordered.
1264     auto* LdInst = dyn_cast<LoadInst>(I);
1265     if (LdInst) {
1266       if (ValueDependOnValue(LdInst->getPointerOperand(), LI)) {
1267         return true;
1268       }
1269     }
1270
1271     // Other instructions that don't affect the reordering.
1272     if (!I->mayHaveSideEffects()) {
1273       continue;
1274     }
1275
1276     // A store whose address depends on 'LI' is also protection.
1277     auto* SI = dyn_cast<StoreInst>(I);
1278     if (SI) {
1279       if (ValueDependOnValue(SI->getPointerOperand(), LI)) {
1280         return true;
1281       }
1282     }
1283
1284     // The following are store/store-like operations. They don't protect later
1285     // stores from being reordered across 'LI', but the analysis can go on if
1286     // they naturally can't be reordered across 'LI' themselves.
1287     {
1288       // Release (or stronger) store.
1289       if (SI) {
1290         auto Order = SI->getOrdering();
1291         if (Order == Release || Order == SequentiallyConsistent) {
1292           continue;
1293         }
1294       }
1295
1296       // Release (or stronger) fetch_*.
1297       if (RMW) {
1298         auto Order = RMW->getOrdering();
1299         if (Order == Release || Order == AcquireRelease ||
1300             Order == SequentiallyConsistent) {
1301           continue;
1302         }
1303       }
1304
1305       // The instruction naturally depends on 'LI'.
1306       if (ValueDependOnValue(I, LI)) {
1307         continue;
1308       }
1309     }
1310     // Otherwise, we need to taint 'LI'.
1311     // XXX-comment: It may be a good idea that we can delay the fake conditional
1312     // branch down to this instruction.
1313     return false;
1314   }
1315
1316   // Just in case, the loop should never end without reaching a return.
1317   return false;
1318 }
1319
1320 // XXX-update: Checks whether the tainting to instruction 'I' can be delayed
1321 // with respects to the relaxed load 'LI'. This usually means 'I' itself already
1322 // depends on the 'LI' or 'I' is a store/store-like atomic operation that has
1323 // release semantics.
1324 bool CanDelayTainting(LoadInst* LI, Instruction* I) {
1325   if (I == I->getParent()->getTerminator()) {
1326     return false;
1327   }
1328
1329   if (!I->mayHaveSideEffects()) {
1330     return true;
1331   }
1332
1333   // The following are store/store-like operations. They don't protect later
1334   // stores from being reordered across 'LI', but the analysis can go on if
1335   // they naturally can't be reordered across 'LI' themselves.
1336
1337   // Release (or stronger) store.
1338   auto* SI = dyn_cast<StoreInst>(I);
1339   if (SI) {
1340     auto Order = SI->getOrdering();
1341     if (Order == Release || Order == SequentiallyConsistent) {
1342       return true;
1343     }
1344   }
1345
1346   // Release (or stronger) fetch_*.
1347   auto* RMW = dyn_cast<AtomicRMWInst>(I);
1348   if (RMW) {
1349     auto Order = RMW->getOrdering();
1350     if (Order == Release || Order == AcquireRelease ||
1351         Order == SequentiallyConsistent) {
1352       return true;
1353     }
1354   }
1355
1356   // The instruction naturally depends on 'LI'.
1357   if (ValueDependOnValue(I, LI)) {
1358     return true;
1359   }
1360
1361   // Otherwise, be conservative and say no!
1362   return false;
1363 }
1364
1365 } // namespace llvm