Keep alias sets with undef pointer values in UndefPointerSets

This allows to scan only these sets when we want to check if a memory
locations was already registered, but importantly, it also keeps these
alias sets referenced.

Author: brunodf-snps

llvm/include/llvm/Analysis/AliasSetTracker.h

@@ -184,6 +184,10 @@ class AliasSetTracker {
   // Map from pointers to their node
   PointerMapType PointerMap;
 
+  // Collection of alias sets that contain undef pointer values
+  DenseSet<AliasSet *> UndefPointerSets;
+  std::vector<AliasSet *> UndefPointerSetsVector;
+
 public:
   /// Create an empty collection of AliasSets, and use the specified alias
   /// analysis object to disambiguate load and store addresses.

llvm/lib/Analysis/AliasSetTracker.cpp

@@ -228,6 +228,8 @@ AliasSet::PointerVector AliasSet::getPointers() const {
 
 void AliasSetTracker::clear() {
   PointerMap.clear();
+  UndefPointerSets.clear();
+  UndefPointerSetsVector.clear();
   AliasSets.clear();
 }
 
@@ -287,7 +289,8 @@ AliasSet &AliasSetTracker::getAliasSetFor(const MemoryLocation &MemLoc) {
   // The PointerMap structure requires that all memory locations for the same
   // value will be in the same alias set, which does not hold for undef values.
   // So we keep UndefValue pointers associated with the nullptr in the pointer
-  // map, and resort to a linear scan of all alias sets in that case.
+  // map, and instead keep a collection with the alias sets that contain memory
+  // locations with undef pointer values.
   auto [It, Inserted] = PointerMap.insert({MemLoc.Ptr, nullptr});
   AliasSet *&MapEntry = It->second;
   if (!Inserted) {
@@ -302,12 +305,27 @@ AliasSet &AliasSetTracker::getAliasSetFor(const MemoryLocation &MemLoc) {
         return *AS;
       }
     } else {
-      for (auto &AS : *this) {
-        if (AS.isForwardingAliasSet())
-          continue;
-        if (llvm::find(AS.MemoryLocs, MemLoc) != AS.MemoryLocs.end())
-          return AS;
-      }
+      // Take opportunity to remove refs to forwarding alias sets from UndefPointerSets.
+      llvm::erase_if(UndefPointerSetsVector, [this](AliasSet *&AS) {
+        if (AS->Forward) {
+          AliasSet *NewAS = AS->getForwardedTarget(*this);
+          UndefPointerSets.erase(AS);
+          AS->dropRef(*this);
+          // Replace earlier entry in UndefPointerSetsVector with forwarded
+          // target, but only if it is new.
+          if (UndefPointerSets.insert(NewAS).second) {
+            NewAS->addRef();
+            AS = NewAS;
+            return false;
+          } else
+            return true;
+        }
+        return false;
+      });
+      // Scan UndefPointerSets for MemLoc.
+      for (AliasSet *AS : UndefPointerSetsVector)
+        if (llvm::find(AS->MemoryLocs, MemLoc) != AS->MemoryLocs.end())
+          return *AS;
     }
   }
 
@@ -349,6 +367,11 @@ AliasSet &AliasSetTracker::getAliasSetFor(const MemoryLocation &MemLoc) {
       AS->addRef();
       MapEntry = AS;
     }
+  } else {
+    if (UndefPointerSets.insert(AS).second) {
+      UndefPointerSetsVector.push_back(AS);
+      AS->addRef();
+    }
   }
   return *AS;
 }