New SubstitutionMap data structure

include/swift/AST/GenericEnvironment.h

@@ -18,6 +18,7 @@
 #define SWIFT_AST_GENERIC_ENVIRONMENT_H
 
 #include "swift/AST/ProtocolConformance.h"
+#include "swift/AST/SubstitutionMap.h"
 
 namespace swift {
 
@@ -60,12 +61,17 @@ class GenericEnvironment final {
   /// Derive a contextual type substitution map from a substitution array.
   /// This is just like GenericSignature::getSubstitutionMap(), except
   /// with contextual types instead of interface types.
+  SubstitutionMap
+  getSubstitutionMap(ModuleDecl *mod,
+                     GenericSignature *sig,
+                     ArrayRef<Substitution> subs) const;
+
+  /// Same as above, but updates an existing map.
   void
   getSubstitutionMap(ModuleDecl *mod,
                      GenericSignature *sig,
                      ArrayRef<Substitution> subs,
-                     TypeSubstitutionMap &subsMap,
-                     ArchetypeConformanceMap &conformanceMap) const;
+                     SubstitutionMap &subMap) const;
 
   ArrayRef<Substitution>
   getForwardingSubstitutions(ModuleDecl *M, GenericSignature *sig) const;

include/swift/AST/GenericSignature.h

@@ -26,10 +26,10 @@
 namespace swift {
 
 class ArchetypeBuilder;
+class ProtocolConformanceRef;
 class ProtocolType;
-
-using TypeConformanceMap
-  = llvm::DenseMap<TypeBase *, ArrayRef<ProtocolConformanceRef>>;
+class Substitution;
+class SubstitutionMap;
 
 /// Iterator that walks the generic parameter types declared in a generic
 /// signature and their dependent members.
@@ -170,14 +170,11 @@ class GenericSignature final : public llvm::FoldingSetNode,
 
   /// Build an interface type substitution map from a vector of Substitutions
   /// that correspond to the generic parameters in this generic signature.
-  /// The order of primary archetypes in the substitution vector must match
-  /// the order of generic parameters in getGenericParams().
-  TypeSubstitutionMap getSubstitutionMap(ArrayRef<Substitution> args) const;
+  SubstitutionMap getSubstitutionMap(ArrayRef<Substitution> args) const;
 
-  /// Variant of the above that also returns conformances.
-  void getSubstitutionMap(ArrayRef<Substitution> subs,
-                          TypeSubstitutionMap &subMap,
-                          TypeConformanceMap &conformanceMap) const;
+  /// Same as above, but updates an existing map.
+  void getSubstitutionMap(ArrayRef<Substitution> args,
+                          SubstitutionMap &subMap) const;
 
   using LookupConformanceFn =
       llvm::function_ref<ProtocolConformanceRef(CanType, Type, ProtocolType *)>;
@@ -192,8 +189,7 @@ class GenericSignature final : public llvm::FoldingSetNode,
   /// Build an array of substitutions from an interface type substitution map,
   /// using the given function to look up conformances.
   void getSubstitutions(ModuleDecl &mod,
-                        const TypeSubstitutionMap &subMap,
-                        const TypeConformanceMap &conformanceMap,
+                        const SubstitutionMap &subMap,
                         SmallVectorImpl<Substitution> &result) const;
 
   /// Return a range that iterates through first all of the generic parameters

include/swift/AST/ProtocolConformance.h

@@ -296,8 +296,7 @@ class alignas(1 << DeclAlignInBits) ProtocolConformance {
   /// applies to the substituted type.
   ProtocolConformance *subst(ModuleDecl *module,
                              Type substType,
-                             TypeSubstitutionMap &subMap,
-                             ArchetypeConformanceMap &conformanceMap);
+                             const SubstitutionMap &subMap) const;
 };
 
 /// Normal protocol conformance, which involves mapping each of the protocol

include/swift/AST/Substitution.h

@@ -17,6 +17,7 @@
 #ifndef SWIFT_AST_SUBSTITUTION_H
 #define SWIFT_AST_SUBSTITUTION_H
 
+#include "swift/AST/SubstitutionMap.h"
 #include "swift/AST/Type.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/Optional.h"
@@ -26,33 +27,7 @@ namespace llvm {
 }
 
 namespace swift {
-  class ArchetypeType;
   class GenericEnvironment;
-  class ProtocolConformanceRef;
-
-/// Data structure type used internally by Substitution::subst to track
-/// conformances applied to archetypes.
-class ArchetypeConformanceMap {
-  using ParentType = std::pair<ArchetypeType *, AssociatedTypeDecl *>;
-
-  llvm::DenseMap<ArchetypeType *, ArrayRef<ProtocolConformanceRef>> map;
-  llvm::DenseMap<ArchetypeType *, SmallVector<ParentType, 1>> parents;
-
-  Optional<ProtocolConformanceRef>
-  lookupArchetypeConformance(ProtocolDecl *proto,
-                             ArrayRef<ProtocolConformanceRef> conformances) const;
-
-public:
-  Optional<ProtocolConformanceRef>
-  lookupArchetypeConformance(ArchetypeType *replacement,
-                             ProtocolDecl *proto) const;
-
-  void addArchetypeConformances(ArchetypeType *replacement,
-                                ArrayRef<ProtocolConformanceRef> conformances);
-
-  void addArchetypeParent(ArchetypeType *replacement, ArchetypeType *parent,
-                          AssociatedTypeDecl *assocType);
-};
 
 /// Substitution - A substitution into a generic specialization.
 class Substitution {
@@ -92,8 +67,7 @@ class Substitution {
                      ArrayRef<Substitution> subs) const;
 
   Substitution subst(ModuleDecl *module,
-                     TypeSubstitutionMap &subMap,
-                     ArchetypeConformanceMap &conformanceMap) const;
+                     const SubstitutionMap &subMap) const;
 
 private:
   friend class ProtocolConformance;

include/swift/AST/SubstitutionMap.h

@@ -0,0 +1,71 @@
+//===--- SubstitutionMap.h - Swift Substitution Map ASTs --------*- C++ -*-===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2014 - 2016 Apple Inc. and the Swift project authors
+// Licensed under Apache License v2.0 with Runtime Library Exception
+//
+// See http://swift.org/LICENSE.txt for license information
+// See http://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
+//
+//===----------------------------------------------------------------------===//
+//
+// This file defines the SubstitutionMap class.
+//
+// This is a data structure type describing the mapping of abstract types to
+// replacement types, together with associated conformances to use for deriving
+// nested types.
+//
+// Depending on how the SubstitutionMap is constructed, the abstract types are
+// either archetypes or interface types. Care must be exercised to only look up
+// one or the other.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef SWIFT_AST_SUBSTITUTION_MAP_H
+#define SWIFT_AST_SUBSTITUTION_MAP_H
+
+#include "swift/AST/ProtocolConformanceRef.h"
+#include "swift/AST/Type.h"
+#include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/Optional.h"
+#include "llvm/ADT/SmallVector.h"
+
+namespace swift {
+
+class SubstitutionMap {
+  using ParentType = std::pair<CanType, AssociatedTypeDecl *>;
+
+  llvm::DenseMap<TypeBase *, Type> subMap;
+  llvm::DenseMap<TypeBase *, ArrayRef<ProtocolConformanceRef>> conformanceMap;
+  llvm::DenseMap<TypeBase *, SmallVector<ParentType, 1>> parentMap;
+
+  Optional<ProtocolConformanceRef>
+  lookupConformance(ProtocolDecl *proto,
+                    ArrayRef<ProtocolConformanceRef> conformances) const;
+
+  template<typename Fn>
+  Optional<ProtocolConformanceRef> forEachParent(CanType type, Fn fn) const;
+
+public:
+  Optional<ProtocolConformanceRef>
+  lookupConformance(CanType type, ProtocolDecl *proto) const;
+
+  const llvm::DenseMap<TypeBase *, Type> &getMap() const {
+    return subMap;
+  }
+
+  void addSubstitution(CanType type, Type replacement);
+
+  void addConformances(CanType type, ArrayRef<ProtocolConformanceRef> conformances);
+
+  void addParent(CanType type, CanType parent,
+                 AssociatedTypeDecl *assocType);
+
+  void removeType(CanType type);
+};
+
+} // end namespace swift
+
+#endif

include/swift/AST/Type.h

@@ -42,6 +42,7 @@ class NormalProtocolConformance;
 enum OptionalTypeKind : unsigned;
 class ProtocolDecl;
 class StructDecl;
+class SubstitutionMap;
 class TypeBase;
 class Type;
 class TypeWalker;
@@ -159,7 +160,7 @@ class Type {
   /// Replace references to substitutable types with new, concrete types and
   /// return the substituted result.
   ///
-  /// \param module The module in which the substitution occurs.
+  /// \param module The module to use for conformance lookups.
   ///
   /// \param substitutions The mapping from substitutable types to their
   /// replacements.
@@ -171,6 +172,18 @@ class Type {
              const TypeSubstitutionMap &substitutions,
              SubstOptions options) const;
 
+  /// Replace references to substitutable types with new, concrete types and
+  /// return the substituted result.
+  ///
+  /// \param substitutions The mapping from substitutable types to their
+  /// replacements and conformances.
+  ///
+  /// \param options Options that affect the substitutions.
+  ///
+  /// \returns the substituted type, or a null type if an error occurred.
+  Type subst(const SubstitutionMap &substitutions,
+             SubstOptions options) const;
+
   bool isPrivateStdlibType(bool whitelistProtocols=true) const;
 
   void dump() const;

include/swift/AST/Types.h

@@ -2480,7 +2480,7 @@ class GenericFunctionType : public AnyFunctionType,
   /// function type and return the resulting non-generic type.
   ///
   /// The order of Substitutions must match the order of generic parameters.
-  FunctionType *substGenericArgs(ModuleDecl *M, ArrayRef<Substitution> subs);
+  FunctionType *substGenericArgs(ArrayRef<Substitution> subs);
 
   void Profile(llvm::FoldingSetNodeID &ID) {
     Profile(ID, getGenericSignature(), getInput(), getResult(),

include/swift/SIL/TypeSubstCloner.h

@@ -55,7 +55,7 @@ class TypeSubstCloner : public SILClonerWithScopes<ImplClass> {
 
   TypeSubstCloner(SILFunction &To,
                   SILFunction &From,
-                  TypeSubstitutionMap &ContextSubs,
+                  const SubstitutionMap &ContextSubs,
                   ArrayRef<Substitution> ApplySubs,
                   SILOpenedArchetypesTracker &OpenedArchetypesTracker,
                   bool Inlining = false)
@@ -68,7 +68,7 @@ class TypeSubstCloner : public SILClonerWithScopes<ImplClass> {
 
   TypeSubstCloner(SILFunction &To,
                   SILFunction &From,
-                  TypeSubstitutionMap &ContextSubs,
+                  const SubstitutionMap &ContextSubs,
                   ArrayRef<Substitution> ApplySubs,
                   bool Inlining = false)
     : SILClonerWithScopes<ImplClass>(To, Inlining),
@@ -81,11 +81,12 @@ class TypeSubstCloner : public SILClonerWithScopes<ImplClass> {
 
 protected:
   SILType remapType(SILType Ty) {
-    return SILType::substType(Original.getModule(), SwiftMod, SubsMap, Ty);
+    return SILType::substType(Original.getModule(), SwiftMod,
+                              SubsMap.getMap(), Ty);
   }
 
   CanType remapASTType(CanType ty) {
-    return ty.subst(SwiftMod, SubsMap, None)->getCanonicalType();
+    return ty.subst(SubsMap, None)->getCanonicalType();
   }
 
   Substitution remapSubstitution(Substitution sub) {
@@ -288,7 +289,7 @@ class TypeSubstCloner : public SILClonerWithScopes<ImplClass> {
   /// The Swift module that the cloned function belongs to.
   Module *SwiftMod;
   /// The substitutions list for the specialization.
-  TypeSubstitutionMap &SubsMap;
+  const SubstitutionMap &SubsMap;
   /// The original function to specialize.
   SILFunction &Original;
   /// The substitutions used at the call site.

include/swift/SILOptimizer/Utils/GenericCloner.h

@@ -41,7 +41,7 @@ class GenericCloner : public TypeSubstCloner<GenericCloner> {
   GenericCloner(SILFunction *F,
                 IsFragile_t Fragile,
                 const ReabstractionInfo &ReInfo,
-                TypeSubstitutionMap &ContextSubs,
+                SubstitutionMap &ContextSubs,
                 ArrayRef<Substitution> ParamSubs,
                 StringRef NewName,
                 CloneCollector::CallbackType Callback)
@@ -56,7 +56,7 @@ class GenericCloner : public TypeSubstCloner<GenericCloner> {
   cloneFunction(SILFunction *F,
                 IsFragile_t Fragile,
                 const ReabstractionInfo &ReInfo,
-                TypeSubstitutionMap &ContextSubs,
+                SubstitutionMap &ContextSubs,
                 ArrayRef<Substitution> ParamSubs,
                 StringRef NewName,
                 CloneCollector::CallbackType Callback =nullptr) {

include/swift/SILOptimizer/Utils/Generics.h

@@ -138,7 +138,7 @@ class GenericFuncSpecializer {
   IsFragile_t Fragile;
   const ReabstractionInfo &ReInfo;
 
-  TypeSubstitutionMap ContextSubs;
+  SubstitutionMap ContextSubs;
   std::string ClonedName;
 public:
   GenericFuncSpecializer(SILFunction *GenericFunc,

include/swift/SILOptimizer/Utils/Local.h

@@ -138,15 +138,15 @@ void replaceDeadApply(ApplySite Old, ValueBase *New);
 
 /// \brief Return true if the substitution map contains a
 /// substitution that is an unbound generic type.
-bool hasUnboundGenericTypes(TypeSubstitutionMap &SubsMap);
+bool hasUnboundGenericTypes(const TypeSubstitutionMap &SubsMap);
 
 /// Return true if the substitution list contains a substitution
 /// that is an unbound generic.
 bool hasUnboundGenericTypes(ArrayRef<Substitution> Subs);
 
 /// \brief Return true if the substitution map contains a
 /// substitution that refers to the dynamic Self type.
-bool hasDynamicSelfTypes(TypeSubstitutionMap &SubsMap);
+bool hasDynamicSelfTypes(const TypeSubstitutionMap &SubsMap);
 
 /// \brief Return true if the substitution list contains a
 /// substitution that refers to the dynamic Self type.

include/swift/SILOptimizer/Utils/SILInliner.h

@@ -46,7 +46,7 @@ class SILInliner : public TypeSubstCloner<SILInliner> {
   };
 
   SILInliner(SILFunction &To, SILFunction &From, InlineKind IKind,
-             TypeSubstitutionMap &ContextSubs, ArrayRef<Substitution> ApplySubs,
+             SubstitutionMap &ContextSubs, ArrayRef<Substitution> ApplySubs,
              SILOpenedArchetypesTracker &OpenedArchetypesTracker,
              CloneCollector::CallbackType Callback = nullptr)
       : TypeSubstCloner<SILInliner>(To, From, ContextSubs, ApplySubs,

lib/AST/CMakeLists.txt

@@ -39,6 +39,7 @@ add_swift_library(swiftAST STATIC
   Stmt.cpp
   SourceEntityWalker.cpp
   Substitution.cpp
+  SubstitutionMap.cpp
   Type.cpp
   TypeJoinMeet.cpp
   TypeRefinementContext.cpp