Add asan hooks to Recycler and ArrayRecycler

[nhaehnle]

Summary

Already fix some use-after-frees.

[nhaehnle]

Extract the MergeStore part to D25601.

include/llvm/CodeGen/SelectionDAGNodes.h

public:
   const DebugLoc &getDebugLoc() const { return DL; }
 };
 
 
 // Define inline functions from the SDValue class.
 
 inline SDValue::SDValue(SDNode *node, unsigned resno)
     : Node(node), ResNo(resno) {
-  assert((!Node || ResNo < Node->getNumValues()) &&
+  // Explicitly check for !ResNo to avoid use-after-free, because there are
+  // callers that use SDValue(N, 0) with a deleted N to indicate successful
+  // combines.
+  assert((!Node || !ResNo || ResNo < Node->getNumValues()) &&
          "Invalid result number for the given node!");
   assert(ResNo < -2U && "Cannot use result numbers reserved for DenseMaps.");
 }
 
 inline unsigned SDValue::getOpcode() const {
   return Node->getOpcode();
 }
 inline EVT SDValue::getValueType() const {

include/llvm/Support/ArrayRecycler.h

class ArrayRecycler {
   // Remove an entry from the free list in Bucket[Idx] and return it.
   // Return NULL if no entries are available.
   T *pop(unsigned Idx) {
     if (Idx >= Bucket.size())
       return nullptr;
     FreeList *Entry = Bucket[Idx];
     if (!Entry)
       return nullptr;
+    __asan_unpoison_memory_region(Entry,
+                                  sizeof(T) * Capacity::get(Idx).getSize());
     Bucket[Idx] = Entry->Next;
     return reinterpret_cast<T*>(Entry);
   }
 
   // Add an entry to the free list at Bucket[Idx].
   void push(unsigned Idx, T *Ptr) {
     assert(Ptr && "Cannot recycle NULL pointer");
     FreeList *Entry = reinterpret_cast<FreeList*>(Ptr);
     if (Idx >= Bucket.size())
       Bucket.resize(size_t(Idx) + 1);
     Entry->Next = Bucket[Idx];
     Bucket[Idx] = Entry;
+    __asan_poison_memory_region(Entry,
+                                sizeof(T) * Capacity::get(Idx).getSize());
   }
 
 public:
   /// The size of an allocated array is represented by a Capacity instance.
   ///
   /// This class is much smaller than a size_t, and it provides methods to work
   /// with the set of legal array capacities.
   class Capacity {

include/llvm/Support/Recycler.h

struct FreeNode {
     FreeNode *Next;
   };
 
   /// List of nodes that have deleted contents and are not in active use.
   FreeNode *FreeList = nullptr;
 
   FreeNode *pop_val() {
     auto *Val = FreeList;
+    // Always unpoison the full Size, because the combiner relies on being able
+    // to morph SDNodes into MachineSDNodes.
+    __asan_unpoison_memory_region(FreeList, Size);
     FreeList = FreeList->Next;
     return Val;
   }
 
   void push(FreeNode *N) {
     N->Next = FreeList;
     FreeList = N;
+    __asan_poison_memory_region(FreeList, Size);
   }
 
 public:
   ~Recycler() {
     // If this fails, either the callee has lost track of some allocation,
     // or the callee isn't tracking allocations and should just call
     // clear() before deleting the Recycler.
     assert(!FreeList && "Non-empty recycler deleted!");

lib/CodeGen/SelectionDAG/DAGCombiner.cpp

private:
     /// need to be narrowed in order to match.
     bool isAndLoadExtLoad(ConstantSDNode *AndC, LoadSDNode *LoadN,
                           EVT LoadResultTy, EVT &ExtVT, EVT &LoadedVT,
                           bool &NarrowLoad);
 
     /// This is a helper function for MergeConsecutiveStores. When the source
     /// elements of the consecutive stores are all constants or all extracted
     /// vector elements, try to merge them into one larger store.
-    /// \return True if a merged store was created.
-    bool MergeStoresOfConstantsOrVecElts(SmallVectorImpl<MemOpLink> &StoreNodes,
-                                         EVT MemVT, unsigned NumStores,
-                                         bool IsConstantSrc, bool UseVector);
+    /// \return number of stores that were merged into a merged store (always
+    /// a prefix of \p StoreNode).
+    unsigned MergeStoresOfConstantsOrVecElts(
+        SmallVectorImpl<MemOpLink> &StoreNodes, EVT MemVT, unsigned NumStores,
 
     /// This is a helper function for MergeConsecutiveStores.
     /// Stores that may be merged are placed in StoreNodes.
     /// Loads that may alias with those stores are placed in AliasLoadNodes.
     void getStoreMergeAndAliasCandidates(
         StoreSDNode* St, SmallVectorImpl<MemOpLink> &StoreNodes,
         SmallVectorImpl<LSBaseSDNode*> &AliasLoadNodes);
 
     /// Helper function for MergeConsecutiveStores. Checks if
     /// Candidate stores have indirect dependency through their
     /// operands. \return True if safe to merge
     bool checkMergeStoreCandidatesForDependencies(
         SmallVectorImpl<MemOpLink> &StoreNodes);
 
     /// Merge consecutive store operations into a wide store.
     /// This optimization uses wide integers or vectors when possible.
-    /// \return True if some memory operations were changed.
-    bool MergeConsecutiveStores(StoreSDNode *N);
+    /// \return number of stores that were merged into a merged store (the
+    /// affected nodes are stored as a prefix in \p StoreNodes).
+    unsigned MergeConsecutiveStores(StoreSDNode *N,
+                                    SmallVectorImpl<MemOpLink> &StoreNodes);
 
     /// \brief Try to transform a truncation where C is a constant:
     ///     (trunc (and X, C)) -> (and (trunc X), (trunc C))
     ///
     /// \p N needs to be a truncation and its first operand an AND. Other
     /// requirements are checked by the function (e.g. that trunc is
     /// single-use) and if missed an empty SDValue is returned.
     SDValue distributeTruncateThroughAnd(SDNode *N);
if ((Opcode == ISD::SDIV) || (Opcode == ISD::UDIV)) {
     if (TLI.isOperationLegalOrCustom(OtherOpcode, VT))
       return SDValue();
   }
 
   SDValue Op0 = Node->getOperand(0);
   SDValue Op1 = Node->getOperand(1);
   SDValue combined;
   for (SDNode::use_iterator UI = Op0.getNode()->use_begin(),
-         UE = Op0.getNode()->use_end(); UI != UE; ++UI) {
-    SDNode *User = *UI;
+         UE = Op0.getNode()->use_end(); UI != UE;) {
+    SDNode *User = *UI++;
     if (User == Node || User->use_empty())
       continue;
     // Convert the other matching node(s), too;
     // otherwise, the DIVREM may get target-legalized into something
     // target-specific that we won't be able to recognize.
     unsigned UserOpc = User->getOpcode();
     if ((UserOpc == Opcode || UserOpc == OtherOpcode || UserOpc == DivRemOpc) &&
         User->getOperand(0) == Op0 &&
for (unsigned I = 0, E = Ty.getVectorNumElements(); I != E; ++I) {
     StoreSDNode *St = cast<StoreSDNode>(Stores[I].MemNode);
     Chains.push_back(St->getChain());
     BuildVector.push_back(St->getValue());
   }
 
   return DAG.getBuildVector(Ty, SL, BuildVector);
 }
 
-bool DAGCombiner::MergeStoresOfConstantsOrVecElts(
+unsigned DAGCombiner::MergeStoresOfConstantsOrVecElts(
                   SmallVectorImpl<MemOpLink> &StoreNodes, EVT MemVT,
                   unsigned NumStores, bool IsConstantSrc, bool UseVector) {
   // Make sure we have something to merge.
   if (NumStores < 2)
-    return false;
+    return 0;
 
   int64_t ElementSizeBytes = MemVT.getSizeInBits() / 8;
   LSBaseSDNode *FirstInChain = StoreNodes[0].MemNode;
   unsigned LatestNodeUsed = 0;
 
   for (unsigned i=0; i < NumStores; ++i) {
     // Find a chain for the new wide-store operand. Notice that some
     // of the store nodes that we found may not be selected for inclusion
for (unsigned i = 0; i < NumStores; ++i) {
       // get CSEed and the net result is that X is now a use of St.
       // Since we know that St is redundant, just iterate.
       while (!St->use_empty())
         DAG.ReplaceAllUsesWith(SDValue(St, 0), St->getChain());
       deleteAndRecombine(St);
     }
   }
 
-  return true;
+  return NumStores;
 }
 
 void DAGCombiner::getStoreMergeAndAliasCandidates(
     StoreSDNode* St, SmallVectorImpl<MemOpLink> &StoreNodes,
     SmallVectorImpl<LSBaseSDNode*> &AliasLoadNodes) {
   // This holds the base pointer, index, and the offset in bytes from the base
   // pointer.
   BaseIndexOffset BasePtr = BaseIndexOffset::match(St->getBasePtr(), DAG);
bool DAGCombiner::checkMergeStoreCandidatesForDependencies(
   // search through DAG. We can stop early if we find a storenode
   for (unsigned i = 0; i < StoreNodes.size(); ++i) {
     if (SDNode::hasPredecessorHelper(StoreNodes[i].MemNode, Visited, Worklist))
       return false;
   }
   return true;
 }
 
-bool DAGCombiner::MergeConsecutiveStores(StoreSDNode* St) {
+unsigned DAGCombiner::MergeConsecutiveStores(
+    StoreSDNode* St, SmallVectorImpl<MemOpLink> &StoreNodes) {
   if (OptLevel == CodeGenOpt::None)
-    return false;
+    return 0;
 
   EVT MemVT = St->getMemoryVT();
   int64_t ElementSizeBytes = MemVT.getSizeInBits() / 8;
   bool NoVectors = DAG.getMachineFunction().getFunction()->hasFnAttribute(
       Attribute::NoImplicitFloat);
 
   // This function cannot currently deal with non-byte-sized memory sizes.
   if (ElementSizeBytes * 8 != MemVT.getSizeInBits())
-    return false;
+    return 0;
 
   if (!MemVT.isSimple())
-    return false;
+    return 0;
 
   // Perform an early exit check. Do not bother looking at stored values that
   // are not constants, loads, or extracted vector elements.
   SDValue StoredVal = St->getValue();
   bool IsLoadSrc = isa<LoadSDNode>(StoredVal);
   bool IsConstantSrc = isa<ConstantSDNode>(StoredVal) ||
                        isa<ConstantFPSDNode>(StoredVal);
   bool IsExtractVecSrc = (StoredVal.getOpcode() == ISD::EXTRACT_VECTOR_ELT ||
                           StoredVal.getOpcode() == ISD::EXTRACT_SUBVECTOR);
 
   if (!IsConstantSrc && !IsLoadSrc && !IsExtractVecSrc)
-    return false;
+    return 0;
 
   // Don't merge vectors into wider vectors if the source data comes from loads.
   // TODO: This restriction can be lifted by using logic similar to the
   // ExtractVecSrc case.
   if (MemVT.isVector() && IsLoadSrc)
-    return false;
+    return 0;
 
   // Only look at ends of store sequences.
   SDValue Chain = SDValue(St, 0);
   if (Chain->hasOneUse() && Chain->use_begin()->getOpcode() == ISD::STORE)
-    return false;
+    return 0;
 
   // Save the LoadSDNodes that we find in the chain.
   // We need to make sure that these nodes do not interfere with
   // any of the store nodes.
   SmallVector<LSBaseSDNode*, 8> AliasLoadNodes;
 
   // Save the StoreSDNodes that we find in the chain.
-  SmallVector<MemOpLink, 8> StoreNodes;
-
   getStoreMergeAndAliasCandidates(St, StoreNodes, AliasLoadNodes);
 
   // Check if there is anything to merge.
   if (StoreNodes.size() < 2)
-    return false;
+    return 0;
 
   // only do dependence check in AA case
   bool UseAA = CombinerAA.getNumOccurrences() > 0 ? CombinerAA
                                                   : DAG.getSubtarget().useAA();
   if (UseAA && !checkMergeStoreCandidatesForDependencies(StoreNodes))
-    return false;
+    return 0;
 
   // Sort the memory operands according to their distance from the
   // base pointer.  As a secondary criteria: make sure stores coming
   // later in the code come first in the list. This is important for
   // the non-UseAA case, because we're merging stores into the FINAL
   // store along a chain which potentially contains aliasing stores.
   // Thus, if there are multiple stores to the same address, the last
   // one can be considered for merging but not the others.
for (unsigned i = 0; i < LastConsecutiveStore + 1; ++i) {
       unsigned StoreValOpcode = St->getValue().getOpcode();
       // This restriction could be loosened.
       // Bail out if any stored values are not elements extracted from a vector.
       // It should be possible to handle mixed sources, but load sources need
       // more careful handling (see the block of code below that handles
       // consecutive loads).
       if (StoreValOpcode != ISD::EXTRACT_VECTOR_ELT &&
           StoreValOpcode != ISD::EXTRACT_SUBVECTOR)
-        return false;
+        return 0;
 
       // Find a legal type for the vector store.
       unsigned Elts = i + 1;
       if (IsVec) {
         // When merging vector stores, get the total number of elements.
         Elts *= MemVT.getVectorNumElements();
       }
       EVT Ty = EVT::getVectorVT(*DAG.getContext(), MemVT.getScalarType(), Elts);
if (LdBasePtr.Base.getNode()) {
       LdBasePtr = LdPtr;
     }
 
     // We found a potential memory operand to merge.
     LoadNodes.push_back(MemOpLink(Ld, LdPtr.Offset, 0));
   }
 
   if (LoadNodes.size() < 2)
-    return false;
+    return 0;
 
   // If we have load/store pair instructions and we only have two values,
   // don't bother.
   unsigned RequiredAlignment;
   if (LoadNodes.size() == 2 && TLI.hasPairedLoad(MemVT, RequiredAlignment) &&
       St->getAlignment() >= RequiredAlignment)
-    return false;
+    return 0;
 
   LoadSDNode *FirstLoad = cast<LoadSDNode>(LoadNodes[0].MemNode);
   unsigned FirstLoadAS = FirstLoad->getAddressSpace();
   unsigned FirstLoadAlign = FirstLoad->getAlignment();
 
   // Scan the memory operations on the chain and find the first non-consecutive
   // load memory address. These variables hold the index in the store node
   // array.
unsigned DAGCombiner::MergeConsecutiveStores(
   unsigned LastLegalType = std::max(LastLegalVectorType, LastLegalIntegerType);
 
   // We add +1 here because the LastXXX variables refer to location while
   // the NumElem refers to array/index size.
   unsigned NumElem = std::min(LastConsecutiveStore, LastConsecutiveLoad) + 1;
   NumElem = std::min(LastLegalType, NumElem);
 
   if (NumElem < 2)
-    return false;
+    return 0;
 
   // Collect the chains from all merged stores.
   SmallVector<SDValue, 8> MergeStoreChains;
   MergeStoreChains.push_back(StoreNodes[0].MemNode->getChain());
 
   // The latest Node in the DAG.
   unsigned LatestNodeUsed = 0;
   for (unsigned i=1; i<NumElem; ++i) {
for (unsigned i = 0; i < NumElem; ++i) {
       if (StoreNodes[i].MemNode == LatestOp)
         continue;
       StoreSDNode *St = cast<StoreSDNode>(StoreNodes[i].MemNode);
       DAG.ReplaceAllUsesOfValueWith(SDValue(St, 0), St->getChain());
       deleteAndRecombine(St);
     }
   }
 
-  return true;
+  return NumElem;
 }
 
 SDValue DAGCombiner::replaceStoreChain(StoreSDNode *ST, SDValue BetterChain) {
   SDLoc SL(ST);
   SDValue ReplStore;
 
   // Replace the chain to avoid dependency.
   if (ST->isTruncatingStore()) {
return DAG.getTruncStore(Chain, SDLoc(N), Value.getOperand(0),
                              Ptr, ST->getMemoryVT(), ST->getMemOperand());
   }
 
   // Only perform this optimization before the types are legal, because we
   // don't want to perform this optimization on every DAGCombine invocation.
   if (!LegalTypes) {
     bool EverChanged = false;
 
-    do {
+    for (;;) {
       // There can be multiple store sequences on the same chain.
       // Keep trying to merge store sequences until we are unable to do so
       // or until we merge the last store on the chain.
-      bool Changed = MergeConsecutiveStores(ST);
-      EverChanged |= Changed;
-      if (!Changed) break;
-    } while (ST->getOpcode() != ISD::DELETED_NODE);
+      SmallVector<MemOpLink, 8> StoreNodes;
+      unsigned NumChanged = MergeConsecutiveStores(ST, StoreNodes);
+      if (!NumChanged)
+        break;
+
+      EverChanged = true;
+
+      bool ChangedST = false;
+      for (unsigned i = 0; i < NumChanged; ++i) {
+        if (StoreNodes[i].MemNode == ST) {
+          ChangedST = true;
+          break;
+        }
+      }
+      if (ChangedST)
+        break;
+    }
 
     if (EverChanged)
       return SDValue(N, 0);
   }
 
   // Turn 'store float 1.0, Ptr' -> 'store int 0x12345678, Ptr'
   //
   // Make sure to do this only after attempting to merge stores in order to