[BasicAA] Fix BatchAA results for phi-phi assumptions

Add a flag that disables caching when computing aliasing results
potentially based on a phi-phi NoAlias assumption. We'll still
insert cache entries temporarily to catch infinite recursion,
but will drop them afterwards, so they won't persist in BatchAA.

Differential Revision: https://reviews.llvm.org/D91936

include/llvm/Analysis/BasicAliasAnalysis.h

@@ -162,6 +162,10 @@ private:
   /// Tracks instructions visited by pointsToConstantMemory.
   SmallPtrSet<const Value *, 16> Visited;
 
+  /// Whether to disable persistent caching in AAQI. This is used to prevent
+  /// caching of results based on temporary assumptions.
+  bool DisableCache = false;
+
   static const Value *
   GetLinearExpression(const Value *V, APInt &Scale, APInt &Offset,
                       unsigned &ZExtBits, unsigned &SExtBits,
@@ -216,6 +220,12 @@ private:
                          LocationSize V2Size, const AAMDNodes &V2AATag,
                          AAQueryInfo &AAQI, const Value *O1 = nullptr,
                          const Value *O2 = nullptr);
+
+  AliasResult aliasCheckRecursive(const Value *V1, LocationSize V1Size,
+                                  const AAMDNodes &V1AATag, const Value *V2,
+                                  LocationSize V2Size, const AAMDNodes &V2AATag,
+                                  AAQueryInfo &AAQI, const Value *O1,
+                                  const Value *O2);
 };
 
 /// Analysis pass providing a never-invalidated alias analysis result.

lib/Analysis/BasicAliasAnalysis.cpp

