[AST] Preserve type sugar for generic typealiases

Introduce a new Type node, BoundNameAliasType, which describes a
reference to a typealias that requires substitutions to produce the
underlying type. This new type node is used both for references to
generic typealiases and for references to (non-generic) typealiases
that occur within generic contexts, e.g., Array<Int>.Element.

At present, the new type node is mainly useful in preserving type
sugar for diagnostics purposes, as well as being reflected in other
tools (indexing, code completion, etc.). The intent is to completely
replace NameAliasType in the future.

Author: DougGregor

include/swift/AST/SubstitutionMap.h

@@ -36,6 +36,10 @@
 #include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/ADT/SmallVector.h"
 
+namespace llvm {
+  class FoldingSetNodeID;
+}
+
 namespace swift {
 
 class GenericSignature;
@@ -174,6 +178,9 @@ class SubstitutionMap {
 
   LLVM_ATTRIBUTE_DEPRECATED(void dump() const, "only for use in the debugger");
 
+  /// Profile the substitution map, for use with LLVM's FoldingSet.
+  void profile(llvm::FoldingSetNodeID &id) const;
+
 private:
   friend class GenericSignature;
   friend class GenericEnvironment;

include/swift/AST/TypeNodes.def

@@ -146,6 +146,7 @@ UNCHECKED_TYPE(TypeVariable, Type)
 ABSTRACT_SUGARED_TYPE(Sugar, Type)
   SUGARED_TYPE(Paren, SugarType)
   SUGARED_TYPE(NameAlias, SugarType)
+  SUGARED_TYPE(BoundNameAlias, SugarType)
   ABSTRACT_SUGARED_TYPE(SyntaxSugar, SugarType)
     ABSTRACT_SUGARED_TYPE(UnarySyntaxSugar, SyntaxSugarType)
       SUGARED_TYPE(ArraySlice, UnarySyntaxSugarType)

include/swift/AST/Types.h

@@ -373,6 +373,16 @@ class alignas(1 << TypeAlignInBits) TypeBase {
     GenericArgCount : 32
   );
 
+  SWIFT_INLINE_BITFIELD_FULL(BoundNameAliasType, SugarType, 1+16,
+    : NumPadBits,
+
+    /// Whether we have a parent type.
+    HasParent : 1,
+
+    /// The number of substitutions.
+    NumSubstitutions : 16
+  );
+
   } Bits;
 
 protected:
@@ -1560,6 +1570,81 @@ class NameAliasType : public SugarType {
   }
 };
 
+/// A reference to a type alias that is somehow generic, along with the
+/// set of substitutions to apply to make the type concrete.
+class BoundNameAliasType final
+  : public SugarType, public llvm::FoldingSetNode,
+    llvm::TrailingObjects<BoundNameAliasType, Type, Substitution>
+{
+  TypeAliasDecl *typealias;
+
+  friend class ASTContext;
+  friend TrailingObjects;
+
+  BoundNameAliasType(TypeAliasDecl *typealias, Type parent,
+                     const SubstitutionMap &substitutions, Type underlying,
+                     RecursiveTypeProperties properties);
+
+  unsigned getNumSubstitutions() const {
+    return Bits.BoundNameAliasType.NumSubstitutions;
+  }
+
+  size_t numTrailingObjects(OverloadToken<Type>) const {
+    return Bits.BoundNameAliasType.HasParent ? 1 : 0;
+  }
+
+  size_t numTrailingObjects(OverloadToken<Substitution>) const {
+    return getNumSubstitutions();
+  }
+
+  /// Retrieve the set of substitutions to be applied to the declaration to
+  /// produce the underlying type.
+  SubstitutionList getSubstitutionList() const {
+    return {getTrailingObjects<Substitution>(), getNumSubstitutions()};
+  }
+
+public:
+  static BoundNameAliasType *get(TypeAliasDecl *typealias, Type parent,
+                                 const SubstitutionMap &substitutions,
+                                 Type underlying);
+
+  /// \brief Returns the declaration that declares this type.
+  TypeAliasDecl *getDecl() const {
+    // Avoid requiring the definition of TypeAliasDecl.
+    return typealias;
+  }
+
+  /// Retrieve the parent of this type as written, e.g., the part that was
+  /// written before ".", if provided.
+  Type getParent() const {
+    return Bits.BoundNameAliasType.HasParent ? *getTrailingObjects<Type>()
+                                             : Type();
+  }
+
+  /// Retrieve the substitution map applied to the declaration's underlying
+  /// to produce the described type.
+  SubstitutionMap getSubstitutionMap() const;
+
+  /// Get the innermost generic arguments, which correspond to the generic
+  /// arguments that are directly applied to the typealias declaration in
+  /// produced by \c getDecl().
+  ///
+  /// The result can be empty, if the declaration itself is non-generic but
+  /// the parent is generic.
+  SmallVector<Type, 2> getInnermostGenericArgs() const;
+
+  // Support for FoldingSet.
+  void Profile(llvm::FoldingSetNodeID &id) const;
+
+  static void Profile(llvm::FoldingSetNodeID &id, TypeAliasDecl *typealias,
+                      Type parent, const SubstitutionMap &substitutions);
+
+  // Implement isa/cast/dyncast/etc.
+  static bool classof(const TypeBase *T) {
+    return T->getKind() == TypeKind::BoundNameAlias;
+  }
+};
+
 // TODO: As part of AST modernization, replace with a proper
 // 'ParameterTypeElt' or similar, and have FunctionTypes only have a list
 // of 'ParameterTypeElt's. Then, this information can be removed from

