AST: Add "re-sugaring" to GenericEnvironment

Sugared GenericTypeParamTypes point to GenericTypeParamDecls,
allowing the name of the parameter as written by the user to be
recovered. Canonical GenericTypeParamTypes on the other hand
only store a depth and index, without referencing the original
declaration.

When printing SIL, we wish to output the original generic parameter
names, even though SIL only uses canonical types. Previously,
we used to accomplish this by mapping the generic parameter to an
archetype and printing the name of the archetype. This was not
adequate if multiple generic parameters mapped to the same
archetype, or if a generic parameter was mapped to a concrete type.

The new approach preserves the original sugared types in the
GenericEnvironment, adding a new GenericEnvironment::getSugaredType()
method.

There are also some other assorted simplifications made possible
by this.

Unfortunately this makes GenericEnvironments use a bit more memory,
however I have more improvements coming that will offset the gains,
in addition to making substitution lists smaller also.

Author: slavapestov

include/swift/AST/ArchetypeBuilder.h

@@ -254,21 +254,11 @@ class ArchetypeBuilder {
                            GenericEnvironment *genericEnv,
                            bool treatRequirementsAsExplicit = false);
 
-  /// \brief Get a generic signature based on the provided complete list
-  /// of generic parameter types.
-  ///
-  /// \returns a generic signature built from the provided list of
-  ///          generic parameter types.
-  GenericSignature *
-  getGenericSignature(ArrayRef<GenericTypeParamType *> genericParamsTypes);
+  /// \brief Build the generic signature.
+  GenericSignature *getGenericSignature();
 
-  /// \brief Get a generic context based on the complete list of generic
-  /// parameter types.
-  ///
-  /// \returns a generic context built from the provided list of
-  ///          generic parameter types.
-  GenericEnvironment *getGenericEnvironment(
-      ArrayRef<GenericTypeParamType *> genericParamsTypes);
+  /// \brief Build the generic environment.
+  GenericEnvironment *getGenericEnvironment();
 
   /// Infer requirements from the given type, recursively.
   ///

include/swift/AST/GenericEnvironment.h

