AST: Move getSugaredType() from GenericEnvironment to GenericSignature

None of this actually involves archetypes, so it can just be an
operation on the GenericSignature itself.

Author: slavapestov

include/swift/AST/GenericEnvironment.h

@@ -149,13 +149,6 @@ class alignas(1 << DeclAlignInBits) GenericEnvironment final
   std::pair<Type, ProtocolConformanceRef>
   mapConformanceRefIntoContext(Type conformingType,
                                ProtocolConformanceRef conformance) const;
-          
-  /// Get the sugared form of a generic parameter type.
-  GenericTypeParamType *getSugaredType(GenericTypeParamType *type) const;
-
-  /// Get the sugared form of a type by substituting any
-  /// generic parameter types by their sugared form.
-  Type getSugaredType(Type type) const;
 
   SubstitutionMap getForwardingSubstitutionMap() const;
 

include/swift/AST/GenericSignature.h

@@ -397,11 +397,18 @@ class alignas(1 << TypeAlignInBits) GenericSignatureImpl final
   ///   <t_0_0, t_0_1, t_1_0>
   /// then this will return 0 for t_0_0, 1 for t_0_1, and 2 for t_1_0.
   unsigned getGenericParamOrdinal(GenericTypeParamType *param) const;
-      
+
   /// Get a substitution map that maps all of the generic signature's
   /// generic parameters to themselves.
   SubstitutionMap getIdentitySubstitutionMap() const;
 
+  /// Get the sugared form of a generic parameter type.
+  GenericTypeParamType *getSugaredType(GenericTypeParamType *type) const;
+
+  /// Get the sugared form of a type by substituting any
+  /// generic parameter types by their sugared form.
+  Type getSugaredType(Type type) const;
+
   /// Whether this generic signature involves a type variable.
   bool hasTypeVariable() const;
 

include/swift/AST/PrintOptions.h

