[AST] Perf: Unite DictionaryType with SyntaxSugarType

1) Move existing SyntaxSugarTypes under a new subclass called UnarySyntaxSugarType.
2) Make DictionaryType subclass SyntaxSugarType.

This helps improve getDesugaredType() performance by ensuring that
ImplOrContext is stored at the same field offset in memory.

This also de-boilerplates some AST walking.

Author: davezarzycki

include/swift/AST/TypeNodes.def

@@ -142,11 +142,13 @@ ARTIFICIAL_TYPE(SILBlockStorage, Type)
 ARTIFICIAL_TYPE(SILBox, Type)
 ARTIFICIAL_TYPE(SILToken, Type)
 ABSTRACT_SUGARED_TYPE(SyntaxSugar, Type)
-  SUGARED_TYPE(ArraySlice, SyntaxSugarType)
-  SUGARED_TYPE(Optional, SyntaxSugarType)
-  SUGARED_TYPE(ImplicitlyUnwrappedOptional, SyntaxSugarType)
-  TYPE_RANGE(SyntaxSugar, ArraySlice, ImplicitlyUnwrappedOptional)
-SUGARED_TYPE(Dictionary, Type)
+  ABSTRACT_SUGARED_TYPE(UnarySyntaxSugar, SyntaxSugarType)
+    SUGARED_TYPE(ArraySlice, UnarySyntaxSugarType)
+    SUGARED_TYPE(Optional, UnarySyntaxSugarType)
+    SUGARED_TYPE(ImplicitlyUnwrappedOptional, UnarySyntaxSugarType)
+    TYPE_RANGE(UnarySyntaxSugar, ArraySlice, ImplicitlyUnwrappedOptional)
+  SUGARED_TYPE(Dictionary, SyntaxSugarType)
+  TYPE_RANGE(SyntaxSugar, ArraySlice, Dictionary)
 TYPE(ProtocolComposition, Type)
 TYPE(LValue, Type)
 TYPE(InOut, Type)

include/swift/AST/Types.h

@@ -3943,40 +3943,69 @@ class SILTokenType final : public TypeBase {
 };
 DEFINE_EMPTY_CAN_TYPE_WRAPPER(SILTokenType, Type)
 
-/// A type with a special syntax that is always sugar for a library type.
+/// A type with a special syntax that is always sugar for a library type. The
+/// library type may have multiple base types. For unary syntax sugar, see
+/// UnarySyntaxSugarType.
 ///
-/// The prime examples are arrays ([T] -> Array<T>) and
-/// optionals (T? -> Optional<T>).
+/// The prime examples are:
+/// Arrays: [T] -> Array<T>
+/// Optionals: T? -> Optional<T>
+/// Dictionaries: [K : V]  -> Dictionary<K, V>
 class SyntaxSugarType : public TypeBase {
-  Type Base;
   llvm::PointerUnion<Type, const ASTContext *> ImplOrContext;
 
+  Type getImplementationTypeSlow();
+
 protected:
   // Syntax sugar types are never canonical.
-  SyntaxSugarType(TypeKind K, const ASTContext &ctx, Type base,
-                  RecursiveTypeProperties properties)
-    : TypeBase(K, nullptr, properties), Base(base), ImplOrContext(&ctx) {}
+  SyntaxSugarType(TypeKind K, const ASTContext &ctx,
+                     RecursiveTypeProperties properties)
+    : TypeBase(K, nullptr, properties), ImplOrContext(&ctx) {}
 
 public:
-  Type getBaseType() const {
-    return Base;
-  }
-
   TypeBase *getSinglyDesugaredType();
 
-  Type getImplementationType();
+  Type getImplementationType() {
+    if (ImplOrContext.is<Type>())
+      return ImplOrContext.get<Type>();
+    return getImplementationTypeSlow();
+  }
 
   static bool classof(const TypeBase *T) {
     return T->getKind() >= TypeKind::First_SyntaxSugarType &&
            T->getKind() <= TypeKind::Last_SyntaxSugarType;
   }
 };