@@ -28,10 +28,15 @@ class GenericTypeParamType;
 /// Describes the mapping between archetypes and interface types for the
 /// generic parameters of a DeclContext.
 class GenericEnvironment final {
+  SmallVector<GenericTypeParamType *, 4> GenericParams;
   TypeSubstitutionMap ArchetypeToInterfaceMap;
   TypeSubstitutionMap InterfaceToArchetypeMap;
 
 public:
+  ArrayRef<GenericTypeParamType *> getGenericParams() const {
+    return GenericParams;
+  }
+
   const TypeSubstitutionMap &getArchetypeToInterfaceMap() const {
     return ArchetypeToInterfaceMap;
   }
@@ -40,10 +45,13 @@ class GenericEnvironment final {
     return InterfaceToArchetypeMap;
   }
 
-  explicit GenericEnvironment(TypeSubstitutionMap interfaceToArchetypeMap);
+  GenericEnvironment(ArrayRef<GenericTypeParamType *> genericParamTypes,
+                     TypeSubstitutionMap interfaceToArchetypeMap);
 
-  static GenericEnvironment *get(ASTContext &ctx,
-                                 TypeSubstitutionMap interfaceToArchetypeMap);
+  static
+  GenericEnvironment * get(ASTContext &ctx,
+                           ArrayRef<GenericTypeParamType *> genericParamTypes,
+                           TypeSubstitutionMap interfaceToArchetypeMap);
 
   /// Make vanilla new/delete illegal.
   void *operator new(size_t Bytes) = delete;
@@ -62,6 +70,9 @@ class GenericEnvironment final {
   /// Map a generic parameter type to a contextual type.
   Type mapTypeIntoContext(GenericTypeParamType *type) const;
 
+  /// Get the sugared form of a generic parameter type.
+  GenericTypeParamType *getSugaredType(GenericTypeParamType *type) const;
+
   /// Derive a contextual type substitution map from a substitution array.
   /// This is just like GenericSignature::getSubstitutionMap(), except
   /// with contextual types instead of interface types.

include/swift/Serialization/DeclTypeRecordNodes.def

@@ -169,6 +169,7 @@ OTHER(GENERIC_PARAM_LIST, 240)
 TRAILING_INFO(GENERIC_PARAM)
 TRAILING_INFO(GENERIC_REQUIREMENT)
 TRAILING_INFO(GENERIC_ENVIRONMENT)
+TRAILING_INFO(SIL_GENERIC_ENVIRONMENT)
 
 OTHER(LOCAL_DISCRIMINATOR, 248)
 OTHER(PRIVATE_DISCRIMINATOR, 249)

include/swift/Serialization/ModuleFormat.h

@@ -54,7 +54,7 @@ const uint16_t VERSION_MAJOR = 0;
 /// in source control, you should also update the comment to briefly
 /// describe what change you made. The content of this comment isn't important;
 /// it just ensures a conflict if two people change the module format.
-const uint16_t VERSION_MINOR = 273; // Last change: partial_apply ownership control
+const uint16_t VERSION_MINOR = 274; // Last change: SIL generic environment sugar
 
 using DeclID = PointerEmbeddedInt<unsigned, 31>;
 using DeclIDField = BCFixed<31>;
@@ -1095,21 +1095,16 @@ namespace decls_block {
     DeclIDField // Typealias
   >;
 
-  // Subtlety here: GENERIC_ENVIRONMENT is serialized for both Decls and
-  // SILFunctions.
-  //
-  // For Decls, the interface type is non-canonical, so it points back
-  // to the GenericParamListDecl. This allows us to use the serialized
-  // GENERIC_ENVIRONMENT records to form the GenericSignature, as well.
-  // The type is canonicalized when forming the actual GenericEnvironment
-  // instance.
-  //
-  // For SILFunctions, the interface type below is always canonical,
-  // since SILFunctions never point back to any original
-  // GenericTypeParamDecls.
   using GenericEnvironmentLayout = BCRecordLayout<
     GENERIC_ENVIRONMENT,
-    TypeIDField,                 // interface type
+    TypeIDField,                 // sugared interface type
+    TypeIDField                  // contextual type
+  >;
+
+  using SILGenericEnvironmentLayout = BCRecordLayout<
+    SIL_GENERIC_ENVIRONMENT,
+    IdentifierIDField,           // generic parameter name
+    TypeIDField,                 // canonical interface type
     TypeIDField                  // contextual type
   >;
 

lib/AST/ASTContext.cpp

@@ -3567,8 +3567,10 @@ GenericSignature *GenericSignature::get(ArrayRef<GenericTypeParamType *> params,
 
 GenericEnvironment *
 GenericEnvironment::get(ASTContext &ctx,
+                        ArrayRef<GenericTypeParamType *> genericParamTypes,
                         TypeSubstitutionMap interfaceToArchetypeMap) {
-  return new (ctx) GenericEnvironment(interfaceToArchetypeMap);
+  return new (ctx) GenericEnvironment(genericParamTypes,
+                                      interfaceToArchetypeMap);
 }
 
 void DeclName::CompoundDeclName::Profile(llvm::FoldingSetNodeID &id,

lib/AST/ASTDumper.cpp

@@ -3090,4 +3090,7 @@ void GenericEnvironment::dump() const {
     pair.first->dump();
     pair.second->dump();
   }
+  llvm::errs() << "Generic parameters:\n";
+  for (auto paramTy : getGenericParams())
+    paramTy->dump();
 }

lib/AST/ASTPrinter.cpp

@@ -3418,14 +3418,6 @@ class TypePrinter : public TypeVisitor<TypePrinter> {
     Printer.printTypePre(TypeLoc::withoutLoc(T));
     SWIFT_DEFER { Printer.printTypePost(TypeLoc::withoutLoc(T)); };
 
-    // If we have an alternate name for this type, use it.
-    if (Options.AlternativeTypeNames) {
-      auto found = Options.AlternativeTypeNames->find(T.getCanonicalTypeOrNull());
-      if (found != Options.AlternativeTypeNames->end()) {
-        Printer << found->second.str();
-        return;
-      }
-    }
     super::visit(T);
   }
 
@@ -3992,6 +3984,14 @@ class TypePrinter : public TypeVisitor<TypePrinter> {
         Printer << ".";
       }
 
+      if (Options.AlternativeTypeNames) {
+        auto found = Options.AlternativeTypeNames->find(T->getCanonicalType());
+        if (found != Options.AlternativeTypeNames->end()) {
+          Printer << found->second.str();
+          return;
+        }
+      }
+
       if (T->getName().empty())
         Printer << "<anonymous>";
       else {
@@ -4012,9 +4012,21 @@ class TypePrinter : public TypeVisitor<TypePrinter> {
   }
 
   void visitGenericTypeParamType(GenericTypeParamType *T) {
-    // Substitute a context archetype if we have context generic params.
-    if (Options.GenericEnv)
-      return visit(Options.GenericEnv->mapTypeIntoContext(T));
+    if (T->getDecl() == nullptr) {
+      // If we have an alternate name for this type, use it.
+      if (Options.AlternativeTypeNames) {
+        auto found = Options.AlternativeTypeNames->find(T->getCanonicalType());
+        if (found != Options.AlternativeTypeNames->end()) {
+          Printer << found->second.str();
+          return;
+        }
+      }
+
+      // When printing SIL types, use a generic environment to map them from
+      // canonical types to sugared types.
+      if (Options.GenericEnv)
+        T = Options.GenericEnv->getSugaredType(T);
+    }
 
     auto Name = T->getName();
     if (Name.empty())

lib/AST/ArchetypeBuilder.cpp

@@ -2096,27 +2096,31 @@ static void collectRequirements(ArchetypeBuilder &builder,
   });
 }
 
-GenericSignature *ArchetypeBuilder::getGenericSignature(
-    ArrayRef<GenericTypeParamType *> genericParamTypes) {
+GenericSignature *ArchetypeBuilder::getGenericSignature() {
   // Collect the requirements placed on the generic parameter types.
+  SmallVector<GenericTypeParamType *, 4> genericParamTypes;
+  for (auto pair : Impl->PotentialArchetypes) {
+    auto paramTy = pair.second->getGenericParam();
+    genericParamTypes.push_back(paramTy);
+  }
+
   SmallVector<Requirement, 4> requirements;
   collectRequirements(*this, genericParamTypes, requirements);
 
   auto sig = GenericSignature::get(genericParamTypes, requirements);
   return sig;
 }
 
-GenericEnvironment *ArchetypeBuilder::getGenericEnvironment(
-    ArrayRef<GenericTypeParamType *> genericParamTypes) {
+GenericEnvironment *ArchetypeBuilder::getGenericEnvironment() {
+  SmallVector<GenericTypeParamType *, 4> genericParamTypes;
   TypeSubstitutionMap interfaceToArchetypeMap;
 
-  for (auto paramTy : genericParamTypes) {
-    auto known = Impl->PotentialArchetypes.find(
-                   GenericTypeParamKey::forType(paramTy));
-    assert(known != Impl->PotentialArchetypes.end());
+  for (auto pair : Impl->PotentialArchetypes) {
+    auto paramTy = pair.second->getGenericParam();
+    genericParamTypes.push_back(paramTy);
 
-    auto archetypeTy = known->second->getType(*this).getAsArchetype();
-    auto concreteTy = known->second->getType(*this).getAsConcreteType();
+    auto archetypeTy = pair.second->getType(*this).getAsArchetype();
+    auto concreteTy = pair.second->getType(*this).getAsConcreteType();
     if (archetypeTy)
       interfaceToArchetypeMap[paramTy] = archetypeTy;
     else if (concreteTy)
@@ -2125,6 +2129,7 @@ GenericEnvironment *ArchetypeBuilder::getGenericEnvironment(
       llvm_unreachable("broken generic parameter");
   }
 
-  return GenericEnvironment::get(Context, interfaceToArchetypeMap);
+  return GenericEnvironment::get(Context, genericParamTypes,
+                                 interfaceToArchetypeMap);
 }
 

lib/AST/Builtins.cpp

@@ -205,7 +205,8 @@ getBuiltinGenericFunction(Identifier Id,
   GenericSignature *Sig =
       GenericSignature::get(GenericParamTypes, requirements);
   GenericEnvironment *Env =
-      GenericEnvironment::get(Context, InterfaceToArchetypeMap);
+      GenericEnvironment::get(Context, GenericParamTypes,
+                              InterfaceToArchetypeMap);
 
   Type InterfaceType = GenericFunctionType::get(Sig, ArgParamType, ResType,
                                                 AnyFunctionType::ExtInfo());

lib/AST/GenericEnvironment.cpp

@@ -20,10 +20,14 @@
 using namespace swift;
 
 GenericEnvironment::GenericEnvironment(
+    ArrayRef<GenericTypeParamType *> genericParamTypes,
     TypeSubstitutionMap interfaceToArchetypeMap) {
 
   assert(!interfaceToArchetypeMap.empty());
 
+  for (auto *paramTy : genericParamTypes)
+    GenericParams.push_back(paramTy);
+
   // Build a mapping in both directions, making sure to canonicalize the
   // interface type where it is used as a key, so that substitution can
   // find them, and to preserve sugar otherwise, so that
@@ -74,6 +78,15 @@ Type GenericEnvironment::mapTypeIntoContext(GenericTypeParamType *type) const {
   return found->second;
 }
 
+GenericTypeParamType *GenericEnvironment::getSugaredType(
+    GenericTypeParamType *type) const {
+  for (auto *sugaredType : GenericParams)
+    if (sugaredType->isEqual(type))
+      return sugaredType;
+
+  llvm_unreachable("missing generic parameter");
+}
+
 ArrayRef<Substitution>
 GenericEnvironment::getForwardingSubstitutions(
     ModuleDecl *M, GenericSignature *sig) const {

lib/ClangImporter/ImportDecl.cpp

@@ -6861,8 +6861,8 @@ buildGenericSignature(GenericParamList *genericParams,
         param->getDeclaredType()->castTo<GenericTypeParamType>());
   }
 
-  auto *sig = builder.getGenericSignature(genericParamTypes);
-  auto *env = builder.getGenericEnvironment(genericParamTypes);
+  auto *sig = builder.getGenericSignature();
+  auto *env = builder.getGenericEnvironment();
 
   return std::make_pair(sig, env);
 }

lib/SIL/SILPrinter.cpp

@@ -30,6 +30,7 @@
 #include "swift/SIL/SILVTable.h"
 #include "swift/AST/Decl.h"
 #include "swift/AST/Expr.h"
+#include "swift/AST/GenericEnvironment.h"
 #include "swift/AST/Module.h"
 #include "swift/AST/PrintOptions.h"
 #include "swift/AST/Types.h"
@@ -1835,21 +1836,25 @@ void SILFunction::print(SILPrintContext &PrintCtx) const {
     llvm::SmallString<16> disambiguatedNameBuf;
     unsigned disambiguatedNameCounter = 1;
     for (auto *paramTy : sig->getGenericParams()) {
-      auto *archetypeTy = mapTypeIntoContext(paramTy)->getAs<ArchetypeType>();
-      if (!archetypeTy)
-        continue;
-
-      Identifier name = archetypeTy->getName();
+      auto sugaredTy = env->getSugaredType(paramTy);
+      Identifier name = sugaredTy->getName();
       while (!UsedNames.insert(name).second) {
         disambiguatedNameBuf.clear();
         {
           llvm::raw_svector_ostream names(disambiguatedNameBuf);
-          names << archetypeTy->getName() << disambiguatedNameCounter++;
+          names << sugaredTy->getName() << disambiguatedNameCounter++;
         }
         name = getASTContext().getIdentifier(disambiguatedNameBuf);
       }
-      if (name != archetypeTy->getName())
-        Aliases[CanType(archetypeTy)] = name;
+      if (name != sugaredTy->getName()) {
+        Aliases[paramTy->getCanonicalType()] = name;
+
+        // Also for the archetype
+        auto archetypeTy = env->mapTypeIntoContext(paramTy)
+            ->getAs<ArchetypeType>();
+        if (archetypeTy)
+          Aliases[archetypeTy->getCanonicalType()] = name;
+      }
     }
   }
 

lib/SILGen/SILGenDecl.cpp

@@ -1780,7 +1780,7 @@ substSelfTypeIntoProtocolRequirementType(SILModule &M,
   if (!allParams.empty()) {
     builder.finalize(SourceLoc());
 
-    auto *sig = builder.getGenericSignature(allParams);
+    auto *sig = builder.getGenericSignature();
 
     return cast<GenericFunctionType>(
       GenericFunctionType::get(sig, input, result, reqtTy->getExtInfo())
@@ -1800,6 +1800,7 @@ getSubstitutedGenericEnvironment(SILModule &M,
     return reqtEnv;
   }
 
+  SmallVector<GenericTypeParamType *, 4> genericParamTypes;
   TypeSubstitutionMap witnessContextParams;
 
   auto selfTy = conformance->getProtocol()->getSelfInterfaceType()
@@ -1809,8 +1810,14 @@ getSubstitutedGenericEnvironment(SILModule &M,
   // the conformance (which might not be the same as the generic
   // context of the witness, if the witness is defined in a
   // superclass, concrete extension or protocol extension).
-  if (auto *outerEnv = conformance->getGenericEnvironment())
+  if (auto *outerEnv = conformance->getGenericEnvironment()) {
     witnessContextParams = outerEnv->getInterfaceToArchetypeMap();
+    for (auto *paramTy : outerEnv->getGenericParams())
+      genericParamTypes.push_back(paramTy);
+  }
+
+  for (auto *paramTy : reqtEnv->getGenericParams().slice(1))
+    genericParamTypes.push_back(paramTy);
 
   // Inner generic parameters come from the requirement and
   // also map to the archetypes of the requirement.
@@ -1824,7 +1831,8 @@ getSubstitutedGenericEnvironment(SILModule &M,
   }
 
   if (!witnessContextParams.empty())
-    return GenericEnvironment::get(M.getASTContext(), witnessContextParams);
+    return GenericEnvironment::get(M.getASTContext(), genericParamTypes,
+                                   witnessContextParams);
 
   return nullptr;
 }