[AST] Perf: Improve getDesugaredType() efficiency

include/swift/AST/TypeNodes.def

@@ -99,8 +99,6 @@ ABSTRACT_TYPE(Builtin, Type)
   BUILTIN_TYPE(BuiltinUnsafeValueBuffer, BuiltinType)
   BUILTIN_TYPE(BuiltinVector, BuiltinType)
   TYPE_RANGE(Builtin, BuiltinInteger, BuiltinVector)
-SUGARED_TYPE(NameAlias, Type)
-SUGARED_TYPE(Paren, Type)
 TYPE(Tuple, Type)
 ABSTRACT_TYPE(ReferenceStorage, Type)
   ARTIFICIAL_TYPE(UnownedStorage, ReferenceStorageType)
@@ -141,19 +139,23 @@ ARTIFICIAL_TYPE(SILFunction, Type)
 ARTIFICIAL_TYPE(SILBlockStorage, Type)
 ARTIFICIAL_TYPE(SILBox, Type)
 ARTIFICIAL_TYPE(SILToken, Type)
-ABSTRACT_SUGARED_TYPE(SyntaxSugar, 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)
 UNCHECKED_TYPE(TypeVariable, Type)
-LAST_TYPE(TypeVariable)
+ABSTRACT_SUGARED_TYPE(Sugar, Type)
+  SUGARED_TYPE(Paren, SugarType)
+  SUGARED_TYPE(NameAlias, SugarType)
+  ABSTRACT_SUGARED_TYPE(SyntaxSugar, SugarType)
+    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_RANGE(Sugar, Paren, Dictionary)
+LAST_TYPE(Dictionary) // Sugared types are last to make isa<SugarType>() fast.
 
 #undef TYPE_RANGE
 #undef ABSTRACT_SUGARED_TYPE

include/swift/AST/Types.h

@@ -286,8 +286,12 @@ class alignas(1 << TypeAlignInBits) TypeBase {
     HasOriginalType : 1
   );
 
+  SWIFT_INLINE_BITFIELD(SugarType, TypeBase, 1,
+    HasCachedType : 1
+  );
+
   enum { NumFlagBits = 5 };
