Simplify and optimize SubstitutionMap

include/swift/AST/SubstitutionMap.h

@@ -73,22 +73,6 @@ class SubstitutionMap {
   /// signature nor any replacement types/conformances.
   Storage *storage = nullptr;
 
-public:
-  /// Retrieve the array of replacement types, which line up with the
-  /// generic parameters.
-  ///
-  /// Note that the types may be null, for cases where the generic parameter
-  /// is concrete but hasn't been queried yet.
-  ///
-  /// Prefer \c getReplacementTypes, this is public for printing purposes.
-  ArrayRef<Type> getReplacementTypesBuffer() const;
-
-private:
-  MutableArrayRef<Type> getReplacementTypesBuffer();
-
-  /// Retrieve a mutable reference to the buffer of conformances.
-  MutableArrayRef<ProtocolConformanceRef> getConformancesBuffer();
-
   /// Form a substitution map for the given generic signature with the
   /// specified replacement types and conformances.
   SubstitutionMap(GenericSignature genericSig,

include/swift/AST/Types.h

@@ -1289,26 +1289,22 @@ class alignas(1 << TypeAlignInBits) TypeBase
   /// Otherwise, it returns the type itself.
   Type getReferenceStorageReferent();
 
-  /// Determine the set of substitutions that should be applied to a
-  /// type spelled within the given DeclContext to treat it as a
-  /// member of this type.
+  /// Assumes this is a nominal type. Returns a substitution map that sends each
+  /// generic parameter of the declaration's generic signature to the corresponding
+  /// generic argument of this nominal type.
   ///
-  /// For example, given:
-  /// \code
-  /// struct X<T, U> { }
-  /// extension X {
-  ///   typealias SomeArray = [T]
-  /// }
-  /// \endcode
+  /// Eg: Array<Int> ---> { Element := Int }
+  SubstitutionMap getContextSubstitutionMap();
+
+  /// More general form of the above that handles additional cases:
   ///
-  /// Asking for the member substitutions of \c X<Int,String> within
-  /// the context of the extension above will produce substitutions T
-  /// -> Int and U -> String suitable for mapping the type of
-  /// \c SomeArray.
+  /// 1) dc is the nominal type itself or an unconstrained extension
+  /// 2) dc is a superclass
+  /// 3) dc is a protocol
   ///
-  /// \param genericEnv If non-null and the type is nested inside of a
-  /// generic function, generic parameters of the outer context are
-  /// mapped to context archetypes of this generic environment.
+  /// In Case 2) and 3), the substitution map has the generic signature of the dc,
+  /// and not the nominal. In Case 1), this is the same as the no-argument overload
+  /// of getContextSubstitutionMap().
   SubstitutionMap getContextSubstitutionMap(const DeclContext *dc,
                                             GenericEnvironment *genericEnv=nullptr);
 
@@ -1530,6 +1526,9 @@ END_CAN_TYPE_WRAPPER(AnyGenericType, Type)
 /// fields exist at the same offset in memory to improve code generation of the
 /// compiler itself.
 class NominalOrBoundGenericNominalType : public AnyGenericType {
+  friend class TypeBase;
+  SubstitutionMap ContextSubMap;
+
 public:
   template <typename... Args>
   NominalOrBoundGenericNominalType(Args &&...args)

include/swift/SILOptimizer/OptimizerBridgingImpl.h

@@ -415,8 +415,7 @@ void BridgedPassContext::loadFunction(BridgedFunction function, bool loadCallees
 
 BridgedSubstitutionMap BridgedPassContext::getContextSubstitutionMap(BridgedType type) const {
   swift::SILType ty = type.unbridged();
-  auto *ntd = ty.getASTType()->getAnyNominal();
-  return ty.getASTType()->getContextSubstitutionMap(ntd);
+  return ty.getASTType()->getContextSubstitutionMap();
 }
 
 BridgedType BridgedPassContext::getBuiltinIntegerType(SwiftInt bitWidth) const {

lib/AST/ASTContext.cpp

@@ -584,6 +584,8 @@ struct ASTContext::Implementation {
     }
 
     size_t getTotalMemory() const;
+
+    void dump(llvm::raw_ostream &out) const;
   };
 
   llvm::DenseMap<ModuleDecl*, ModuleType*> ModuleTypes;
@@ -672,6 +674,8 @@ struct ASTContext::Implementation {
   BuiltinTupleType *TheTupleType = nullptr;
 
   std::array<ProtocolDecl *, NumInvertibleProtocols> InvertibleProtocolDecls = {};
+
+  void dump(llvm::raw_ostream &out) const;
 };
 
 ASTContext::Implementation::Implementation()
@@ -803,9 +807,59 @@ ASTContext::ASTContext(
   }
 }
 
+void ASTContext::Implementation::dump(llvm::raw_ostream &os) const {
+  os << "-------------------------------------------------\n";
+  os << "Arena\t0\t" << Allocator.getBytesAllocated() << "\n";
+  Permanent.dump(os);
+
+#define SIZE(Name) os << #Name << "\t" << Name.size() << "\t0\n"
+#define SIZE_AND_BYTES(Name) os << #Name << "\t"                          \
+                                << Name.size() << "\t"                    \
+                                << llvm::capacity_in_bytes(Name) << "\n"
+
+  SIZE(LoadedModules);
+  SIZE(IdentifierTable);
+  SIZE(Cleanups);
+  SIZE_AND_BYTES(ModuleLoaders);
+  SIZE_AND_BYTES(ExternalSourceLocs);
+  SIZE_AND_BYTES(ForeignErrorConventions);
+  SIZE_AND_BYTES(ForeignAsyncConventions);
+  SIZE_AND_BYTES(AssociativityCache);
+  SIZE_AND_BYTES(DelayedConformanceDiags);
+  SIZE_AND_BYTES(LazyContexts);
+  SIZE_AND_BYTES(ExistentialSignatures);
+  SIZE_AND_BYTES(ElementSignatures);
+  SIZE_AND_BYTES(Overrides);
+  SIZE_AND_BYTES(DefaultWitnesses);
+  SIZE_AND_BYTES(DefaultTypeWitnesses);
+  SIZE_AND_BYTES(DefaultAssociatedConformanceWitnesses);
+  SIZE_AND_BYTES(DefaultTypeRequestCaches);
+  SIZE_AND_BYTES(PropertyWrapperBackingVarTypes);
+  SIZE_AND_BYTES(OriginalWrappedProperties);
+  SIZE_AND_BYTES(BuiltinInitWitness);
+  SIZE_AND_BYTES(OriginalBodySourceRanges);
+  SIZE_AND_BYTES(NextMacroDiscriminator);
+  SIZE_AND_BYTES(NextDiscriminator);
+  SIZE_AND_BYTES(ModuleTypes);
+  SIZE_AND_BYTES(GenericParamTypes);
+  SIZE_AND_BYTES(SILBlockStorageTypes);
+  SIZE_AND_BYTES(SILMoveOnlyWrappedTypes);
+  SIZE_AND_BYTES(IntegerTypes);
+  SIZE_AND_BYTES(OpenedExistentialEnvironments);
+  SIZE_AND_BYTES(OpenedElementEnvironments);
+  SIZE_AND_BYTES(ForeignRepresentableCache);
+  SIZE(SearchPathsSet);
+
+#undef SIZE
+#undef SIZE_AND_BYTES
+}
+
 ASTContext::~ASTContext() {
-  if (LangOpts.AnalyzeRequestEvaluator)
+  if (LangOpts.AnalyzeRequestEvaluator) {
     evaluator.dump(llvm::dbgs());
+    getImpl().dump(llvm::dbgs());
+  }
+
   getImpl().~Implementation();
 }
 
@@ -3064,6 +3118,56 @@ size_t ASTContext::Implementation::Arena::getTotalMemory() const {
     // BuiltinConformances ?
 }
 
+void ASTContext::Implementation::Arena::dump(llvm::raw_ostream &os) const {
+#define SIZE(Name) os << #Name << "\t" << Name.size() << "\t0\n"
+#define SIZE_AND_BYTES(Name) os << #Name << "\t"                          \
+                                << Name.size() << "\t"                    \
+                                << llvm::capacity_in_bytes(Name) << "\n"
+
+    SIZE_AND_BYTES(ErrorTypesWithOriginal);
+    SIZE(TypeAliasTypes);
+    SIZE(TupleTypes);
+    SIZE(PackTypes);
+    SIZE(PackExpansionTypes);
+    SIZE(PackElementTypes);
+    SIZE_AND_BYTES(MetatypeTypes);
+    SIZE_AND_BYTES(ExistentialMetatypeTypes);
+    SIZE_AND_BYTES(ArraySliceTypes);
+    SIZE_AND_BYTES(VariadicSequenceTypes);
+    SIZE_AND_BYTES(DictionaryTypes);
+    SIZE_AND_BYTES(OptionalTypes);
+    SIZE_AND_BYTES(ParenTypes);
+    SIZE_AND_BYTES(ReferenceStorageTypes);
+    SIZE_AND_BYTES(LValueTypes);
+    SIZE_AND_BYTES(InOutTypes);
+    SIZE_AND_BYTES(DependentMemberTypes);
+    SIZE(ErrorUnionTypes);
+    SIZE_AND_BYTES(PlaceholderTypes);
+    SIZE_AND_BYTES(DynamicSelfTypes);
+    SIZE_AND_BYTES(EnumTypes);
+    SIZE_AND_BYTES(StructTypes);
+    SIZE_AND_BYTES(ClassTypes);
+    SIZE_AND_BYTES(ProtocolTypes);
+    SIZE_AND_BYTES(ExistentialTypes);
+    SIZE(UnboundGenericTypes);
+    SIZE(BoundGenericTypes);
+    SIZE(ProtocolCompositionTypes);
+    SIZE(ParameterizedProtocolTypes);
+    SIZE(LayoutConstraints);
+    SIZE_AND_BYTES(OpaqueArchetypeEnvironments);
+    SIZE(FunctionTypes);
+    SIZE(NormalConformances);
+    SIZE(SelfConformances);
+    SIZE(SpecializedConformances);
+    SIZE(InheritedConformances);
+    SIZE_AND_BYTES(BuiltinConformances);
+    SIZE(PackConformances);
+    SIZE(SubstitutionMaps);
+
+#undef SIZE
+#undef SIZE_AND_BYTES
+}
+
 void AbstractFunctionDecl::setForeignErrorConvention(
                                          const ForeignErrorConvention &conv) {
   assert(hasThrows() && "setting error convention on non-throwing decl");
@@ -5192,21 +5296,11 @@ void SubstitutionMap::Storage::Profile(
   id.AddPointer(genericSig.getPointer());
   if (!genericSig) return;
 
-  // Profile those replacement types that corresponding to canonical generic
-  // parameters within the generic signature.
-  id.AddInteger(replacementTypes.size());
-
-  unsigned i = 0;
-  genericSig->forEachParam([&](GenericTypeParamType *gp, bool canonical) {
-    if (canonical)
-      id.AddPointer(replacementTypes[i].getPointer());
-    else
-      id.AddPointer(nullptr);
-    ++i;
-  });
+  // Replacement types.
+  for (auto replacementType : replacementTypes)
+    id.AddPointer(replacementType.getPointer());
 
   // Conformances.
-  id.AddInteger(conformances.size());
   for (auto conformance : conformances)
     id.AddPointer(conformance.getOpaqueValue());
 }

lib/AST/ASTDumper.cpp

@@ -3816,16 +3816,13 @@ class PrintConformance : public PrintBase {
 
     auto genericParams = genericSig.getGenericParams();
     auto replacementTypes =
-    static_cast<const SubstitutionMap &>(map).getReplacementTypesBuffer();
+    static_cast<const SubstitutionMap &>(map).getReplacementTypes();
     for (unsigned i : indices(genericParams)) {
       if (style == SubstitutionMap::DumpStyle::Minimal) {
         printFieldRaw([&](raw_ostream &out) {
           genericParams[i]->print(out);
           out << " -> ";
-          if (replacementTypes[i])
-            out << replacementTypes[i];
-          else
-            out << "<unresolved concrete type>";
+          out << replacementTypes[i];
         }, "");
       } else {
         printRecArbitrary([&](StringRef label) {
@@ -3834,8 +3831,7 @@ class PrintConformance : public PrintBase {
             genericParams[i]->print(out);
             out << " -> ";
           }, "");
-          if (replacementTypes[i])
-            printRec(replacementTypes[i]);
+          printRec(replacementTypes[i]);
           printFoot();
         });
       }

lib/AST/ASTMangler.cpp

@@ -1976,10 +1976,9 @@ void ASTMangler::appendRetroactiveConformances(Type type, GenericSignature sig)
     if (type->hasUnboundGenericType())
       return;
 
-    auto nominal = type->getAnyNominal();
-    if (!nominal) return;
+    if (!type->getAnyNominal()) return;
 
-    subMap = type->getContextSubstitutionMap(nominal);
+    subMap = type->getContextSubstitutionMap();
   }
 
   appendRetroactiveConformances(subMap, sig);

lib/AST/RequirementMachine/RequirementLowering.cpp

@@ -637,7 +637,9 @@ struct InferRequirementsWalker : public TypeWalker {
       }
     }
 
-    if (!ty->isSpecialized())
+    // Both is<ExistentialType>() and isSpecialized() end up being true if we
+    // have invalid code where a protocol is nested inside a generic nominal.
+    if (ty->is<ExistentialType>() || !ty->isSpecialized())
       return Action::Continue;
 
     // Infer from generic nominal types.