From 659e145e5bc2758bf9ab341b168b0bea354d0cd5 Mon Sep 17 00:00:00 2001
From: Robert Widmann <devteam.codafi@gmail.com>
Date: Tue, 11 Jan 2022 15:01:42 -0800
Subject: [PATCH] Propagate Pattern Arena When Allocating PackExpansionTypes

This may be the cause of the instability we've been seeing on the bots
here. If a type variable is used as the pattern type across allocation
arenas, we may incorrectly reuse a cache entry and it could be
substituted a second time for a bogus type.

rdar://87413287
---
 lib/AST/ASTContext.cpp | 13 +++++++++----
 1 file changed, 9 insertions(+), 4 deletions(-)

diff --git a/lib/AST/ASTContext.cpp b/lib/AST/ASTContext.cpp
index 42d1e81a4ad9e..84982ee54c39a 100644
--- a/lib/AST/ASTContext.cpp
+++ b/lib/AST/ASTContext.cpp
@@ -389,6 +389,7 @@ struct ASTContext::Implementation {
     llvm::FoldingSet<TypeAliasType> TypeAliasTypes;
     llvm::FoldingSet<TupleType> TupleTypes;
     llvm::FoldingSet<PackType> PackTypes;
+    llvm::FoldingSet<PackExpansionType> PackExpansionTypes;
     llvm::DenseMap<llvm::PointerIntPair<TypeBase*, 3, unsigned>,
                    MetatypeType*> MetatypeTypes;
     llvm::DenseMap<llvm::PointerIntPair<TypeBase*, 3, unsigned>,
@@ -471,7 +472,6 @@ struct ASTContext::Implementation {
   llvm::FoldingSet<SILBoxType> SILBoxTypes;
   llvm::DenseMap<BuiltinIntegerWidth, BuiltinIntegerType*> IntegerTypes;
   llvm::FoldingSet<BuiltinVectorType> BuiltinVectorTypes;
-  llvm::FoldingSet<PackExpansionType> PackExpansionTypes;
   llvm::FoldingSet<DeclName::CompoundDeclName> CompoundNames;
   llvm::DenseMap<UUID, OpenedArchetypeType *> OpenedExistentialArchetypes;
   llvm::FoldingSet<IndexSubset> IndexSubsets;
@@ -2960,20 +2960,25 @@ Type TupleTypeElt::getType() const {
 PackExpansionType *PackExpansionType::get(Type patternTy) {
   assert(patternTy && "Missing pattern type in expansion");
 
+  auto properties = patternTy->getRecursiveProperties();
+  auto arena = getArena(properties);
+
   auto &context = patternTy->getASTContext();
   llvm::FoldingSetNodeID id;
   PackExpansionType::Profile(id, patternTy);
 
   void *insertPos;
   if (PackExpansionType *expType =
-          context.getImpl().PackExpansionTypes.FindNodeOrInsertPos(id,
-                                                                   insertPos))
+          context.getImpl()
+              .getArena(arena)
+              .PackExpansionTypes.FindNodeOrInsertPos(id, insertPos))
     return expType;
 
   const ASTContext *canCtx = patternTy->isCanonical() ? &context : nullptr;
   PackExpansionType *expansionTy = new (context, AllocationArena::Permanent)
       PackExpansionType(patternTy, canCtx);
-  context.getImpl().PackExpansionTypes.InsertNode(expansionTy, insertPos);
+  context.getImpl().getArena(arena).PackExpansionTypes.InsertNode(expansionTy,
+                                                                  insertPos);
   return expansionTy;
 }