Merge pull request #15416 from DougGregor/generic-typealias-sugar

[AST] Preserve type sugar for generic typealiases

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);
+  }
+}
+