[AST] Unique and ASTContext-allocate SubstitutionMaps.

Prepare for SubstitutionMaps to be stored in other AST nodes by making
them ASTContext-allocated and uniqued (via a FoldingSet). They are now
cheap to copy and have trivial destructors.

Author: DougGregor

include/swift/AST/SubstitutionMap.h

@@ -22,9 +22,9 @@
 #include "swift/AST/SubstitutionList.h"
 #include "swift/AST/Type.h"
 #include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/FoldingSet.h"
 #include "llvm/ADT/Optional.h"
 #include "llvm/Support/TrailingObjects.h"
-#include <memory>
 
 namespace llvm {
   class FoldingSetNodeID;
@@ -56,10 +56,12 @@ enum class CombineSubstitutionMaps {
 /// on a GenericSignature or (equivalently) by calling one of the static
 /// \c SubstitutionMap::get() methods.
 class SubstitutionMap {
+public:
   /// Stored data for a substitution map, which uses tail allocation for the
   /// replacement types and conformances.
   class Storage final
-    : llvm::TrailingObjects<Storage, Type, ProtocolConformanceRef>
+    : public llvm::FoldingSetNode,
+      llvm::TrailingObjects<Storage, Type, ProtocolConformanceRef>
   {
     friend TrailingObjects;
 
@@ -126,13 +128,26 @@ class SubstitutionMap {
                                 getTrailingObjects<ProtocolConformanceRef>(),
                                 numConformanceRequirements);
     }
+
+    /// Profile the substitution map storage, for use with LLVM's FoldingSet.
+    void Profile(llvm::FoldingSetNodeID &id) const {
+      Profile(id, getGenericSignature(), getReplacementTypes(),
+              getConformances());
+    }
+
+    /// Profile the substitution map storage, for use with LLVM's FoldingSet.
+    static void Profile(llvm::FoldingSetNodeID &id,
+                        GenericSignature *genericSig,
+                        ArrayRef<Type> replacementTypes,
+                        ArrayRef<ProtocolConformanceRef> conformances);
   };
 
+private:
   /// The storage needed to describe the set of substitutions.
   ///
   /// When null, this substitution map is empty, having neither a generic
   /// signature nor any replacement types/conformances.
-  std::shared_ptr<Storage> storage;
+  Storage *storage = nullptr;
 
   /// Retrieve the array of replacement types, which line up with the
   /// generic parameters.

lib/AST/ASTContext.cpp

@@ -309,6 +309,9 @@ FOR_KNOWN_FOUNDATION_TYPES(CACHE_FOUNDATION_DECL)
     /// The set of inherited protocol conformances.
     llvm::FoldingSet<InheritedProtocolConformance> InheritedConformances;
 
+    /// The set of substitution maps (uniqued by their storage).
+    llvm::FoldingSet<SubstitutionMap::Storage> SubstitutionMaps;
+
     ~Arena() {
       for (auto &conformance : SpecializedConformances)
         conformance.~SpecializedProtocolConformance();
@@ -2907,11 +2910,9 @@ NameAliasType::NameAliasType(TypeAliasDecl *typealias, Type parent,
   }
 }
 
-NameAliasType *NameAliasType::get(
-                                        TypeAliasDecl *typealias,
-                                        Type parent,
-                                        const SubstitutionMap &substitutions,
-                                        Type underlying) {
+NameAliasType *NameAliasType::get(TypeAliasDecl *typealias, Type parent,
+                                  const SubstitutionMap &substitutions,
+                                  Type underlying) {
   // Compute the recursive properties.
   //
   auto properties = underlying->getRecursiveProperties();
@@ -4293,6 +4294,64 @@ CapturingTypeCheckerDebugConsumer::~CapturingTypeCheckerDebugConsumer() {
   delete Log;
 }
 
+void SubstitutionMap::Storage::Profile(
+                               llvm::FoldingSetNodeID &id,
+                               GenericSignature *genericSig,
+                               ArrayRef<Type> replacementTypes,
+                               ArrayRef<ProtocolConformanceRef> conformances) {
+  id.AddPointer(genericSig);
+  id.AddInteger(replacementTypes.size());
+  for (auto type : replacementTypes)
+    id.AddPointer(type.getPointer());
+  id.AddInteger(conformances.size());
+  for (auto conformance : conformances)
+    id.AddPointer(conformance.getOpaqueValue());
+}
+
+SubstitutionMap::Storage *SubstitutionMap::Storage::get(
+                            GenericSignature *genericSig,
+                            ArrayRef<Type> replacementTypes,
+                            ArrayRef<ProtocolConformanceRef> conformances) {
+  // If there is no generic signature, we need no storage.
+  if (!genericSig) {
+    assert(replacementTypes.empty());
+    assert(conformances.empty());
+    return nullptr;
+  }
+
+  // Figure out which arena this should go in.
+  RecursiveTypeProperties properties;
+  for (auto type : replacementTypes) {
+    if (type)
+      properties |= type->getRecursiveProperties();
+  }
+
+  // Profile the substitution map.
+  llvm::FoldingSetNodeID id;
+  SubstitutionMap::Storage::Profile(id, genericSig, replacementTypes,
+                                    conformances);
+
+  auto arena = getArena(properties);
+
+  // Did we already record this substitution map?
+  auto &ctx = genericSig->getASTContext();
+  void *insertPos;
+  auto &substitutionMaps = ctx.Impl.getArena(arena).SubstitutionMaps;
+  if (auto result = substitutionMaps.FindNodeOrInsertPos(id, insertPos))
+    return result;
+
+  // Allocate the appropriate amount of storage for the signature and its
+  // replacement types and conformances.
+  auto size = Storage::totalSizeToAlloc<Type, ProtocolConformanceRef>(
+                                                      replacementTypes.size(),
+                                                      conformances.size());
+  auto mem = ctx.Allocate(size, alignof(Storage), arena);
+
+  auto result = new (mem) Storage(genericSig, replacementTypes, conformances);
+  substitutionMaps.InsertNode(result, insertPos);
+  return result;
+}
+
 void GenericSignature::Profile(llvm::FoldingSetNodeID &ID,
                               TypeArrayView<GenericTypeParamType> genericParams,
                               ArrayRef<Requirement> requirements) {

lib/AST/SubstitutionMap.cpp

@@ -50,26 +50,6 @@ SubstitutionMap::Storage::Storage(
             getConformances().data());
 }
 
-SubstitutionMap::Storage *SubstitutionMap::Storage::get(
-                            GenericSignature *genericSig,
-                            ArrayRef<Type> replacementTypes,
-                            ArrayRef<ProtocolConformanceRef> conformances) {
-  // If there is no generic signature, we need no storage.
-  if (!genericSig) {
-    assert(replacementTypes.empty());
-    assert(conformances.empty());
-    return nullptr;
-  }
-
-  // Allocate the appropriate amount of storage for the signature and its
-  // replacement types and conformances.
-  auto size = Storage::totalSizeToAlloc<Type, ProtocolConformanceRef>(
-                                                      replacementTypes.size(),
-                                                      conformances.size());
-  auto mem = malloc(size);
-  return new (mem) Storage(genericSig, replacementTypes, conformances);
-}
-
 bool SubstitutionMap::hasArchetypes() const {
   for (Type replacementTy : getReplacementTypes()) {
     if (replacementTy && replacementTy->hasArchetype())
@@ -587,20 +567,6 @@ void SubstitutionMap::dump() const {
 }
 
 void SubstitutionMap::profile(llvm::FoldingSetNodeID &id) const {
-  // Generic signature.
-  auto genericSig = getGenericSignature();
-  id.AddPointer(genericSig);
-
-  if (empty() || !genericSig) return;
-
-  // Replacement types.
-  for (Type gp : genericSig->getGenericParams()) {
-    id.AddPointer(gp.subst(*this).getPointer());
-  }
-
-  // Conformance requirements.
-  for (const auto conformance : getConformances()) {
-    id.AddPointer(conformance.getOpaqueValue());
-  }
+  id.AddPointer(storage);
 }