6 files changed, 280 insertions, 42 deletions

diff --git a/include/clang/AST/ASTFwd.h b/include/clang/AST/ASTFwd.h
new file mode 100644
index 0000000000..27e7faa015
--- /dev/null
+++ b/include/clang/AST/ASTFwd.h
@@ -0,0 +1,27 @@
+//===--- ASTFwd.h ----------------------------------------*- C++ -*-===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===--------------------------------------------------------------===//
+///
+/// \file
+/// \brief Forward declaration of all AST node types.
+///
+//===-------------------------------------------------------------===//
+
+namespace clang {
+
+class Decl;
+#define DECL(DERIVED, BASE) class DERIVED##Decl;
+#include "clang/AST/DeclNodes.inc"
+class Stmt;
+#define STMT(DERIVED, BASE) class DERIVED;
+#include "clang/AST/StmtNodes.inc"
+class Type;
+#define TYPE(DERIVED, BASE) class DERIVED##Type;
+#include "clang/AST/TypeNodes.def"
+
+} // end namespace clang
diff --git a/include/clang/AST/ASTTypeTraits.h b/include/clang/AST/ASTTypeTraits.h
index 1139e0e185..d98ad33a2a 100644
--- a/include/clang/AST/ASTTypeTraits.h
+++ b/include/clang/AST/ASTTypeTraits.h
@@ -7,22 +7,117 @@
 //
 //===----------------------------------------------------------------------===//
 //
-//  Provides a dynamically typed node container that can be used to store
-//  an AST base node at runtime in the same storage in a type safe way.
+//  Provides a dynamic type identifier and a dynamically typed node container
+//  that can be used to store an AST base node at runtime in the same storage in
+//  a type safe way.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_CLANG_AST_AST_TYPE_TRAITS_H
 #define LLVM_CLANG_AST_AST_TYPE_TRAITS_H
 
+#include "clang/AST/ASTFwd.h"
 #include "clang/AST/Decl.h"
 #include "clang/AST/Stmt.h"
 #include "clang/AST/TypeLoc.h"