+
+/// A type with a special syntax that is always sugar for a library type that
+/// wraps a single other type.
+///
+/// The prime examples are arrays ([T] -> Array<T>) and
+/// optionals (T? -> Optional<T>).
+class UnarySyntaxSugarType : public SyntaxSugarType {
+  Type Base;
+
+protected:
+  UnarySyntaxSugarType(TypeKind K, const ASTContext &ctx, Type base,
+                       RecursiveTypeProperties properties)
+    : SyntaxSugarType(K, ctx, properties), Base(base) {}
+
+public:
+  Type getBaseType() const {
+    return Base;
+  }
+
+  static bool classof(const TypeBase *T) {
+    return T->getKind() >= TypeKind::First_UnarySyntaxSugarType &&
+           T->getKind() <= TypeKind::Last_UnarySyntaxSugarType;
+  }
+};
 
 /// The type [T], which is always sugar for a library type.
-class ArraySliceType : public SyntaxSugarType {
+class ArraySliceType : public UnarySyntaxSugarType {
   ArraySliceType(const ASTContext &ctx, Type base,
                  RecursiveTypeProperties properties)
-    : SyntaxSugarType(TypeKind::ArraySlice, ctx, base, properties) {}
+    : UnarySyntaxSugarType(TypeKind::ArraySlice, ctx, base, properties) {}
 
 public:
   /// Return a uniqued array slice type with the specified base type.
@@ -3988,10 +4017,10 @@ class ArraySliceType : public SyntaxSugarType {
 };
 
 /// The type T?, which is always sugar for a library type.
-class OptionalType : public SyntaxSugarType {
+class OptionalType : public UnarySyntaxSugarType {
   OptionalType(const ASTContext &ctx,Type base,
                RecursiveTypeProperties properties)
-    : SyntaxSugarType(TypeKind::Optional, ctx, base, properties) {}
+    : UnarySyntaxSugarType(TypeKind::Optional, ctx, base, properties) {}
 
 public:
   /// Return a uniqued optional type with the specified base type.
@@ -4014,10 +4043,11 @@ class OptionalType : public SyntaxSugarType {
 };
 
 /// The type T!, which is always sugar for a library type.
-class ImplicitlyUnwrappedOptionalType : public SyntaxSugarType {
+class ImplicitlyUnwrappedOptionalType : public UnarySyntaxSugarType {
   ImplicitlyUnwrappedOptionalType(const ASTContext &ctx, Type base,
                         RecursiveTypeProperties properties)
-    : SyntaxSugarType(TypeKind::ImplicitlyUnwrappedOptional, ctx, base, properties) {}
+    : UnarySyntaxSugarType(TypeKind::ImplicitlyUnwrappedOptional, ctx, base,
+                           properties) {}
 
 public:
   /// Return a uniqued optional type with the specified base type.
@@ -4035,17 +4065,16 @@ class ImplicitlyUnwrappedOptionalType : public SyntaxSugarType {
 /// \code
 /// var dict: [String : Int] = ["hello" : 0, "world" : 1]
 /// \endcode
-class DictionaryType : public TypeBase {
+class DictionaryType : public SyntaxSugarType {
   Type Key;
   Type Value;
-  llvm::PointerUnion<Type, const ASTContext *> ImplOrContext;
 
 protected:
   // Syntax sugar types are never canonical.
   DictionaryType(const ASTContext &ctx, Type key, Type value,
                  RecursiveTypeProperties properties)
-    : TypeBase(TypeKind::Dictionary, nullptr, properties), 
-      Key(key), Value(value), ImplOrContext(&ctx) {}
+    : SyntaxSugarType(TypeKind::Dictionary, ctx, properties), 
+      Key(key), Value(value) {}
 
 public:
   /// Return a uniqued dictionary type with the specified key and value types.
@@ -4054,10 +4083,6 @@ class DictionaryType : public TypeBase {
   Type getKeyType() const { return Key; }
   Type getValueType() const { return Value; }
 
-  TypeBase *getSinglyDesugaredType();
-
-  Type getImplementationType();
-
   static bool classof(const TypeBase *T) {
     return T->getKind() == TypeKind::Dictionary;
   }

lib/AST/ASTMangler.cpp

@@ -734,9 +734,8 @@ void ASTMangler::appendType(Type type) {
       return appendSugaredType<ParenType>(type);
     case TypeKind::ArraySlice: /* fallthrough */
     case TypeKind::Optional:
-      return appendSugaredType<SyntaxSugarType>(type);
     case TypeKind::Dictionary:
-      return appendSugaredType<DictionaryType>(type);
+      return appendSugaredType<SyntaxSugarType>(type);
 
     case TypeKind::ImplicitlyUnwrappedOptional: {
       assert(DWARFMangling && "sugared types are only legal for the debugger");

lib/AST/ASTPrinter.cpp

@@ -3381,8 +3381,6 @@ class TypePrinter : public TypeVisitor<TypePrinter> {
 
     if (auto sugarType = dyn_cast<SyntaxSugarType>(T.getPointer()))
       T = sugarType->getImplementationType();
-    else if (auto dictType = dyn_cast<DictionaryType>(T.getPointer()))
-      T = dictType->getImplementationType();
 
     TypePrinter(Printer, innerOptions).visit(T);
   }

lib/AST/Type.cpp

@@ -1319,47 +1319,46 @@ TypeBase *SyntaxSugarType::getSinglyDesugaredType() {
   return getImplementationType().getPointer();
 }
 
-Type SyntaxSugarType::getImplementationType() {
-  if (ImplOrContext.is<Type>())
-    return ImplOrContext.get<Type>();
-
+Type SyntaxSugarType::getImplementationTypeSlow() {
   // Find the generic type that implements this syntactic sugar type.
   auto &ctx = *ImplOrContext.get<const ASTContext *>();
   NominalTypeDecl *implDecl;
 
-  if (isa<ArraySliceType>(this)) {
+  // XXX -- If the Decl and Type class hierarchies agreed on spelling, then
+  // we could handle the entire switch statement via macros.
+  switch (getKind()) {
+#define TYPE(Id, Parent) \
+  case TypeKind::Id: llvm_unreachable("non-sugared type?");
+#define SUGARED_TYPE(Id, Parent)
+#include "swift/AST/TypeNodes.def"
+  case TypeKind::NameAlias:
+  case TypeKind::Paren:
+    llvm_unreachable("typealiases and parens are sugar, but not syntax sugar");
+  case TypeKind::ArraySlice:
     implDecl = ctx.getArrayDecl();
-    assert(implDecl && "Array type has not been set yet");
-  } else if (isa<OptionalType>(this)) {
+    break;
+  case TypeKind::Optional:
     implDecl = ctx.getOptionalDecl();
-    assert(implDecl && "Optional type has not been set yet");
-  } else if (isa<ImplicitlyUnwrappedOptionalType>(this)) {
+    break;
+  case TypeKind::ImplicitlyUnwrappedOptional:
     implDecl = ctx.getImplicitlyUnwrappedOptionalDecl();
-    assert(implDecl && "Optional type has not been set yet");
-  } else {
-    llvm_unreachable("Unhandled syntax sugar type");
+    break;
+  case TypeKind::Dictionary:
+    implDecl = ctx.getDictionaryDecl();
+    break;
   }
+  assert(implDecl && "Type has not been set yet");
 
-  // Record the implementation type.
-  ImplOrContext = BoundGenericType::get(implDecl, Type(), Base);
-  return ImplOrContext.get<Type>();
-}
-
-TypeBase *DictionaryType::getSinglyDesugaredType() {
-  return getImplementationType().getPointer();
-}
-
-Type DictionaryType::getImplementationType() {
-  if (ImplOrContext.is<Type>())
-    return ImplOrContext.get<Type>();
-
-  // Find the generic type that implements this syntactic sugar type.
-  auto &ctx = *ImplOrContext.get<const ASTContext *>();
-  NominalTypeDecl *implDecl = ctx.getDictionaryDecl();
-  assert(implDecl && "Dictionary type has not been set yet");
+  if (auto Ty = dyn_cast<UnarySyntaxSugarType>(this)) {
+    ImplOrContext = BoundGenericType::get(implDecl, Type(), Ty->getBaseType());
+  } else if (auto Ty = dyn_cast<DictionaryType>(this)) {
+    ImplOrContext = BoundGenericType::get(implDecl, Type(),
+                                      { Ty->getKeyType(), Ty->getValueType() });
+  } else {
+    llvm_unreachable("Not UnarySyntaxSugarType or DictionaryType?");
+  }
 
   // Record the implementation type.
-  ImplOrContext = BoundGenericType::get(implDecl, Type(), { Key, Value });
   return ImplOrContext.get<Type>();
 }
 

lib/AST/TypeWalker.cpp

@@ -119,7 +119,7 @@ class Traversal : public TypeVisitor<Traversal, bool>
     return false;
   }
 
-  bool visitSyntaxSugarType(SyntaxSugarType *ty) {
+  bool visitUnarySyntaxSugarType(UnarySyntaxSugarType *ty) {
     return doIt(ty->getBaseType());
   }
 

lib/IRGen/IRGenDebugInfo.cpp

@@ -1298,6 +1298,7 @@ class IRGenDebugInfoImpl : public IRGenDebugInfo {
     }
 
     // SyntaxSugarType derivations.
+    case TypeKind::Dictionary:
     case TypeKind::ArraySlice:
     case TypeKind::Optional:
     case TypeKind::ImplicitlyUnwrappedOptional: {
@@ -1306,12 +1307,6 @@ class IRGenDebugInfoImpl : public IRGenDebugInfo {
       return getOrCreateDesugaredType(CanTy, DbgTy);
     }
 
-    case TypeKind::Dictionary: {
-      auto *DictionaryTy = cast<DictionaryType>(BaseTy);
-      auto *CanTy = DictionaryTy->getDesugaredType();
-      return getOrCreateDesugaredType(CanTy, DbgTy);
-    }
-
     case TypeKind::GenericTypeParam: {
       auto *ParamTy = cast<GenericTypeParamType>(BaseTy);
       // FIXME: Provide a more meaningful debug type.

lib/PrintAsObjC/PrintAsObjC.cpp

@@ -1824,11 +1824,6 @@ class ObjCPrinter : private DeclVisitor<ObjCPrinter>,
     visitPart(SST->getSinglyDesugaredType(), optionalKind);
   }
 
-  void visitDictionaryType(DictionaryType *DT, 
-                           Optional<OptionalTypeKind> optionalKind) {
-    visitPart(DT->getSinglyDesugaredType(), optionalKind);
-  }
-
   void visitDynamicSelfType(DynamicSelfType *DST, 
                             Optional<OptionalTypeKind> optionalKind) {
     printNullability(optionalKind, NullabilityPrintKind::ContextSensitive);
@@ -1961,10 +1956,6 @@ class ReferencedTypeFinder : public TypeVisitor<ReferencedTypeFinder> {
     visit(sugar->getSinglyDesugaredType());
   }
 
-  void visitDictionaryType(DictionaryType *DT) {
-    visit(DT->getSinglyDesugaredType());
-  }
-
   void visitProtocolCompositionType(ProtocolCompositionType *composition) {
     auto layout = composition->getExistentialLayout();
     if (layout.superclass)