include/swift/Serialization/DeclTypeRecordNodes.def

@@ -106,6 +106,7 @@ TYPE(OPENED_EXISTENTIAL)
 TYPE(EXISTENTIAL_METATYPE)
 TYPE(SIL_BLOCK_STORAGE)
 TYPE(SIL_BOX)
+TYPE(BOUND_NAME_ALIAS)
 
 FIRST_DECL(TYPE_ALIAS, 60)
 DECL(GENERIC_TYPE_PARAM)

include/swift/Serialization/ModuleFormat.h

@@ -647,6 +647,14 @@ namespace decls_block {
     TypeIDField  // canonical type (a fallback)
   >;
 
+  using BoundNameAliasTypeLayout = BCRecordLayout<
+    BOUND_NAME_ALIAS_TYPE,
+    DeclIDField, // typealias decl
+    TypeIDField, // parent type
+    TypeIDField  // underlying type
+    // trailing substitutions
+  >;
+
   using GenericTypeParamTypeLayout = BCRecordLayout<
     GENERIC_TYPE_PARAM_TYPE,
     DeclIDField, // generic type parameter decl or depth

lib/AST/ASTContext.cpp

@@ -278,6 +278,7 @@ FOR_KNOWN_FOUNDATION_TYPES(CACHE_FOUNDATION_DECL)
   /// arenas.
   struct Arena {
     llvm::DenseMap<Type, ErrorType *> ErrorTypesWithOriginal;
+    llvm::FoldingSet<BoundNameAliasType> BoundNameAliasTypes;
     llvm::FoldingSet<TupleType> TupleTypes;
     llvm::DenseMap<std::pair<Type,char>, MetatypeType*> MetatypeTypes;
     llvm::DenseMap<std::pair<Type,char>,
@@ -2886,6 +2887,88 @@ StringRef ASTContext::getSwiftName(KnownFoundationEntity kind) {
 // Type manipulation routines.
 //===----------------------------------------------------------------------===//
 
+BoundNameAliasType::BoundNameAliasType(TypeAliasDecl *typealias, Type parent,
+                                       const SubstitutionMap &substitutions,
+                                       Type underlying,
+                                       RecursiveTypeProperties properties)
+    : SugarType(TypeKind::BoundNameAlias, underlying, properties),
+      typealias(typealias) {
+  // Record the parent (or absence of a parent).
+  if (parent) {
+    Bits.BoundNameAliasType.HasParent = true;
+    *getTrailingObjects<Type>() = parent;
+  } else {
+    Bits.BoundNameAliasType.HasParent = false;
+  }
+
+  // Record the substitutions.
+  if (auto genericSig = typealias->getGenericSignature()) {
+    SmallVector<Substitution, 4> flatSubs;
+    genericSig->getSubstitutions(substitutions, flatSubs);
+    Bits.BoundNameAliasType.NumSubstitutions = flatSubs.size();
+    std::copy(flatSubs.begin(), flatSubs.end(),
+              getTrailingObjects<Substitution>());
+  } else {
+    Bits.BoundNameAliasType.NumSubstitutions = 0;
+  }
+}
+
+BoundNameAliasType *BoundNameAliasType::get(
+                                        TypeAliasDecl *typealias,
+                                        Type parent,
+                                        const SubstitutionMap &substitutions,
+                                        Type underlying) {
+  // Compute the recursive properties.
+  auto properties = underlying->getRecursiveProperties();
+  if (parent)
+    properties |= parent->getRecursiveProperties();
+  auto genericSig = typealias->getGenericSignature();
+  if (genericSig) {
+    for (Type gp : genericSig->getGenericParams()) {
+      properties |= gp.subst(substitutions, SubstFlags::UseErrorType)
+                      ->getRecursiveProperties();
+    }
+  }
+
+  // Figure out which arena this type will go into.
+  auto &ctx = underlying->getASTContext();
+  auto arena = getArena(properties);
+
+  // Profile the type.
+  llvm::FoldingSetNodeID id;
+  BoundNameAliasType::Profile(id, typealias, parent, substitutions);
+
+  // Did we already record this type?
+  void *insertPos;
+  auto &types = ctx.Impl.getArena(arena).BoundNameAliasTypes;
+  if (auto result = types.FindNodeOrInsertPos(id, insertPos))
+    return result;
+
+  // Build a new type.
+  unsigned numSubstitutions =
+    genericSig ? genericSig->getSubstitutionListSize() : 0;
+  auto size =
+    totalSizeToAlloc<Type, Substitution>(parent ? 1 : 0, numSubstitutions);
+  auto mem = ctx.Allocate(size, alignof(BoundNameAliasType), arena);
+  auto result = new (mem) BoundNameAliasType(typealias, parent, substitutions,
+                                             underlying, properties);
+  types.InsertNode(result, insertPos);
+  return result;
+}
+
+void BoundNameAliasType::Profile(llvm::FoldingSetNodeID &id) const {
+  Profile(id, getDecl(), getParent(), getSubstitutionMap());
+}
+
+void BoundNameAliasType::Profile(
+                           llvm::FoldingSetNodeID &id,
+                           TypeAliasDecl *typealias,
+                           Type parent, const SubstitutionMap &substitutions) {
+  id.AddPointer(typealias);
+  id.AddPointer(parent.getPointer());
+  substitutions.profile(id);
+}
+
 // Simple accessors.
 Type ErrorType::get(const ASTContext &C) { return C.TheErrorType; }
 

lib/AST/ASTDumper.cpp

@@ -3035,6 +3035,17 @@ namespace {
       OS << ")";
     }
 
+    void visitBoundNameAliasType(BoundNameAliasType *T, StringRef label) {
+      printCommon(label, "bound_name_alias_type");
+      printField("decl", T->getDecl()->printRef());
+      if (T->getParent())
+        printRec("parent", T->getParent());
+
+      for (auto arg : T->getInnermostGenericArgs())
+        printRec(arg);
+      OS << ")";
+    }
+
     void visitParenType(ParenType *T, StringRef label) {
       printCommon(label, "paren_type");
       dumpParameterFlags(T->getParameterFlags());

lib/AST/ASTMangler.cpp

@@ -746,6 +746,13 @@ void ASTMangler::appendType(Type type) {
       // unless the type alias references a builtin type.
       return appendAnyGenericType(decl);
     }
+    case TypeKind::BoundNameAlias: {
+      assert(DWARFMangling && "sugared types are only legal for the debugger");
+      auto boundAliasTy = cast<BoundNameAliasType>(tybase);
+
+      // FIXME: Mangle as a generic type.
+      return appendType(boundAliasTy->getSinglyDesugaredType());
+    }
 
     case TypeKind::Paren:
       return appendSugaredType<ParenType>(type);

lib/AST/ASTPrinter.cpp

@@ -3283,6 +3283,21 @@ class TypePrinter : public TypeVisitor<TypePrinter> {
     printTypeDeclName(T);
   }
 
+  void visitBoundNameAliasType(BoundNameAliasType *T) {
+    if (Options.PrintForSIL || Options.PrintNameAliasUnderlyingType) {
+      visit(T->getSinglyDesugaredType());
+      return;
+    }
+
+    if (auto parent = T->getParent()) {
+      visit(parent);
+      Printer << ".";
+    }
+
+    printTypeDeclName(T);
+    printGenericArgs(T->getInnermostGenericArgs());
+  }
+
   void visitParenType(ParenType *T) {
     Printer << "(";
     printParameterFlags(Printer, Options, T->getParameterFlags());

lib/AST/DiagnosticEngine.cpp

@@ -327,6 +327,9 @@ static void formatSelectionArgument(StringRef ModifierArguments,
 }
 
 static bool isInterestingTypealias(Type type) {
+  // Bound name alias types are always interesting.
+  if (isa<BoundNameAliasType>(type.getPointer())) return true;
+
   auto aliasTy = dyn_cast<NameAliasType>(type.getPointer());
   if (!aliasTy)
     return false;

lib/AST/SubstitutionMap.cpp

@@ -533,3 +533,24 @@ void SubstitutionMap::dump(llvm::raw_ostream &out) const {
 void SubstitutionMap::dump() const {
   return dump(llvm::errs());
 }
+
+void SubstitutionMap::profile(llvm::FoldingSetNodeID &id) const {
+  if (empty() || !genericSig) return;
+
+  // Replacement types.
+  for (Type gp : genericSig->getGenericParams()) {
+    id.AddPointer(gp.subst(*this).getPointer());
+  }
+
+  // Conformance requirements.
+  for (const auto &req : genericSig->getRequirements()) {
+    if (req.getKind() != RequirementKind::Conformance)
+      continue;
+
+    auto conformance =
+      lookupConformance(req.getFirstType()->getCanonicalType(),
+                        req.getSecondType()->castTo<ProtocolType>()->getDecl());
+    id.AddPointer(conformance ? conformance->getOpaqueValue() : nullptr);
+  }
+}
+