AST: Allow non-canonical opaque archetypes

include/swift/AST/GenericEnvironment.h

@@ -61,13 +61,13 @@ struct OpaqueEnvironmentData {
 };
 
 /// Extra data in a generic environment for an opened existential.
-struct OpenedExistentialEnvironmentData {
+struct ExistentialEnvironmentData {
   Type existential;
   UUID uuid;
 };
 
 /// Extra data in a generic environment for an opened pack element.
-struct OpenedElementEnvironmentData {
+struct ElementEnvironmentData {
   UUID uuid;
   CanGenericTypeParamType shapeClass;
 };
@@ -85,37 +85,38 @@ class alignas(1 << DeclAlignInBits) GenericEnvironment final
         GenericEnvironment,
         SubstitutionMap,
         OpaqueEnvironmentData,
-        OpenedExistentialEnvironmentData,
-        OpenedElementEnvironmentData,
+        ExistentialEnvironmentData,
+        ElementEnvironmentData,
         Type> {
 public:
-  enum class Kind {
+  enum class Kind: uint8_t {
     /// A normal generic environment, determined only by its generic
     /// signature.
     Primary,
     /// A generic environment describing an opaque type archetype.
     Opaque,
     /// A generic environment describing an opened existential archetype.
-    OpenedExistential,
+    Existential,
     /// A generic environment describing an opened element type of a
     /// pack archetype inside a pack expansion expression.
-    OpenedElement,
+    Element,
   };
 
   class NestedTypeStorage;
 
 private:
-  mutable llvm::PointerIntPair<GenericSignature, 2, Kind> SignatureAndKind{
-      GenericSignature(), Kind::Primary};
+  GenericSignature sig;
   NestedTypeStorage *nestedTypeStorage = nullptr;
+  Kind kind;
+  bool canonical;
 
   friend TrailingObjects;
   friend OpaqueTypeArchetypeType;
 
   size_t numTrailingObjects(OverloadToken<SubstitutionMap>) const;
   size_t numTrailingObjects(OverloadToken<OpaqueEnvironmentData>) const;
-  size_t numTrailingObjects(OverloadToken<OpenedExistentialEnvironmentData>) const;
-  size_t numTrailingObjects(OverloadToken<OpenedElementEnvironmentData>) const;
+  size_t numTrailingObjects(OverloadToken<ExistentialEnvironmentData>) const;
+  size_t numTrailingObjects(OverloadToken<ElementEnvironmentData>) const;
   size_t numTrailingObjects(OverloadToken<Type>) const;
 
   /// Retrieve the array containing the context types associated with the
@@ -168,10 +169,13 @@ class alignas(1 << DeclAlignInBits) GenericEnvironment final
 
 public:
   GenericSignature getGenericSignature() const {
-    return SignatureAndKind.getPointer();
+    return sig;
   }
 
-  Kind getKind() const { return SignatureAndKind.getInt(); }
+  Kind getKind() const { return kind; }
+
+  /// Returns if the archetypes from this environment are canonical types.
+  bool isCanonical() const { return canonical; }
 
   ArrayRef<GenericTypeParamType *> getGenericParams() const;
 

include/swift/AST/TypeNodes.def

@@ -176,7 +176,7 @@ TYPE(DynamicSelf, Type)
 ABSTRACT_TYPE(Substitutable, Type)
   ABSTRACT_TYPE(Archetype, SubstitutableType)
     ALWAYS_CANONICAL_TYPE(PrimaryArchetype, ArchetypeType)
-    ALWAYS_CANONICAL_TYPE(OpaqueTypeArchetype, ArchetypeType)
+    TYPE(OpaqueTypeArchetype, ArchetypeType)
     ABSTRACT_TYPE(LocalArchetype, ArchetypeType)
       ALWAYS_CANONICAL_TYPE(ExistentialArchetype, LocalArchetypeType)
       ALWAYS_CANONICAL_TYPE(ElementArchetype, LocalArchetypeType)

include/swift/AST/Types.h

@@ -7009,7 +7009,7 @@ class ArchetypeType : public SubstitutableType,
   }
 protected:
   ArchetypeType(TypeKind Kind,
-                const ASTContext &C,
+                const ASTContext *C,
                 RecursiveTypeProperties properties,
                 Type InterfaceType,
                 ArrayRef<ProtocolDecl *> ConformsTo,
@@ -7088,7 +7088,8 @@ class OpaqueTypeArchetypeType final : public ArchetypeType,
   }
 
 private:
-  OpaqueTypeArchetypeType(GenericEnvironment *environment,
+  OpaqueTypeArchetypeType(const ASTContext *ctx,
+                          GenericEnvironment *environment,
                           RecursiveTypeProperties properties,
                           Type interfaceType,
                           ArrayRef<ProtocolDecl*> conformsTo,
@@ -7232,11 +7233,13 @@ class ExistentialArchetypeType final : public LocalArchetypeType,
   }
 
 private:
-  ExistentialArchetypeType(GenericEnvironment *environment, Type interfaceType,
-                      ArrayRef<ProtocolDecl *> conformsTo,
-                      Type superclass,
-                      LayoutConstraint layout,
-                      RecursiveTypeProperties properties);
+  ExistentialArchetypeType(const ASTContext *ctx,
+                           GenericEnvironment *environment,
+                           Type interfaceType,
+                           ArrayRef<ProtocolDecl *> conformsTo,
+                           Type superclass,
+                           LayoutConstraint layout,
+                           RecursiveTypeProperties properties);
 };
 BEGIN_CAN_TYPE_WRAPPER(ExistentialArchetypeType, LocalArchetypeType)
 END_CAN_TYPE_WRAPPER(ExistentialArchetypeType, LocalArchetypeType)
@@ -7309,7 +7312,7 @@ class ElementArchetypeType final : public LocalArchetypeType,
   }
 
 private:
-  ElementArchetypeType(const ASTContext &ctx,
+  ElementArchetypeType(const ASTContext *ctx,
                        GenericEnvironment *environment, Type interfaceType,
                        ArrayRef<ProtocolDecl *> conformsTo, Type superclass,
                        LayoutConstraint layout);

lib/AST/ASTContext.cpp

@@ -5731,7 +5731,8 @@ OpaqueTypeArchetypeType *OpaqueTypeArchetypeType::getNew(
   ASTContext &ctx = interfaceType->getASTContext();
   auto mem = ctx.Allocate(size, alignof(OpaqueTypeArchetypeType), arena);
   return ::new (mem)
-      OpaqueTypeArchetypeType(environment, properties, interfaceType,
+      OpaqueTypeArchetypeType(environment->isCanonical() ? &ctx : nullptr,
+                              environment, properties, interfaceType,
                               conformsTo, superclass, layout);
 }
 
@@ -5759,6 +5760,7 @@ CanTypeWrapper<ExistentialArchetypeType> ExistentialArchetypeType::getNew(
   void *mem = ctx.Allocate(size, alignof(ExistentialArchetypeType), arena);
 
   return CanExistentialArchetypeType(::new (mem) ExistentialArchetypeType(
+      environment->isCanonical() ? &ctx : nullptr,
       environment, interfaceType, conformsTo, superclass, layout,
       properties));
 }
@@ -6014,8 +6016,8 @@ GenericEnvironment *GenericEnvironment::forPrimary(GenericSignature signature) {
   unsigned numGenericParams = signature.getGenericParams().size();
   size_t bytes = totalSizeToAlloc<SubstitutionMap,
                                   OpaqueEnvironmentData,
-                                  OpenedExistentialEnvironmentData,
-                                  OpenedElementEnvironmentData, Type>(
+                                  ExistentialEnvironmentData,
+                                  ElementEnvironmentData, Type>(
       0, 0, 0, 0, numGenericParams);
   void *mem = ctx.Allocate(bytes, alignof(GenericEnvironment));
   return new (mem) GenericEnvironment(signature);
@@ -6025,11 +6027,6 @@ GenericEnvironment *GenericEnvironment::forPrimary(GenericSignature signature) {
 /// outer substitutions.
 GenericEnvironment *GenericEnvironment::forOpaqueType(
     OpaqueTypeDecl *opaque, SubstitutionMap subs) {
-  // TODO: We could attempt to preserve type sugar in the substitution map.
-  // Currently archetypes are assumed to be always canonical in many places,
-  // though, so doing so would require fixing those places.
-  subs = subs.getCanonical();
-
   auto &ctx = opaque->getASTContext();
 
   auto properties = ArchetypeType::archetypeProperties(
@@ -6044,9 +6041,9 @@ GenericEnvironment *GenericEnvironment::forOpaqueType(
     auto signature = opaque->getOpaqueInterfaceGenericSignature();
     unsigned numGenericParams = signature.getGenericParams().size();
     size_t bytes = totalSizeToAlloc<SubstitutionMap,
-                                  OpaqueEnvironmentData,
-                                    OpenedExistentialEnvironmentData,
-                                    OpenedElementEnvironmentData, Type>(
+                                    OpaqueEnvironmentData,
+                                    ExistentialEnvironmentData,
+                                    ElementEnvironmentData, Type>(
         1, 1, 0, 0, numGenericParams);
     void *mem = ctx.Allocate(bytes, alignof(GenericEnvironment), arena);
     env = new (mem) GenericEnvironment(signature, opaque, subs);
@@ -6108,8 +6105,8 @@ GenericEnvironment::forOpenedExistential(
   unsigned numGenericParams = signature.getGenericParams().size();
   size_t bytes = totalSizeToAlloc<SubstitutionMap,
                                   OpaqueEnvironmentData,
-                                  OpenedExistentialEnvironmentData,
-                                  OpenedElementEnvironmentData, Type>(
+                                  ExistentialEnvironmentData,
+                                  ElementEnvironmentData, Type>(
       1, 0, 1, 0, numGenericParams);
   void *mem = ctx.Allocate(bytes, alignof(GenericEnvironment));
   auto *genericEnv =
@@ -6146,8 +6143,8 @@ GenericEnvironment::forOpenedElement(GenericSignature signature,
   unsigned numOpenedParams = signature.getInnermostGenericParams().size();
   size_t bytes = totalSizeToAlloc<SubstitutionMap,
                                   OpaqueEnvironmentData,
-                                  OpenedExistentialEnvironmentData,
-                                  OpenedElementEnvironmentData,
+                                  ExistentialEnvironmentData,
+                                  ElementEnvironmentData,
                                   Type>(
       1, 0, 0, 1, numGenericParams + numOpenedParams);
   void *mem = ctx.Allocate(bytes, alignof(GenericEnvironment));

lib/AST/ASTPrinter.cpp

@@ -3634,16 +3634,6 @@ void PrintAST::visitTypeAliasDecl(TypeAliasDecl *decl) {
 
   if (ShouldPrint) {
     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 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.
-    PrintOptions::OverrideScope scope(Options);
-    OVERRIDE_PRINT_OPTION(scope, GenericSig, decl->getGenericSignature().getPointer());
     printTypeLoc(TypeLoc(decl->getUnderlyingTypeRepr(), Ty));
     printDeclGenericRequirements(decl);
   }
@@ -7500,11 +7490,6 @@ class TypePrinter : public TypeVisitor<TypePrinter, void, NonRecursivePrintOptio
       if (auto existential = constraint->getAs<ExistentialType>())
         constraint = existential->getConstraintType();
 
-      // Opaque archetype substitutions are always canonical, so re-sugar the
-      // constraint type using the owning declaration's generic parameter names.
-      if (genericSig)
-        constraint = genericSig->getSugaredType(constraint);
-
       visit(constraint);
       return;
     }

lib/AST/CaptureInfo.cpp

@@ -77,7 +77,7 @@ CaptureInfo::CaptureInfo(ASTContext &ctx, ArrayRef<CapturedValue> captures,
   // This is the only kind of local generic environment we can capture right now.
 #ifndef NDEBUG
   for (auto *env : genericEnv) {
-    assert(env->getKind() == GenericEnvironment::Kind::OpenedElement);
+    assert(env->getKind() == GenericEnvironment::Kind::Element);
   }
 #endif