@@ -1489,8 +1489,10 @@ AliasResult BasicAAResult::aliasPHI(const PHINode *PN, LocationSize PNSize,
       // operand from outside the PHIs' cycle that is MayAlias/MustAlias or
       // there must be an operation on the PHIs within the PHIs' value cycle
       // that causes a MayAlias.
-      // Pretend the phis do not alias.
-      AliasResult Alias = NoAlias;
+      // Disable persistent caching, so intermediate results based on a
+      // possibly incorrect assumption do not get cached.
+      bool OrigDisableCache = DisableCache;
+      DisableCache = true;
       AliasResult OrigAliasResult;
       {
         // Limited lifetime iterator invalidated by the aliasCheck call below.
@@ -1501,6 +1503,7 @@ AliasResult BasicAAResult::aliasPHI(const PHINode *PN, LocationSize PNSize,
         CacheIt->second = NoAlias;
       }
 
+      AliasResult Alias = NoAlias;
       for (unsigned i = 0, e = PN->getNumIncomingValues(); i != e; ++i) {
         AliasResult ThisAlias =
             aliasCheck(PN->getIncomingValue(i), PNSize, PNAAInfo,
@@ -1514,6 +1517,7 @@ AliasResult BasicAAResult::aliasPHI(const PHINode *PN, LocationSize PNSize,
       // Reset if speculation failed.
       if (Alias != NoAlias)
         AAQI.updateResult(Locs, OrigAliasResult);
+      DisableCache = OrigDisableCache;
       return Alias;
     }
 
@@ -1753,59 +1757,73 @@ AliasResult BasicAAResult::aliasCheck(const Value *V1, LocationSize V1Size,
   if (!Pair.second)
     return Pair.first->second;
 
+  AliasResult Result = aliasCheckRecursive(V1, V1Size, V1AAInfo, V2, V2Size,
+                                           V2AAInfo, AAQI, O1, O2);
+
+  // If caching is disabled, remove the entry once the recursive checks are
+  // done. We only needed it to prevent infinite recursion.
+  if (DisableCache)
+    AAQI.AliasCache.erase(AAQI.AliasCache.find(Locs));
+  else if (Result != MayAlias)
+    AAQI.updateResult(Locs, Result);
+  return Result;
+}
+
+AliasResult BasicAAResult::aliasCheckRecursive(
+    const Value *V1, LocationSize V1Size, const AAMDNodes &V1AAInfo,
+    const Value *V2, LocationSize V2Size, const AAMDNodes &V2AAInfo,
+    AAQueryInfo &AAQI, const Value *O1, const Value *O2) {
   if (const GEPOperator *GV1 = dyn_cast<GEPOperator>(V1)) {
     AliasResult Result =
         aliasGEP(GV1, V1Size, V1AAInfo, V2, V2Size, V2AAInfo, O1, O2, AAQI);
     if (Result != MayAlias)
-      return AAQI.updateResult(Locs, Result);
+      return Result;
   } else if (const GEPOperator *GV2 = dyn_cast<GEPOperator>(V2)) {
     AliasResult Result =
         aliasGEP(GV2, V2Size, V2AAInfo, V1, V1Size, V1AAInfo, O2, O1, AAQI);
     if (Result != MayAlias)
-      return AAQI.updateResult(Locs, Result);
+      return Result;
   }
 
   if (const PHINode *PN = dyn_cast<PHINode>(V1)) {
     AliasResult Result =
         aliasPHI(PN, V1Size, V1AAInfo, V2, V2Size, V2AAInfo, O2, AAQI);
     if (Result != MayAlias)
-      return AAQI.updateResult(Locs, Result);
+      return Result;
   } else if (const PHINode *PN = dyn_cast<PHINode>(V2)) {
     AliasResult Result =
         aliasPHI(PN, V2Size, V2AAInfo, V1, V1Size, V1AAInfo, O1, AAQI);
     if (Result != MayAlias)
-      return AAQI.updateResult(Locs, Result);
+      return Result;
   }
 
   if (const SelectInst *S1 = dyn_cast<SelectInst>(V1)) {
     AliasResult Result =
         aliasSelect(S1, V1Size, V1AAInfo, V2, V2Size, V2AAInfo, O2, AAQI);
     if (Result != MayAlias)
-      return AAQI.updateResult(Locs, Result);
+      return Result;
   } else if (const SelectInst *S2 = dyn_cast<SelectInst>(V2)) {
     AliasResult Result =
         aliasSelect(S2, V2Size, V2AAInfo, V1, V1Size, V1AAInfo, O1, AAQI);
     if (Result != MayAlias)
-      return AAQI.updateResult(Locs, Result);
+      return Result;
   }
 
   // If both pointers are pointing into the same object and one of them
   // accesses the entire object, then the accesses must overlap in some way.
-  if (O1 == O2)
+  if (O1 == O2) {
+    bool NullIsValidLocation = NullPointerIsDefined(&F);
     if (V1Size.isPrecise() && V2Size.isPrecise() &&
         (isObjectSize(O1, V1Size.getValue(), DL, TLI, NullIsValidLocation) ||
          isObjectSize(O2, V2Size.getValue(), DL, TLI, NullIsValidLocation)))
-      return AAQI.updateResult(Locs, PartialAlias);
+      return PartialAlias;
+  }
 
   // Recurse back into the best AA results we have, potentially with refined
   // memory locations. We have already ensured that BasicAA has a MayAlias
   // cache result for these, so any recursion back into BasicAA won't loop.
-  AliasResult Result = getBestAAResults().alias(Locs.first, Locs.second, AAQI);
-  if (Result != MayAlias)
-    return AAQI.updateResult(Locs, Result);
-
-  // MayAlias is already in the cache.
-  return MayAlias;
+  return getBestAAResults().alias(MemoryLocation(V1, V1Size, V1AAInfo),
+                                  MemoryLocation(V2, V2Size, V2AAInfo), AAQI);
 }
 
 /// Check whether two Values can be considered equivalent.

unittests/Analysis/AliasAnalysisTest.cpp

@@ -295,8 +295,7 @@ TEST_F(AliasAnalysisTest, BatchAAPhiAssumption) {
 
   BatchAAResults BatchAA(AA);
   EXPECT_EQ(MayAlias, BatchAA.alias(ALoc, BLoc));
-  // TODO: This is incorrect.
-  EXPECT_EQ(NoAlias, BatchAA.alias(ANextLoc, BNextLoc));
+  EXPECT_EQ(MayAlias, BatchAA.alias(ANextLoc, BNextLoc));
 }
 
 class AAPassInfraTest : public testing::Test {