+#include "clang/Basic/LLVM.h"
 #include "llvm/Support/AlignOf.h"
 
 namespace clang {
 namespace ast_type_traits {
 
+/// \brief Kind identifier.
+///
+/// It can be constructed from any node kind and allows for runtime type
+/// hierarchy checks.
+/// Use getFromNodeKind<T>() to construct them.
+class ASTNodeKind {
+public:
+  /// \brief Empty identifier. It matches nothing.
+  ASTNodeKind() : KindId(NKI_None) {}
+
+  /// \brief Construct an identifier for T.
+  template <class T>
+  static ASTNodeKind getFromNodeKind() {
+    return ASTNodeKind(KindToKindId<T>::Id);
+  }
+
+  /// \brief Returns \c true if \c this and \c Other represent the same kind.
+  bool isSame(ASTNodeKind Other);
+
+  /// \brief Returns \c true if \c this is a base kind of (or same as) \c Other
+  bool isBaseOf(ASTNodeKind Other);
+
+  /// \brief String representation of the kind.
+  StringRef asStringRef() const;
+
+private:
+  /// \brief Kind ids.
+  ///
+  /// Includes all possible base and derived kinds.
+  enum NodeKindId {
+    NKI_None,
+    NKI_NestedNameSpecifier,
+    NKI_NestedNameSpecifierLoc,
+    NKI_QualType,
+    NKI_TypeLoc,
+    NKI_Decl,
+#define DECL(DERIVED, BASE) NKI_##DERIVED##Decl,
+#include "clang/AST/DeclNodes.inc"
+    NKI_Stmt,
+#define STMT(DERIVED, BASE) NKI_##DERIVED,
+#include "clang/AST/StmtNodes.inc"
+    NKI_Type,
+#define TYPE(DERIVED, BASE) NKI_##DERIVED##Type,
+#include "clang/AST/TypeNodes.def"
+    NKI_NumberOfKinds
+  };
+
+  /// \brief Use getFromNodeKind<T>() to construct the kind.
+  ASTNodeKind(NodeKindId KindId) : KindId(KindId) {}
+
+  /// \brief Returns \c true if \c Base is a base kind of (or same as) \c
+  ///   Derived
+  static bool isBaseOf(NodeKindId Base, NodeKindId Derived);
+
+  /// \brief Helper meta-function to convert a kind T to its enum value.
+  ///
+  /// This struct is specialized below for all known kinds.
+  template <class T> struct KindToKindId {
+    static const NodeKindId Id = NKI_None;
+  };
+
+  /// \brief Per kind info.
+  struct KindInfo {
+    /// \brief The id of the parent kind, or None if it has no parent.
+    NodeKindId ParentId;
+    /// \brief Name of the kind.
+    const char *Name;
+  };
+  static const KindInfo AllKindInfo[NKI_NumberOfKinds];
+
+  NodeKindId KindId;
+};
+
+#define KIND_TO_KIND_ID(Class)                                                 \
+  template <> struct ASTNodeKind::KindToKindId<Class> {                        \
+    static const NodeKindId Id = NKI_##Class;                                  \
+  };
+KIND_TO_KIND_ID(NestedNameSpecifier)
+KIND_TO_KIND_ID(NestedNameSpecifierLoc)
+KIND_TO_KIND_ID(QualType)
+KIND_TO_KIND_ID(TypeLoc)
+KIND_TO_KIND_ID(Decl)
+KIND_TO_KIND_ID(Stmt)
+KIND_TO_KIND_ID(Type)
+#define DECL(DERIVED, BASE) KIND_TO_KIND_ID(DERIVED##Decl)
+#include "clang/AST/DeclNodes.inc"
+#define STMT(DERIVED, BASE) KIND_TO_KIND_ID(DERIVED)
+#include "clang/AST/StmtNodes.inc"
+#define TYPE(DERIVED, BASE) KIND_TO_KIND_ID(DERIVED##Type)
+#include "clang/AST/TypeNodes.def"
+#undef KIND_TO_KIND_ID
+
 /// \brief A dynamically typed AST node container.
 ///
 /// Stores an AST node in a type safe way. This allows writing code that
@@ -57,7 +152,7 @@ public:
   /// use the pointer outside the scope of the DynTypedNode.
   template <typename T>
   const T *get() const {
-    return BaseConverter<T>::get(Tag, Storage.buffer);
+    return BaseConverter<T>::get(NodeKind, Storage.buffer);
   }
 
   /// \brief Returns a pointer that identifies the stored AST node.
@@ -90,16 +185,7 @@ private:
   /// \brief Takes care of converting from and to \c T.
   template <typename T, typename EnablerT = void> struct BaseConverter;
 
-  /// \brief Supported base node types.
-  enum NodeTypeTag {
-    NT_Decl,
-    NT_Stmt,
-    NT_NestedNameSpecifier,
-    NT_NestedNameSpecifierLoc,
-    NT_QualType,
-    NT_Type,
-    NT_TypeLoc
-  } Tag;
+  ASTNodeKind NodeKind;
 
   /// \brief Stores the data of the node.
   ///
@@ -107,103 +193,105 @@ private:
   /// guaranteed to be unique pointers pointing to dedicated storage in the
   /// AST. \c QualTypes on the other hand do not have storage or unique
   /// pointers and thus need to be stored by value.
-  llvm::AlignedCharArrayUnion<Decl *, Stmt *, NestedNameSpecifier,
-                              NestedNameSpecifierLoc, QualType, Type,
+  llvm::AlignedCharArrayUnion<Decl *, Stmt *, NestedNameSpecifier *,
+                              NestedNameSpecifierLoc, QualType, Type *,
                               TypeLoc> Storage;
 };
 
 // FIXME: Pull out abstraction for the following.
 template<typename T> struct DynTypedNode::BaseConverter<T,
     typename llvm::enable_if<llvm::is_base_of<Decl, T> >::type> {
-  static const T *get(NodeTypeTag Tag, const char Storage[]) {
-    if (Tag == NT_Decl)
+  static const T *get(ASTNodeKind NodeKind, const char Storage[]) {
+    if (ASTNodeKind::getFromNodeKind<Decl>().isBaseOf(NodeKind))
       return dyn_cast<T>(*reinterpret_cast<Decl*const*>(Storage));
     return NULL;
   }
   static DynTypedNode create(const Decl &Node) {
     DynTypedNode Result;
-    Result.Tag = NT_Decl;
+    Result.NodeKind = ASTNodeKind::getFromNodeKind<T>();
     new (Result.Storage.buffer) const Decl*(&Node);
     return Result;
   }
 };
 template<typename T> struct DynTypedNode::BaseConverter<T,
     typename llvm::enable_if<llvm::is_base_of<Stmt, T> >::type> {
-  static const T *get(NodeTypeTag Tag, const char Storage[]) {
-    if (Tag == NT_Stmt)
+  static const T *get(ASTNodeKind NodeKind, const char Storage[]) {
+    if (ASTNodeKind::getFromNodeKind<Stmt>().isBaseOf(NodeKind))
       return dyn_cast<T>(*reinterpret_cast<Stmt*const*>(Storage));
     return NULL;
   }
   static DynTypedNode create(const Stmt &Node) {
     DynTypedNode Result;
-    Result.Tag = NT_Stmt;
+    Result.NodeKind = ASTNodeKind::getFromNodeKind<T>();
     new (Result.Storage.buffer) const Stmt*(&Node);
     return Result;
   }
 };
 template<typename T> struct DynTypedNode::BaseConverter<T,
     typename llvm::enable_if<llvm::is_base_of<Type, T> >::type> {
-  static const T *get(NodeTypeTag Tag, const char Storage[]) {
-    if (Tag == NT_Type)
+  static const T *get(ASTNodeKind NodeKind, const char Storage[]) {
+    if (ASTNodeKind::getFromNodeKind<Type>().isBaseOf(NodeKind))
       return dyn_cast<T>(*reinterpret_cast<Type*const*>(Storage));
     return NULL;
   }
   static DynTypedNode create(const Type &Node) {
     DynTypedNode Result;
-    Result.Tag = NT_Type;
+    Result.NodeKind = ASTNodeKind::getFromNodeKind<T>();
     new (Result.Storage.buffer) const Type*(&Node);
     return Result;
   }
 };
 template<> struct DynTypedNode::BaseConverter<NestedNameSpecifier, void> {
-  static const NestedNameSpecifier *get(NodeTypeTag Tag, const char Storage[]) {
-    if (Tag == NT_NestedNameSpecifier)
+  static const NestedNameSpecifier *get(ASTNodeKind NodeKind,
+                                        const char Storage[]) {
+    if (ASTNodeKind::getFromNodeKind<NestedNameSpecifier>().isBaseOf(NodeKind))
       return *reinterpret_cast<NestedNameSpecifier*const*>(Storage);
     return NULL;
   }
   static DynTypedNode create(const NestedNameSpecifier &Node) {
     DynTypedNode Result;
-    Result.Tag = NT_NestedNameSpecifier;
+    Result.NodeKind = ASTNodeKind::getFromNodeKind<NestedNameSpecifier>();
     new (Result.Storage.buffer) const NestedNameSpecifier*(&Node);
     return Result;
   }
 };
 template<> struct DynTypedNode::BaseConverter<NestedNameSpecifierLoc, void> {
-  static const NestedNameSpecifierLoc *get(NodeTypeTag Tag,
+  static const NestedNameSpecifierLoc *get(ASTNodeKind NodeKind,
                                            const char Storage[]) {
-    if (Tag == NT_NestedNameSpecifierLoc)
+    if (ASTNodeKind::getFromNodeKind<NestedNameSpecifierLoc>().isBaseOf(
+            NodeKind))
       return reinterpret_cast<const NestedNameSpecifierLoc*>(Storage);
     return NULL;
   }
   static DynTypedNode create(const NestedNameSpecifierLoc &Node) {
     DynTypedNode Result;
-    Result.Tag = NT_NestedNameSpecifierLoc;
+    Result.NodeKind = ASTNodeKind::getFromNodeKind<NestedNameSpecifierLoc>();
     new (Result.Storage.buffer) NestedNameSpecifierLoc(Node);
     return Result;
   }
 };
 template<> struct DynTypedNode::BaseConverter<QualType, void> {
-  static const QualType *get(NodeTypeTag Tag, const char Storage[]) {
-    if (Tag == NT_QualType)
+  static const QualType *get(ASTNodeKind NodeKind, const char Storage[]) {
+    if (ASTNodeKind::getFromNodeKind<QualType>().isBaseOf(NodeKind))
       return reinterpret_cast<const QualType*>(Storage);
     return NULL;
   }
   static DynTypedNode create(const QualType &Node) {
     DynTypedNode Result;
-    Result.Tag = NT_QualType;
+    Result.NodeKind = ASTNodeKind::getFromNodeKind<QualType>();
     new (Result.Storage.buffer) QualType(Node);
     return Result;
   }
 };
 template<> struct DynTypedNode::BaseConverter<TypeLoc, void> {
-  static const TypeLoc *get(NodeTypeTag Tag, const char Storage[]) {
-    if (Tag == NT_TypeLoc)
+  static const TypeLoc *get(ASTNodeKind NodeKind, const char Storage[]) {
+    if (ASTNodeKind::getFromNodeKind<TypeLoc>().isBaseOf(NodeKind))
       return reinterpret_cast<const TypeLoc*>(Storage);
     return NULL;
   }
   static DynTypedNode create(const TypeLoc &Node) {
     DynTypedNode Result;
-    Result.Tag = NT_TypeLoc;
+    Result.NodeKind = ASTNodeKind::getFromNodeKind<TypeLoc>();
     new (Result.Storage.buffer) TypeLoc(Node);
     return Result;
   }
@@ -213,15 +301,18 @@ template<> struct DynTypedNode::BaseConverter<TypeLoc, void> {
 // AST node that is not supported, but prevents misuse - a user cannot create
 // a DynTypedNode from arbitrary types.
 template <typename T, typename EnablerT> struct DynTypedNode::BaseConverter {
-  static const T *get(NodeTypeTag Tag, const char Storage[]) { return NULL; }
+  static const T *get(ASTNodeKind NodeKind, const char Storage[]) {
+    return NULL;
+  }
 };
 
 inline const void *DynTypedNode::getMemoizationData() const {
-  switch (Tag) {
-    case NT_Decl: return BaseConverter<Decl>::get(Tag, Storage.buffer);
-    case NT_Stmt: return BaseConverter<Stmt>::get(Tag, Storage.buffer);
-    default: return NULL;
-  };
+  if (ASTNodeKind::getFromNodeKind<Decl>().isBaseOf(NodeKind)) {
+    return BaseConverter<Decl>::get(NodeKind, Storage.buffer);
+  } else if (ASTNodeKind::getFromNodeKind<Stmt>().isBaseOf(NodeKind)) {
+    return BaseConverter<Stmt>::get(NodeKind, Storage.buffer);
+  }
+  return NULL;
 }
 
 } // end namespace ast_type_traits
diff --git a/lib/AST/ASTTypeTraits.cpp b/lib/AST/ASTTypeTraits.cpp
new file mode 100644
index 0000000000..40e669d7ad
--- /dev/null
+++ b/lib/AST/ASTTypeTraits.cpp
@@ -0,0 +1,56 @@
+//===--- ASTTypeTraits.cpp --------------------------------------*- C++ -*-===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+//
+//  Provides a dynamic type identifier and a dynamically typed node container
+//  that can be used to store an AST base node at runtime in the same storage in
+//  a type safe way.
+//
+//===----------------------------------------------------------------------===//
+
+#include "clang/AST/ASTTypeTraits.h"
+
+namespace clang {
+namespace ast_type_traits {
+
+const ASTNodeKind::KindInfo ASTNodeKind::AllKindInfo[] = {
+  { NKI_None, "<None>" },
+  { NKI_None, "NestedNameSpecifier" },
+  { NKI_None, "NestedNameSpecifierLoc" },
+  { NKI_None, "QualType" },
+  { NKI_None, "TypeLoc" },
+  { NKI_None, "Decl" },
+#define DECL(DERIVED, BASE) { NKI_##BASE, #DERIVED "Decl" },
+#include "clang/AST/DeclNodes.inc"
+  { NKI_None, "Stmt" },
+#define STMT(DERIVED, BASE) { NKI_##BASE, #DERIVED },
+#include "clang/AST/StmtNodes.inc"
+  { NKI_None, "Type" },
+#define TYPE(DERIVED, BASE) { NKI_##BASE, #DERIVED "Type" },
+#include "clang/AST/TypeNodes.def"
+};
+
+bool ASTNodeKind::isBaseOf(ASTNodeKind Other) {
+  return isBaseOf(KindId, Other.KindId);
+}
+
+bool ASTNodeKind::isSame(ASTNodeKind Other) {
+  return KindId != NKI_None && KindId == Other.KindId;
+}
+
+bool ASTNodeKind::isBaseOf(NodeKindId Base, NodeKindId Derived) {
+  if (Base == NKI_None || Derived == NKI_None) return false;
+  while (Derived != Base && Derived != NKI_None)
+    Derived = AllKindInfo[Derived].ParentId;
+  return Derived == Base;
+}
+
+StringRef ASTNodeKind::asStringRef() const { return AllKindInfo[KindId].Name; }
+
+} // end namespace ast_type_traits
+} // end namespace clang
diff --git a/lib/AST/CMakeLists.txt b/lib/AST/CMakeLists.txt
index e804fe7205..b508144684 100644
--- a/lib/AST/CMakeLists.txt
+++ b/lib/AST/CMakeLists.txt
@@ -7,6 +7,7 @@ add_clang_library(clangAST
   ASTDiagnostic.cpp
   ASTDumper.cpp
   ASTImporter.cpp
+  ASTTypeTraits.cpp
   AttrImpl.cpp
   CXXInheritance.cpp
   Comment.cpp
diff --git a/unittests/AST/ASTTypeTraitsTest.cpp b/unittests/AST/ASTTypeTraitsTest.cpp
new file mode 100644
index 0000000000..5e3abe25d1
--- /dev/null
+++ b/unittests/AST/ASTTypeTraitsTest.cpp
@@ -0,0 +1,62 @@
+//===- unittest/AST/ASTTypeTraits.cpp - AST type traits unit tests ------===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===--------------------------------------------------------------------===//
+
+
+#include "clang/AST/ASTTypeTraits.h"
+#include "gtest/gtest.h"
+
+namespace clang {
+namespace ast_type_traits {
+
+TEST(ASTNodeKind, NoKind) {
+  EXPECT_FALSE(ASTNodeKind().isBaseOf(ASTNodeKind()));
+  EXPECT_FALSE(ASTNodeKind().isSame(ASTNodeKind()));
+}
+
+template <typename T> static ASTNodeKind DNT() {
+  return ASTNodeKind::getFromNodeKind<T>();
+}
+
+TEST(ASTNodeKind, Bases) {
+  EXPECT_TRUE(DNT<Decl>().isBaseOf(DNT<VarDecl>()));
+  EXPECT_FALSE(DNT<Decl>().isSame(DNT<VarDecl>()));
+  EXPECT_FALSE(DNT<VarDecl>().isBaseOf(DNT<Decl>()));
+
+  EXPECT_TRUE(DNT<Decl>().isSame(DNT<Decl>()));
+}
+
+TEST(ASTNodeKind, SameBase) {
+  EXPECT_TRUE(DNT<Expr>().isBaseOf(DNT<CallExpr>()));
+  EXPECT_TRUE(DNT<Expr>().isBaseOf(DNT<BinaryOperator>()));
+  EXPECT_FALSE(DNT<CallExpr>().isBaseOf(DNT<BinaryOperator>()));
+  EXPECT_FALSE(DNT<BinaryOperator>().isBaseOf(DNT<CallExpr>()));
+}
+
+TEST(ASTNodeKind, DiffBase) {
+  EXPECT_FALSE(DNT<Expr>().isBaseOf(DNT<ArrayType>()));
+  EXPECT_FALSE(DNT<QualType>().isBaseOf(DNT<FunctionDecl>()));
+  EXPECT_FALSE(DNT<Type>().isSame(DNT<QualType>()));
+}
+
+struct Foo {};
+
+TEST(ASTNodeKind, UnknownKind) {
+  // We can construct one, but it is nowhere in the hierarchy.
+  EXPECT_FALSE(DNT<Foo>().isSame(DNT<Foo>()));
+}
+
+TEST(ASTNodeKind, Name) {
+  EXPECT_EQ("Decl", DNT<Decl>().asStringRef());
+  EXPECT_EQ("CallExpr", DNT<CallExpr>().asStringRef());
+  EXPECT_EQ("ConstantArrayType", DNT<ConstantArrayType>().asStringRef());
+  EXPECT_EQ("<None>", ASTNodeKind().asStringRef());
+}
+
+}  // namespace ast_type_traits
+}  // namespace clang
diff --git a/unittests/AST/CMakeLists.txt b/unittests/AST/CMakeLists.txt
index 3ef2a5e153..c414ae3b23 100644
--- a/unittests/AST/CMakeLists.txt
+++ b/unittests/AST/CMakeLists.txt
@@ -1,5 +1,6 @@
 add_clang_unittest(ASTTests
   ASTContextParentMapTest.cpp
+  ASTTypeTraitsTest.cpp
   CommentLexer.cpp
   CommentParser.cpp
   DeclPrinterTest.cpp