[PATCH] D25601: DAGCombiner: fix use-after-free when merging consecutive stores

[Justin Bogner via llvm-commits]

Nicolai Hähnle <nhaehnle at gmail.com> writes:
> nhaehnle updated this revision to Diff 76719.
> nhaehnle added a comment.
>
> Keep returning a boolean.
>
> Since not all nodes are necessarily merged, the success-path now needs to
> truncate the StoreNodes vector to just those nodes that were actually
> removed.

Looks pretty close. A couple of comments below.

>
> https://reviews.llvm.org/D25601
>
> Files:
>   lib/CodeGen/SelectionDAG/DAGCombiner.cpp
>
>
> Index: lib/CodeGen/SelectionDAG/DAGCombiner.cpp
> ===================================================================
> --- lib/CodeGen/SelectionDAG/DAGCombiner.cpp
> +++ lib/CodeGen/SelectionDAG/DAGCombiner.cpp
> @@ -450,10 +450,11 @@
>      /// 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).
> +    bool MergeStoresOfConstantsOrVecElts(
> +        SmallVectorImpl<MemOpLink> &StoreNodes, EVT MemVT, unsigned NumStores,
> +        bool IsConstantSrc, bool UseVector);
>  
>      /// This is a helper function for MergeConsecutiveStores.
>      /// Stores that may be merged are placed in StoreNodes.
> @@ -470,8 +471,10 @@
>  
>      /// 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).
> +    bool 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))
> @@ -11513,6 +11516,7 @@
>      }
>    }
>  
> +  StoreNodes.erase(StoreNodes.begin() + NumStores, StoreNodes.end());
>    return true;
>  }
>  
> @@ -11655,7 +11659,8 @@
>    return true;
>  }
>  
> -bool DAGCombiner::MergeConsecutiveStores(StoreSDNode* St) {
> +bool DAGCombiner::MergeConsecutiveStores(
> +    StoreSDNode* St, SmallVectorImpl<MemOpLink> &StoreNodes) {
>    if (OptLevel == CodeGenOpt::None)
>      return false;
>  
> @@ -11699,9 +11704,6 @@
>    // 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.
> @@ -12060,6 +12062,7 @@
>      }
>    }
>  
> +  StoreNodes.erase(StoreNodes.begin() + NumElem, StoreNodes.end());
>    return true;
>  }
>  
> @@ -12305,14 +12308,25 @@
>    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);
> +      SmallVector<MemOpLink, 8> StoreNodes;
> +      bool Changed = MergeConsecutiveStores(ST, StoreNodes);
>        EverChanged |= Changed;
>        if (!Changed) break;
> -    } while (ST->getOpcode() != ISD::DELETED_NODE);
> +
> +      bool ChangedST = false;
> +      for (const auto &Link : StoreNodes) {
> +        if (Link.MemNode == ST) {
> +          ChangedST = true;
> +          break;
> +        }
> +      }
> +      if (ChangedST)
> +        break;
> +    }
>  
>      if (EverChanged)
>        return SDValue(N, 0);
>
>
> Index: lib/CodeGen/SelectionDAG/DAGCombiner.cpp
> ===================================================================
> --- lib/CodeGen/SelectionDAG/DAGCombiner.cpp
> +++ lib/CodeGen/SelectionDAG/DAGCombiner.cpp
> @@ -450,10 +450,11 @@
>      /// 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).
> +    bool MergeStoresOfConstantsOrVecElts(
> +        SmallVectorImpl<MemOpLink> &StoreNodes, EVT MemVT, unsigned NumStores,
> +        bool IsConstantSrc, bool UseVector);
>  
>      /// This is a helper function for MergeConsecutiveStores.
>      /// Stores that may be merged are placed in StoreNodes.
> @@ -470,8 +471,10 @@
>  
>      /// 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).
> +    bool 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))
> @@ -11513,6 +11516,7 @@
>      }
>    }
>  
> +  StoreNodes.erase(StoreNodes.begin() + NumStores, StoreNodes.end());
>    return true;
>  }
>  
> @@ -11655,7 +11659,8 @@
>    return true;
>  }
>  
> -bool DAGCombiner::MergeConsecutiveStores(StoreSDNode* St) {
> +bool DAGCombiner::MergeConsecutiveStores(
> +    StoreSDNode* St, SmallVectorImpl<MemOpLink> &StoreNodes) {
>    if (OptLevel == CodeGenOpt::None)
>      return false;
>  
> @@ -11699,9 +11704,6 @@
>    // 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.
> @@ -12060,6 +12062,7 @@
>      }
>    }
>  
> +  StoreNodes.erase(StoreNodes.begin() + NumElem, StoreNodes.end());
>    return true;
>  }
>  
> @@ -12305,14 +12308,25 @@
>    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);
> +      SmallVector<MemOpLink, 8> StoreNodes;
> +      bool Changed = MergeConsecutiveStores(ST, StoreNodes);

What do you think of returning the SmallVector instead of the bool, then
just doing `EverChanged |= !StoreNodes.empty()`? It seems a bit simpler
to me, and move semantics should mean it's just as efficient.

>        EverChanged |= Changed;
>        if (!Changed) break;
> -    } while (ST->getOpcode() != ISD::DELETED_NODE);
> +
> +      bool ChangedST = false;
> +      for (const auto &Link : StoreNodes) {
> +        if (Link.MemNode == ST) {
> +          ChangedST = true;
> +          break;
> +        }
> +      }
> +      if (ChangedST)
> +        break;

I think this can be simplified with any_of. Something like this:

      if (any_of(StoreNodes,
                 [&ST](const MemOpLink &Link) { return Link.MemNode == ST; }))
        return;

> +    }
>  
>      if (EverChanged)
>        return SDValue(N, 0);