-  SWIFT_INLINE_BITFIELD(ParenType, TypeBase, NumFlagBits,
+  SWIFT_INLINE_BITFIELD(ParenType, SugarType, NumFlagBits,
     /// Whether there is an original type.
     Flags : NumFlagBits
   );
@@ -1461,13 +1465,62 @@ class BuiltinFloatType : public BuiltinType {
 };
 DEFINE_EMPTY_CAN_TYPE_WRAPPER(BuiltinFloatType, BuiltinType)
 
+/// An abstract type for all sugared types to make getDesugaredType() fast by
+/// sharing field offsets and logic for the fast path.
+class SugarType : public TypeBase {
+  // The state of this union is known via Bits.SugarType.HasCachedType so that
+  // we can avoid masking the pointer on the fast path.
+  union {
+    TypeBase *UnderlyingType;
+    const ASTContext *Context;
+  };
+
+protected:
+  // Sugar types are never canonical.
+  SugarType(TypeKind K, const ASTContext *ctx,
+            RecursiveTypeProperties properties)
+      : TypeBase(K, nullptr, properties), Context(ctx) {
+    Bits.SugarType.HasCachedType = false;
+  }
+
+  // Sugar types are never canonical.
+  SugarType(TypeKind K, Type type, RecursiveTypeProperties properties)
+      : TypeBase(K, nullptr, properties), UnderlyingType(type.getPointer()) {
+    Bits.SugarType.HasCachedType = true;
+  }
+
+  void setUnderlyingType(Type type) {
+    assert(!Bits.SugarType.HasCachedType && "Cached type already set");
+    Bits.SugarType.HasCachedType = true;
+    UnderlyingType = type.getPointer();
+  }
+
+public:
+  /// Remove one level of top-level sugar from this type.
+  Type getSinglyDesugaredTypeSlow();
+  TypeBase *getSinglyDesugaredType() const {
+    if (LLVM_LIKELY(Bits.SugarType.HasCachedType))
+      return UnderlyingType;
+    auto Ty = const_cast<SugarType*>(this);
+    return Ty->getSinglyDesugaredTypeSlow().getPointer();
+  }
+
+  static bool classof(const TypeBase *T) {
+    if (TypeKind::Last_Type == TypeKind::Last_SugarType)
+      return T->getKind() >= TypeKind::First_SugarType;
+    return T->getKind() >= TypeKind::First_SugarType &&
+           T->getKind() <= TypeKind::Last_SugarType;
+  }
+};
+
 /// NameAliasType - An alias type is a name for another type, just like a
 /// typedef in C.
-class NameAliasType : public TypeBase {
+class NameAliasType : public SugarType {
   friend class TypeAliasDecl;
   // NameAliasType are never canonical.
   NameAliasType(TypeAliasDecl *d) 
-    : TypeBase(TypeKind::NameAlias, nullptr, RecursiveTypeProperties()),
+    : SugarType(TypeKind::NameAlias, (ASTContext*)nullptr,
+                RecursiveTypeProperties()),
       TheDecl(d) {}
   TypeAliasDecl *const TheDecl;
 
@@ -1476,9 +1529,6 @@ class NameAliasType : public TypeBase {
 
   using TypeBase::setRecursiveProperties;
 
-  /// Remove one level of top-level sugar from this type.
-  TypeBase *getSinglyDesugaredType();
-
   // Implement isa/cast/dyncast/etc.
   static bool classof(const TypeBase *T) {
     return T->getKind() == TypeKind::NameAlias;
@@ -1561,23 +1611,18 @@ class ParameterTypeFlags {
 };
 
 /// ParenType - A paren type is a type that's been written in parentheses.
-class ParenType : public TypeBase {
-  Type UnderlyingType;
-
+class ParenType : public SugarType {
   friend class ASTContext;
 
   ParenType(Type UnderlyingType, RecursiveTypeProperties properties,
             ParameterTypeFlags flags);
 
 public:
-  Type getUnderlyingType() const { return UnderlyingType; }
+  Type getUnderlyingType() const { return getSinglyDesugaredType(); }
 
   static ParenType *get(const ASTContext &C, Type underlying,
                         ParameterTypeFlags flags = {});
 
-  /// Remove one level of top-level sugar from this type.
-  TypeBase *getSinglyDesugaredType();
-
   /// Get the parameter flags
   ParameterTypeFlags getParameterFlags() const {
     return ParameterTypeFlags::fromRaw(Bits.ParenType.Flags);
@@ -3951,25 +3996,15 @@ DEFINE_EMPTY_CAN_TYPE_WRAPPER(SILTokenType, Type)
 /// Arrays: [T] -> Array<T>
 /// Optionals: T? -> Optional<T>
 /// Dictionaries: [K : V]  -> Dictionary<K, V>
-class SyntaxSugarType : public TypeBase {
-  llvm::PointerUnion<Type, const ASTContext *> ImplOrContext;
-
-  Type getImplementationTypeSlow();
-
+class SyntaxSugarType : public SugarType {
 protected:
   // Syntax sugar types are never canonical.
   SyntaxSugarType(TypeKind K, const ASTContext &ctx,
-                     RecursiveTypeProperties properties)
-    : TypeBase(K, nullptr, properties), ImplOrContext(&ctx) {}
+                  RecursiveTypeProperties properties)
+    : SugarType(K, &ctx, properties) {}
 
 public:
-  TypeBase *getSinglyDesugaredType();
-
-  Type getImplementationType() {
-    if (ImplOrContext.is<Type>())
-      return ImplOrContext.get<Type>();
-    return getImplementationTypeSlow();
-  }
+  Type getImplementationType() const { return getSinglyDesugaredType(); }
 
   static bool classof(const TypeBase *T) {
     return T->getKind() >= TypeKind::First_SyntaxSugarType &&
@@ -5169,6 +5204,12 @@ constexpr bool TypeBase::isSugaredType<id##Type>() { \
 inline GenericParamKey::GenericParamKey(const GenericTypeParamType *p)
   : Depth(p->getDepth()), Index(p->getIndex()) { }
 
+inline TypeBase *TypeBase::getDesugaredType() {
+  if (!isa<SugarType>(this))
+    return this;
+  return cast<SugarType>(this)->getSinglyDesugaredType()->getDesugaredType();
+}
+
 } // end namespace swift
 
 namespace llvm {

include/swift/Basic/InlineBitfield.h

@@ -38,22 +38,27 @@ namespace swift {
   class T##Bitfield { \
     friend class T; \
     uint64_t __VA_ARGS__; \
+    uint64_t : 64 - (C); /* Better code gen */ \
   } T; \
   LLVM_PACKED_END \
   enum { Num##T##Bits = (C) }; \
   static_assert(sizeof(T##Bitfield) <= 8, "Bitfield overflow")
 
 /// Define an bitfield for type 'T' with parent class 'U' and 'C' bits used.
-#define SWIFT_INLINE_BITFIELD(T, U, C, ...) \
+#define SWIFT_INLINE_BITFIELD_TEMPLATE(T, U, C, HC, ...) \
   LLVM_PACKED_START \
   class T##Bitfield { \
     friend class T; \
     uint64_t : Num##U##Bits, __VA_ARGS__; \
+    uint64_t : 64 - (Num##U##Bits + (HC) + (C)); /* Better code gen */ \
   } T; \
   LLVM_PACKED_END \
   enum { Num##T##Bits = Num##U##Bits + (C) }; \
   static_assert(sizeof(T##Bitfield) <= 8, "Bitfield overflow")
 
+#define SWIFT_INLINE_BITFIELD(T, U, C, ...) \
+  SWIFT_INLINE_BITFIELD_TEMPLATE(T, U, C, 0, __VA_ARGS__)
+
 /// Define a full bitfield for type 'T' that uses all of the remaining bits in
 /// the inline bitfield.
 ///

include/swift/SIL/SILNode.h

@@ -142,7 +142,7 @@ class alignas(8) SILNode {
   );
 
 #define IBWTO_BITFIELD(T, U, C, ...) \
-  SWIFT_INLINE_BITFIELD(T, U, (C), __VA_ARGS__, : 32)
+  SWIFT_INLINE_BITFIELD_TEMPLATE(T, U, (C), 32, __VA_ARGS__)
 #define IBWTO_BITFIELD_EMPTY(T, U) \
   SWIFT_INLINE_BITFIELD_EMPTY(T, U)
 

lib/AST/Type.cpp

@@ -1249,56 +1249,16 @@ TypeBase *TypeBase::reconstituteSugar(bool Recursive) {
     return Func(this).getPointer();
 }
 
-TypeBase *TypeBase::getDesugaredType() {
-  switch (getKind()) {
-#define ALWAYS_CANONICAL_TYPE(id, parent) case TypeKind::id:
-#define UNCHECKED_TYPE(id, parent) case TypeKind::id:
-#define TYPE(id, parent)
-#include "swift/AST/TypeNodes.def"
-  case TypeKind::Error:
-  case TypeKind::Tuple:
-  case TypeKind::Function:
-  case TypeKind::GenericFunction:
-  case TypeKind::SILBlockStorage:
-  case TypeKind::SILBox:
-  case TypeKind::SILFunction:
-  case TypeKind::SILToken:
-  case TypeKind::LValue:
-  case TypeKind::InOut:
-  case TypeKind::ProtocolComposition:
-  case TypeKind::ExistentialMetatype:
-  case TypeKind::Metatype:
-  case TypeKind::BoundGenericClass:
-  case TypeKind::BoundGenericEnum:
-  case TypeKind::BoundGenericStruct:
-  case TypeKind::Enum:
-  case TypeKind::Struct:
-  case TypeKind::Class:
-  case TypeKind::Protocol:
-  case TypeKind::GenericTypeParam:
-  case TypeKind::DependentMember:
-  case TypeKind::UnownedStorage:
-  case TypeKind::UnmanagedStorage:
-  case TypeKind::WeakStorage:
-  case TypeKind::DynamicSelf:
-    // None of these types have sugar at the outer level.
-    return this;
-#define SUGARED_TYPE(ID, PARENT) \
-  case TypeKind::ID: \
-    return cast<ID##Type>(this)->getSinglyDesugaredType()->getDesugaredType();
-#define TYPE(id, parent)
+#define TYPE(Id, Parent)
+#define SUGARED_TYPE(Id, Parent) \
+  static_assert(std::is_base_of<SugarType, Id##Type>::value, "Sugar mismatch");
 #include "swift/AST/TypeNodes.def"
-  }
-
-  llvm_unreachable("Unknown type kind");
-}
 
 ParenType::ParenType(Type baseType, RecursiveTypeProperties properties,
                      ParameterTypeFlags flags)
-  : TypeBase(TypeKind::Paren, nullptr, properties),
-    UnderlyingType(flags.isInOut()
-                     ? InOutType::get(baseType)
-                     : baseType) {
+  : SugarType(TypeKind::Paren,
+              flags.isInOut() ? InOutType::get(baseType) : baseType,
+              properties) {
   Bits.ParenType.Flags = flags.toRaw();
   if (flags.isInOut())
     assert(!baseType->is<InOutType>() && "caller did not pass a base type");
@@ -1307,21 +1267,9 @@ ParenType::ParenType(Type baseType, RecursiveTypeProperties properties,
 }
 
 
-TypeBase *ParenType::getSinglyDesugaredType() {
-  return getUnderlyingType().getPointer();
-}
-
-TypeBase *NameAliasType::getSinglyDesugaredType() {
-  return getDecl()->getUnderlyingTypeLoc().getType().getPointer();
-}
-
-TypeBase *SyntaxSugarType::getSinglyDesugaredType() {
-  return getImplementationType().getPointer();
-}
-
-Type SyntaxSugarType::getImplementationTypeSlow() {
+Type SugarType::getSinglyDesugaredTypeSlow() {
   // Find the generic type that implements this syntactic sugar type.
-  auto &ctx = *ImplOrContext.get<const ASTContext *>();
+  auto &ctx = *Context;
   NominalTypeDecl *implDecl;
 
   // XXX -- If the Decl and Type class hierarchies agreed on spelling, then
@@ -1331,9 +1279,17 @@ Type SyntaxSugarType::getImplementationTypeSlow() {
   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");
+    llvm_unreachable("parenthesis are sugar, but not syntax sugar");
+  case TypeKind::NameAlias: {
+    auto Ty = cast<NameAliasType>(this);
+    auto UTy = Ty->getDecl()->getUnderlyingTypeLoc().getType().getPointer();
+    if (isa<NominalOrBoundGenericNominalType>(UTy)) {
+      Bits.SugarType.HasCachedType = true;
+      UnderlyingType = UTy;
+    }
+    return UTy;
+  }
   case TypeKind::ArraySlice:
     implDecl = ctx.getArrayDecl();
     break;
@@ -1349,17 +1305,18 @@ Type SyntaxSugarType::getImplementationTypeSlow() {
   }
   assert(implDecl && "Type has not been set yet");
 
+  Bits.SugarType.HasCachedType = true;
   if (auto Ty = dyn_cast<UnarySyntaxSugarType>(this)) {
-    ImplOrContext = BoundGenericType::get(implDecl, Type(), Ty->getBaseType());
+    UnderlyingType = BoundGenericType::get(implDecl, Type(), Ty->getBaseType());
   } else if (auto Ty = dyn_cast<DictionaryType>(this)) {
-    ImplOrContext = BoundGenericType::get(implDecl, Type(),
+    UnderlyingType = BoundGenericType::get(implDecl, Type(),
                                       { Ty->getKeyType(), Ty->getValueType() });
   } else {
     llvm_unreachable("Not UnarySyntaxSugarType or DictionaryType?");
   }
 
   // Record the implementation type.
-  return ImplOrContext.get<Type>();
+  return UnderlyingType;
 }
 
 unsigned GenericTypeParamType::getDepth() const {