[GSB] Store the root potential archetype in RequirementSource.

Use TrailingObjects to help us efficiently store the root potential
archetype within requirement sources, so we can reconstruct the
complete path from the point where a requirement was created to the
potential archetype it affects.

The test changes are because we are now dumping the root potential
archetype as part of -debug-generic-signatures.

Author: DougGregor

include/swift/AST/GenericSignatureBuilder.h

@@ -31,6 +31,7 @@
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/MapVector.h"
 #include "llvm/ADT/TinyPtrVector.h"
+#include "llvm/Support/TrailingObjects.h"
 #include <functional>
 #include <memory>
 
@@ -402,7 +403,10 @@ class GenericSignatureBuilder {
 /// other implications (e.g., introduced by a particular protocol).
 ///
 /// Requirement sources are uniqued within a generic signature builder.
-class GenericSignatureBuilder::RequirementSource : public llvm::FoldingSetNode {
+class GenericSignatureBuilder::RequirementSource final
+  : public llvm::FoldingSetNode,
+    private llvm::TrailingObjects<RequirementSource, PotentialArchetype *> {
+
 public:
   enum Kind : uint8_t {
     /// A requirement stated explicitly, e.g., in a where clause or type
@@ -468,7 +472,7 @@ class GenericSignatureBuilder::RequirementSource : public llvm::FoldingSetNode {
 private:
   /// The kind of storage we have.
   enum class StorageKind : uint8_t {
-    None,
+    RootArchetype,
     TypeRepr,
     RequirementRepr,
     ProtocolDecl,
@@ -481,6 +485,9 @@ class GenericSignatureBuilder::RequirementSource : public llvm::FoldingSetNode {
 
   /// The actual storage, described by \c storageKind.
   union {
+    /// The root archetype.
+    PotentialArchetype *rootArchetype;
+
     /// The type representation describing where the requirement came from.
     const TypeRepr *typeRepr;
 
@@ -497,13 +504,38 @@ class GenericSignatureBuilder::RequirementSource : public llvm::FoldingSetNode {
     AssociatedTypeDecl *assocType;
   } storage;
 
+  friend TrailingObjects;
+
+  /// The trailing potential archetype, for
+  size_t numTrailingObjects(OverloadToken<PotentialArchetype *>) const {
+    switch (kind) {
+    case RequirementSignatureSelf:
+      return 1;
+
+    case Explicit:
+    case Inferred:
+      return storageKind == StorageKind::RootArchetype ? 0 : 1;
+
+    case NestedTypeNameMatch:
+    case ProtocolRequirement:
+    case Superclass:
+    case Parent:
+    case Concrete:
+      return 0;
+    }
+  }
+
   /// Determines whether we have been provided with an acceptable storage kind
   /// for the given requirement source kind.
   static bool isAcceptableStorageKind(Kind kind, StorageKind storageKind);
 
   /// Retrieve the opaque storage as a single pointer, for use in uniquing.
   const void *getOpaqueStorage() const;
 
+  /// Retrieve the extra opaque storage as a single pointer, for use in
+  /// uniquing.
+  const void *getExtraOpaqueStorage() const;
+
   /// Whether this kind of requirement source is a root.
   static bool isRootKind(Kind kind) {
     switch (kind) {
@@ -528,15 +560,15 @@ class GenericSignatureBuilder::RequirementSource : public llvm::FoldingSetNode {
   /// requirement source with one of the "root" kinds.
   const RequirementSource * const parent;
 
-  RequirementSource(Kind kind, const RequirementSource *parent)
-    : kind(kind), storageKind(StorageKind::None), parent(parent) {
+  RequirementSource(Kind kind, const RequirementSource *parent,
+                    PotentialArchetype *rootArchetype)
+    : kind(kind), storageKind(StorageKind::RootArchetype), parent(parent) {
     assert((static_cast<bool>(parent) != isRootKind(kind)) &&
            "Root RequirementSource should not have parent (or vice versa)");
     assert(isAcceptableStorageKind(kind, storageKind) &&
            "RequirementSource kind/storageKind mismatch");
 
-    // Prevent uninitialized memory.
-    storage.typeRepr = nullptr;
+    storage.rootArchetype = rootArchetype;
   }
 
   RequirementSource(Kind kind, const RequirementSource *parent,
@@ -651,6 +683,9 @@ class GenericSignatureBuilder::RequirementSource : public llvm::FoldingSetNode {
   const RequirementSource *viaParent(GenericSignatureBuilder &builder,
                                      AssociatedTypeDecl *assocType) const;
 
+  /// Retrieve the potential archetype at the root.
+  PotentialArchetype *getRootPotentialArchetype() const;
+
   /// Whether the requirement can be derived from something in its path.
   ///
   /// Derived requirements will not be recorded in a minimized generic
@@ -704,15 +739,17 @@ class GenericSignatureBuilder::RequirementSource : public llvm::FoldingSetNode {
 
   /// Profiling support for \c FoldingSet.
   void Profile(llvm::FoldingSetNodeID &ID) {
-    Profile(ID, kind, parent, getOpaqueStorage());
+    Profile(ID, kind, parent, getOpaqueStorage(), getExtraOpaqueStorage());
   }
 
   /// Profiling support for \c FoldingSet.
   static void Profile(llvm::FoldingSetNodeID &ID, Kind kind,
-                      const RequirementSource *parent, const void *storage) {
+                      const RequirementSource *parent, const void *storage,
+                      const void *extraStorage) {
     ID.AddInteger(kind);
     ID.AddPointer(parent);
     ID.AddPointer(storage);
+    ID.AddPointer(extraStorage);
   }
 
   LLVM_ATTRIBUTE_DEPRECATED(

lib/AST/GenericSignatureBuilder.cpp

@@ -82,7 +82,7 @@ bool RequirementSource::isAcceptableStorageKind(Kind kind,
   switch (kind) {
   case Explicit:
     switch (storageKind) {
-    case StorageKind::None:
+    case StorageKind::RootArchetype:
     case StorageKind::TypeRepr:
     case StorageKind::RequirementRepr:
       return true;
@@ -95,7 +95,7 @@ bool RequirementSource::isAcceptableStorageKind(Kind kind,
 
   case Inferred:
     switch (storageKind) {
-    case StorageKind::None:
+    case StorageKind::RootArchetype:
     case StorageKind::TypeRepr:
       return true;
 
@@ -108,7 +108,7 @@ bool RequirementSource::isAcceptableStorageKind(Kind kind,
 
   case NestedTypeNameMatch:
     switch (storageKind) {
-    case StorageKind::None:
+    case StorageKind::RootArchetype:
         return true;
 
     case StorageKind::TypeRepr:
@@ -124,7 +124,7 @@ bool RequirementSource::isAcceptableStorageKind(Kind kind,
     case StorageKind::AssociatedTypeDecl:
         return true;
 
-    case StorageKind::None:
+    case StorageKind::RootArchetype:
     case StorageKind::TypeRepr:
     case StorageKind::ProtocolDecl:
     case StorageKind::ProtocolConformance:
@@ -138,7 +138,7 @@ bool RequirementSource::isAcceptableStorageKind(Kind kind,
     case StorageKind::ProtocolDecl:
       return true;
 
-    case StorageKind::None:
+    case StorageKind::RootArchetype:
     case StorageKind::TypeRepr:
     case StorageKind::ProtocolConformance:
     case StorageKind::RequirementRepr:
@@ -152,7 +152,7 @@ bool RequirementSource::isAcceptableStorageKind(Kind kind,
     case StorageKind::ProtocolConformance:
       return true;
 
-    case StorageKind::None:
+    case StorageKind::RootArchetype:
     case StorageKind::ProtocolDecl:
     case StorageKind::TypeRepr:
     case StorageKind::RequirementRepr:
@@ -164,9 +164,8 @@ bool RequirementSource::isAcceptableStorageKind(Kind kind,
 
 const void *RequirementSource::getOpaqueStorage() const {
   switch (storageKind) {
-  case StorageKind::None:
-    // Note: always null.
-    return storage.typeRepr;
+  case StorageKind::RootArchetype:
+    return storage.rootArchetype;
 
   case StorageKind::TypeRepr:
     return storage.typeRepr;
@@ -185,6 +184,13 @@ const void *RequirementSource::getOpaqueStorage() const {
   }
 }
 
+const void *RequirementSource::getExtraOpaqueStorage() const {
+  if (numTrailingObjects(OverloadToken<PotentialArchetype *>()) == 1)
+    return getTrailingObjects<PotentialArchetype *>()[0];
+
+  return nullptr;
+}
+
 bool RequirementSource::isDerivedRequirement() const {
   switch (kind) {
   case Explicit:
@@ -228,110 +234,126 @@ bool RequirementSource::isDerivedViaConcreteConformance() const {
   return false;
 }
 
-#define REQUIREMENT_SOURCE_FACTORY_BODY(SourceKind, Parent, Storage)       \
-  llvm::FoldingSetNodeID nodeID;                                           \
-  Profile(nodeID, Kind::SourceKind, Parent, Storage);                      \
-                                                                           \
-  void *insertPos = nullptr;                                               \
-  if (auto known =                                                         \
-        builder.Impl->RequirementSources.FindNodeOrInsertPos(nodeID,       \
-                                                            insertPos))    \
-    return known;                                                          \
-                                                                           \
-  auto result = new RequirementSource(Kind::SourceKind, Parent, Storage);  \
-  builder.Impl->RequirementSources.InsertNode(result, insertPos);          \
-  return result
-
-#define REQUIREMENT_SOURCE_FACTORY_BODY_NOSTORAGE(SourceKind, Parent)      \
+#define REQUIREMENT_SOURCE_FACTORY_BODY(                                   \
+          SourceKind, Parent, Storage, ExtraStorage,                       \
+          NumPotentialArchetypes)                                          \
   llvm::FoldingSetNodeID nodeID;                                           \
-  Profile(nodeID, Kind::SourceKind, Parent, nullptr);                      \
+  Profile(nodeID, Kind::SourceKind, Parent, Storage, ExtraStorage);        \
                                                                            \
   void *insertPos = nullptr;                                               \
   if (auto known =                                                         \
         builder.Impl->RequirementSources.FindNodeOrInsertPos(nodeID,       \
                                                             insertPos))    \
     return known;                                                          \
                                                                            \
-  auto result = new RequirementSource(Kind::SourceKind, Parent);           \
+  unsigned size =                                                          \
+    totalSizeToAlloc<PotentialArchetype *>(NumPotentialArchetypes);        \
+  void *mem = malloc(size);                                                \
+  auto result = new (mem) RequirementSource(Kind::SourceKind, Parent,      \
+                                            Storage);                      \
+  if (NumPotentialArchetypes > 0)                                          \
+    result->getTrailingObjects<PotentialArchetype *>()[0] = ExtraStorage;  \
   builder.Impl->RequirementSources.InsertNode(result, insertPos);          \
   return result
 
 const RequirementSource *RequirementSource::forAbstract(
                                                     PotentialArchetype *root) {
   auto &builder = *root->getBuilder();
-  // FIXME: Store the root
-  REQUIREMENT_SOURCE_FACTORY_BODY_NOSTORAGE(Explicit, nullptr);
+  REQUIREMENT_SOURCE_FACTORY_BODY(Explicit, nullptr, root, nullptr, 0);
 }
 
 const RequirementSource *RequirementSource::forExplicit(
                                              PotentialArchetype *root,
                                              const TypeRepr *typeRepr) {
+  // If the type representation is NULL, we have an abstract requirement
+  // source.
+  if (!typeRepr)
+    return forAbstract(root);
+
   auto &builder = *root->getBuilder();
   // FIXME: Store the root
-  REQUIREMENT_SOURCE_FACTORY_BODY(Explicit, nullptr, typeRepr);
+  REQUIREMENT_SOURCE_FACTORY_BODY(Explicit, nullptr, typeRepr, root, 1);
 }
 
 const RequirementSource *RequirementSource::forExplicit(
                                       PotentialArchetype *root,
                                       const RequirementRepr *requirementRepr) {
+  // If the requirement representation is NULL, we have an abstract requirement
+  // source.
+  if (!requirementRepr)
+    return forAbstract(root);
+
   auto &builder = *root->getBuilder();
-  // FIXME: Store the root
-  REQUIREMENT_SOURCE_FACTORY_BODY(Explicit, nullptr, requirementRepr);
+  REQUIREMENT_SOURCE_FACTORY_BODY(Explicit, nullptr, requirementRepr, root, 1);
 }
 
 const RequirementSource *RequirementSource::forInferred(
                                               PotentialArchetype *root,
                                               const TypeRepr *typeRepr) {
   auto &builder = *root->getBuilder();
-  // FIXME: Store the root
-  REQUIREMENT_SOURCE_FACTORY_BODY(Inferred, nullptr, typeRepr);
+  REQUIREMENT_SOURCE_FACTORY_BODY(Inferred, nullptr, typeRepr, root, 1);
 }
 
 const RequirementSource *RequirementSource::forRequirementSignature(
                                               PotentialArchetype *root,
                                               ProtocolDecl *protocol) {
   auto &builder = *root->getBuilder();
-  // FIXME: Store the root
-  REQUIREMENT_SOURCE_FACTORY_BODY(RequirementSignatureSelf, nullptr, protocol);
+  REQUIREMENT_SOURCE_FACTORY_BODY(RequirementSignatureSelf, nullptr, protocol,
+                                  root, 1);
 }
 
 const RequirementSource *RequirementSource::forNestedTypeNameMatch(
                                              PotentialArchetype *root) {
   auto &builder = *root->getBuilder();
-  // FIXME: Store the root
-  REQUIREMENT_SOURCE_FACTORY_BODY_NOSTORAGE(NestedTypeNameMatch, nullptr);
+  REQUIREMENT_SOURCE_FACTORY_BODY(NestedTypeNameMatch, nullptr, root, nullptr,
+                                  0);
 }
 
 const RequirementSource *RequirementSource::viaAbstractProtocolRequirement(
                                                GenericSignatureBuilder &builder,
                                                ProtocolDecl *protocol) const {
-  REQUIREMENT_SOURCE_FACTORY_BODY(ProtocolRequirement, this, protocol);
+  REQUIREMENT_SOURCE_FACTORY_BODY(ProtocolRequirement, this, protocol,
+                                  nullptr, 0);
 }
 
 const RequirementSource *RequirementSource::viaSuperclass(
                                       GenericSignatureBuilder &builder,
                                       ProtocolConformance *conformance) const {
-  REQUIREMENT_SOURCE_FACTORY_BODY(Superclass, this, conformance);
+  REQUIREMENT_SOURCE_FACTORY_BODY(Superclass, this, conformance, nullptr, 0);
 }
 
 const RequirementSource *RequirementSource::viaConcrete(
                                       GenericSignatureBuilder &builder,
                                       ProtocolConformance *conformance) const {
-  REQUIREMENT_SOURCE_FACTORY_BODY(Concrete, this, conformance);
+  REQUIREMENT_SOURCE_FACTORY_BODY(Concrete, this, conformance, nullptr, 0);
 }
 
 const RequirementSource *RequirementSource::viaParent(
                                       GenericSignatureBuilder &builder,
                                       AssociatedTypeDecl *assocType) const {
-  REQUIREMENT_SOURCE_FACTORY_BODY(Parent, this, assocType);
+  REQUIREMENT_SOURCE_FACTORY_BODY(Parent, this, assocType, nullptr, 0);
 }
 
-#undef REQUIREMENT_SOURCE_FACTORY_BODY_NOSTORAGE
 #undef REQUIREMENT_SOURCE_FACTORY_BODY
 
+PotentialArchetype *RequirementSource::getRootPotentialArchetype() const {
+  /// Find the root.
+  auto root = this;
+  while (auto parent = root->parent)
+    root = parent;
+
+  // If the root archetype is in extra storage, grab it from there.
+  if (root->numTrailingObjects(OverloadToken<PotentialArchetype *>()) == 1)
+    return root->getTrailingObjects<PotentialArchetype *>()[0];
+
+  // Otherwise, it's in inline storage.
+  assert(storageKind == StorageKind::RootArchetype);
+  return storage.rootArchetype;
+}
+
 ProtocolDecl *RequirementSource::getProtocolDecl() const {
   switch (storageKind) {
-  case StorageKind::None:
+  case StorageKind::RootArchetype:
   case StorageKind::TypeRepr:
   case StorageKind::RequirementRepr:
     return nullptr;
@@ -405,6 +427,9 @@ void RequirementSource::print(llvm::raw_ostream &out,
   if (parent) {
     parent->print(out, srcMgr);
     out << " -> ";
+  } else {
+    auto pa = getRootPotentialArchetype();
+    out << pa->getDebugName() << ": ";
   }
 
   switch (kind) {
@@ -453,7 +478,7 @@ void RequirementSource::print(llvm::raw_ostream &out,
   };
 
   switch (storageKind) {
-  case StorageKind::None:
+  case StorageKind::RootArchetype:
     break;
 
   case StorageKind::TypeRepr: