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.

lib/AST/ASTContext.cpp

@@ -3559,8 +3559,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

@@ -3082,4 +3082,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

@@ -6860,8 +6860,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;
+      }
     }
   }