@@ -24,7 +24,7 @@
 
 namespace swift {
 class ASTPrinter;
-class GenericEnvironment;
+class GenericSignatureImpl;
 class CanType;
 class Decl;
 class Pattern;
@@ -423,8 +423,8 @@ struct PrintOptions {
   /// Replaces the name of private and internal properties of types with '_'.
   bool OmitNameOfInaccessibleProperties = false;
 
-  /// Print dependent types as references into this generic environment.
-  GenericEnvironment *GenericEnv = nullptr;
+  /// Use this signature to re-sugar dependent types.
+  const GenericSignatureImpl *GenericSig = nullptr;
 
   /// Print types with alternative names from their canonical names.
   llvm::DenseMap<CanType, Identifier> *AlternativeTypeNames = nullptr;

lib/AST/ASTPrinter.cpp

@@ -24,7 +24,7 @@
 #include "swift/AST/Decl.h"
 #include "swift/AST/Expr.h"
 #include "swift/AST/FileUnit.h"
-#include "swift/AST/GenericEnvironment.h"
+#include "swift/AST/GenericSignature.h"
 #include "swift/AST/Module.h"
 #include "swift/AST/NameLookup.h"
 #include "swift/AST/ParameterList.h"
@@ -2488,14 +2488,14 @@ void PrintAST::visitTypeAliasDecl(TypeAliasDecl *decl) {
     Printer << " = ";
     // FIXME: An inferred associated type witness type alias may reference
     // an opaque type, but OpaqueTypeArchetypes are always canonicalized
-    // so lose type sugar for generic params. Bind the generic environment so
-    // we can map params back into the generic environment and print them
+    // so lose type sugar for generic params. Bind the generic signature so
+    // we can map params back into the generic signature and print them
     // correctly.
     //
     // Remove this when we have a way to represent non-canonical archetypes
     // preserving sugar.
-    llvm::SaveAndRestore<GenericEnvironment*> setGenericEnv(Options.GenericEnv,
-                                                decl->getGenericEnvironment());
+    llvm::SaveAndRestore<const GenericSignatureImpl *> setGenericSig(
+        Options.GenericSig, decl->getGenericSignature().getPointer());
     printTypeLoc(TypeLoc(decl->getUnderlyingTypeRepr(), Ty));
     printDeclGenericRequirements(decl);
   }
@@ -4319,7 +4319,7 @@ class TypePrinter : public TypeVisitor<TypePrinter> {
     Optional<TypePrinter> subBuffer;
     PrintOptions subOptions = Options;
     if (auto substitutions = T->getPatternSubstitutions()) {
-      subOptions.GenericEnv = nullptr;
+      subOptions.GenericSig = nullptr;
       subBuffer.emplace(Printer, subOptions);
       sub = &*subBuffer;
 
@@ -4411,7 +4411,7 @@ class TypePrinter : public TypeVisitor<TypePrinter> {
       // A box layout has its own independent generic environment. Don't try
       // to print it with the environment's generic params.
       PrintOptions subOptions = Options;
-      subOptions.GenericEnv = nullptr;
+      subOptions.GenericSig = nullptr;
       TypePrinter sub(Printer, subOptions);
 
       // Capture list used here to ensure we don't print anything using `this`
@@ -4594,10 +4594,10 @@ class TypePrinter : public TypeVisitor<TypePrinter> {
         }
       }
 
-      // When printing SIL types, use a generic environment to map them from
+      // When printing SIL types, use a generic signature to map them from
       // canonical types to sugared types.
-      if (Options.GenericEnv)
-        T = Options.GenericEnv->getSugaredType(T);
+      if (Options.GenericSig)
+        T = Options.GenericSig->getSugaredType(T);
     }
 
     auto Name = T->getName();

lib/AST/Attr.cpp

@@ -590,15 +590,15 @@ static void printDifferentiableAttrArguments(
       stream << ' ';
     stream << "where ";
     std::function<Type(Type)> getInterfaceType;
-    if (!original || !original->getGenericEnvironment()) {
+    if (!original || !original->getGenericSignature()) {
       getInterfaceType = [](Type Ty) -> Type { return Ty; };
     } else {
-      // Use GenericEnvironment to produce user-friendly
+      // Use GenericSignature to produce user-friendly
       // names instead of something like 't_0_0'.
-      auto *genericEnv = original->getGenericEnvironment();
-      assert(genericEnv);
+      auto genericSig = original->getGenericSignature();
+      assert(genericSig);
       getInterfaceType = [=](Type Ty) -> Type {
-        return genericEnv->getSugaredType(Ty);
+        return genericSig->getSugaredType(Ty);
       };
     }
     interleave(requirementsToPrint, [&](Requirement req) {
@@ -933,20 +933,20 @@ bool DeclAttribute::printImpl(ASTPrinter &Printer, const PrintOptions &Options,
 
     std::function<Type(Type)> GetInterfaceType;
     auto *FnDecl = dyn_cast_or_null<AbstractFunctionDecl>(D);
-    if (!FnDecl || !FnDecl->getGenericEnvironment())
+    if (!FnDecl || !FnDecl->getGenericSignature())
       GetInterfaceType = [](Type Ty) -> Type { return Ty; };
     else {
-      // Use GenericEnvironment to produce user-friendly
+      // Use GenericSignature to produce user-friendly
       // names instead of something like t_0_0.
-      auto *GenericEnv = FnDecl->getGenericEnvironment();
-      assert(GenericEnv);
+      auto GenericSig = FnDecl->getGenericSignature();
+      assert(GenericSig);
       GetInterfaceType = [=](Type Ty) -> Type {
-        return GenericEnv->getSugaredType(Ty);
+        return GenericSig->getSugaredType(Ty);
       };
 
       if (auto sig = attr->getSpecializedSgnature()) {
         requirementsScratch = sig->requirementsNotSatisfiedBy(
-            GenericEnv->getGenericSignature());
+            GenericSig);
         requirements = requirementsScratch;
       }
     }

lib/AST/GenericEnvironment.cpp

@@ -171,27 +171,6 @@ Type GenericEnvironment::mapTypeIntoContext(GenericTypeParamType *type) const {
   return result;
 }
 
-GenericTypeParamType *GenericEnvironment::getSugaredType(
-    GenericTypeParamType *type) const {
-  for (auto *sugaredType : getGenericParams())
-    if (sugaredType->isEqual(type))
-      return sugaredType;
-
-  llvm_unreachable("missing generic parameter");
-}
-
-Type GenericEnvironment::getSugaredType(Type type) const {
-  if (!type->hasTypeParameter())
-    return type;
-
-  return type.transform([this](Type Ty) -> Type {
-    if (auto GP = dyn_cast<GenericTypeParamType>(Ty.getPointer())) {
-      return Type(getSugaredType(GP));
-    }
-    return Ty;
-  });
-}
-
 SubstitutionMap GenericEnvironment::getForwardingSubstitutionMap() const {
   auto genericSig = getGenericSignature();
   return SubstitutionMap::get(genericSig,

lib/AST/GenericSignature.cpp

@@ -1037,10 +1037,27 @@ SubstitutionMap GenericSignatureImpl::getIdentitySubstitutionMap() const {
                               MakeAbstractConformanceForGenericType());
 }
 
+GenericTypeParamType *GenericSignatureImpl::getSugaredType(
+    GenericTypeParamType *type) const {
+  unsigned ordinal = getGenericParamOrdinal(type);
+  return getGenericParams()[ordinal];
+}
+
+Type GenericSignatureImpl::getSugaredType(Type type) const {
+  if (!type->hasTypeParameter())
+    return type;
+
+  return type.transform([this](Type Ty) -> Type {
+    if (auto GP = dyn_cast<GenericTypeParamType>(Ty.getPointer())) {
+      return Type(getSugaredType(GP));
+    }
+    return Ty;
+  });
+}
+
 unsigned GenericSignatureImpl::getGenericParamOrdinal(
     GenericTypeParamType *param) const {
-  return GenericParamKey(param->getDepth(), param->getIndex())
-    .findIndexIn(getGenericParams());
+  return GenericParamKey(param).findIndexIn(getGenericParams());
 }
 
 bool GenericSignatureImpl::hasTypeVariable() const {

lib/PrintAsObjC/DeclAndTypePrinter.cpp

@@ -1909,7 +1909,7 @@ class DeclAndTypePrinter::Implementation
       assert(extension->getGenericSignature().getCanonicalSignature() ==
                  extendedClass->getGenericSignature().getCanonicalSignature() &&
              "constrained extensions or custom generic parameters?");
-      type = extendedClass->getGenericEnvironment()->getSugaredType(type);
+      type = extendedClass->getGenericSignature()->getSugaredType(type);
       decl = type->getDecl();
     }
 

lib/SIL/IR/SILPrinter.cpp

@@ -462,18 +462,20 @@ static void printSILFunctionNameAndType(
     llvm::DenseMap<CanType, Identifier> &sugaredTypeNames) {
   function->printName(OS);
   OS << " : $";
-  auto genSig =
-    function->getLoweredFunctionType()->getInvocationGenericSignature();
   auto *genEnv = function->getGenericEnvironment();
-  // If `genSig` and `genEnv` are both defined, get sugared names of generic
+  const GenericSignatureImpl *genSig = nullptr;
+
+  // If `genEnv` is defined, get sugared names of generic
   // parameter types for printing.
-  if (genSig && genEnv) {
+  if (genEnv) {
+    genSig = genEnv->getGenericSignature().getPointer();
+
     llvm::DenseSet<Identifier> usedNames;
     llvm::SmallString<16> disambiguatedNameBuf;
     unsigned disambiguatedNameCounter = 1;
     for (auto *paramTy : genSig->getGenericParams()) {
       // Get a uniqued sugared name for the generic parameter type.
-      auto sugaredTy = genEnv->getSugaredType(paramTy);
+      auto sugaredTy = genEnv->getGenericSignature()->getSugaredType(paramTy);
       Identifier name = sugaredTy->getName();
       while (!usedNames.insert(name).second) {
         disambiguatedNameBuf.clear();
@@ -495,7 +497,7 @@ static void printSILFunctionNameAndType(
     }
   }
   auto printOptions = PrintOptions::printSIL();
-  printOptions.GenericEnv = genEnv;
+  printOptions.GenericSig = genSig;
   printOptions.AlternativeTypeNames =
       sugaredTypeNames.empty() ? nullptr : &sugaredTypeNames;
   function->getLoweredFunctionType()->print(OS, printOptions);
@@ -3245,8 +3247,9 @@ void SILWitnessTable::print(llvm::raw_ostream &OS, bool Verbose) const {
     OS << "[serialized] ";
 
   getConformance()->printName(OS, Options);
-  Options.GenericEnv =
-    getConformance()->getDeclContext()->getGenericEnvironmentOfContext();
+  Options.GenericSig =
+    getConformance()->getDeclContext()->getGenericSignatureOfContext()
+      .getPointer();
 
   if (isDeclaration()) {
     OS << "\n\n";
@@ -3289,7 +3292,7 @@ void SILDefaultWitnessTable::print(llvm::raw_ostream &OS, bool Verbose) const {
   OS << getProtocol()->getName() << " {\n";
 
   PrintOptions options = PrintOptions::printSIL();
-  options.GenericEnv = Protocol->getGenericEnvironmentOfContext();
+  options.GenericSig = Protocol->getGenericSignatureOfContext().getPointer();
 
   for (auto &witness : getEntries()) {
     witness.print(OS, Verbose, options);
@@ -3440,12 +3443,17 @@ void SILSpecializeAttr::print(llvm::raw_ostream &OS) const {
   OS << "exported: " << exported << ", ";
   OS << "kind: " << kind << ", ";
 
+  auto *genericEnv = getFunction()->getGenericEnvironment();
+  GenericSignature genericSig;
+  if (genericEnv)
+    genericSig = genericEnv->getGenericSignature();
+
   ArrayRef<Requirement> requirements;
   SmallVector<Requirement, 4> requirementsScratch;
   if (auto specializedSig = getSpecializedSignature()) {
-    if (auto env = getFunction()->getGenericEnvironment()) {
+    if (genericSig) {
       requirementsScratch = specializedSig->requirementsNotSatisfiedBy(
-          env->getGenericSignature());
+          genericSig);
       requirements = requirementsScratch;
     } else {
       requirements = specializedSig->getRequirements();
@@ -3455,19 +3463,18 @@ void SILSpecializeAttr::print(llvm::raw_ostream &OS) const {
     OS << "where ";
     SILFunction *F = getFunction();
     assert(F);
-    auto GenericEnv = F->getGenericEnvironment();
     interleave(requirements,
                [&](Requirement req) {
-                 if (!GenericEnv) {
+                 if (!genericSig) {
                    req.print(OS, SubPrinter);
                    return;
                  }
                  // Use GenericEnvironment to produce user-friendly
                  // names instead of something like t_0_0.
-                 auto FirstTy = GenericEnv->getSugaredType(req.getFirstType());
+                 auto FirstTy = genericSig->getSugaredType(req.getFirstType());
                  if (req.getKind() != RequirementKind::Layout) {
                    auto SecondTy =
-                       GenericEnv->getSugaredType(req.getSecondType());
+                       genericSig->getSugaredType(req.getSecondType());
                    Requirement ReqWithDecls(req.getKind(), FirstTy, SecondTy);
                    ReqWithDecls.print(OS, SubPrinter);
                  } else {