[refactor] Add the AST source selection component

This commit adds the base AST source selection component to the refactoring
library. AST selection is represented using a tree of SelectedASTNode values.
Each selected node gets its own selection kind, which can actually be None even
in the middle of tree (e.g. statement in a macro whose child is in a macro
argument). The initial version constructs a "raw" selection tree, without
applying filters and canonicalisation operations to the nodes.

Differential Revision: https://reviews.llvm.org/D35012


git-svn-id: https://llvm.org/svn/llvm-project/cfe/trunk@311655 91177308-0d34-0410-b5e6-96231b3b80d8

5 files changed, 235 insertions, 3 deletions

diff --git a/include/clang/AST/LexicallyOrderedRecursiveASTVisitor.h b/include/clang/AST/LexicallyOrderedRecursiveASTVisitor.h
new file mode 100644
index 0000000000..7d2029de02
--- /dev/null
+++ b/include/clang/AST/LexicallyOrderedRecursiveASTVisitor.h
@@ -0,0 +1,137 @@
+//===--- LexicallyOrderedRecursiveASTVisitor.h - ----------------*- C++ -*-===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+//
+//  This file defines the LexicallyOrderedRecursiveASTVisitor interface, which
+//  recursively traverses the entire AST in a lexical order.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_CLANG_AST_LEXICALLY_ORDERED_RECURSIVEASTVISITOR_H
+#define LLVM_CLANG_AST_LEXICALLY_ORDERED_RECURSIVEASTVISITOR_H
+
+#include "clang/AST/RecursiveASTVisitor.h"
+#include "clang/Basic/LLVM.h"
+#include "clang/Basic/SourceManager.h"
+#include "llvm/Support/SaveAndRestore.h"
+
+namespace clang {
+
+/// A RecursiveASTVisitor subclass that guarantees that AST traversal is
+/// performed in a lexical order (i.e. the order in which declarations are
+/// written in the source).
+///
+/// RecursiveASTVisitor doesn't guarantee lexical ordering because there are
+/// some declarations, like Objective-C @implementation declarations
+/// that might be represented in the AST differently to how they were written
+/// in the source.
+/// In particular, Objective-C @implementation declarations may contain
+/// non-Objective-C declarations, like functions:
+///
+///   @implementation MyClass
+///
+///   - (void) method { }
+///   void normalFunction() { }
+///
+///   @end
+///
+/// Clang's AST stores these declarations outside of the @implementation
+/// declaration, so the example above would be represented using the following
+/// AST:
+///   |-ObjCImplementationDecl ... MyClass
+///   | `-ObjCMethodDecl ... method
+///   |    ...
+///   `-FunctionDecl ... normalFunction
+///       ...
+///
+/// This class ensures that these declarations are traversed before the
+/// corresponding TraverseDecl for the @implementation returns. This ensures
+/// that the lexical parent relationship between these declarations and the
+/// @implementation is preserved while traversing the AST. Note that the
+/// current implementation doesn't mix these declarations with the declarations
+/// contained in the @implementation, so the traversal of all of the
+/// declarations in the @implementation still doesn't follow the lexical order.
+template <typename Derived>
+class LexicallyOrderedRecursiveASTVisitor
+    : public RecursiveASTVisitor<Derived> {
+  using BaseType = RecursiveASTVisitor<Derived>;
+
+public:
+  LexicallyOrderedRecursiveASTVisitor(const SourceManager &SM) : SM(SM) {}
+
+  bool TraverseObjCImplementationDecl(ObjCImplementationDecl *D) {
+    // Objective-C @implementation declarations should not trigger early exit
+    // until the additional decls are traversed as their children are not
+    // lexically ordered.
+    bool Result = BaseType::TraverseObjCImplementationDecl(D);
+    return TraverseAdditionalLexicallyNestedDeclarations() ? Result : false;
+  }
+
+  bool TraverseObjCCategoryImplDecl(ObjCCategoryImplDecl *D) {
+    bool Result = BaseType::TraverseObjCCategoryImplDecl(D);
+    return TraverseAdditionalLexicallyNestedDeclarations() ? Result : false;
+  }
+
+  bool TraverseDeclContextHelper(DeclContext *DC) {
+    if (!DC)
+      return true;
+
+    for (auto I = DC->decls_begin(), E = DC->decls_end(); I != E;) {
+      Decl *Child = *I;
+      if (BaseType::canIgnoreChildDeclWhileTraversingDeclContext(Child)) {
+        ++I;
+        continue;
+      }
+      if (!isa<ObjCImplementationDecl>(Child) &&
+          !isa<ObjCCategoryImplDecl>(Child)) {
+        if (!BaseType::getDerived().TraverseDecl(Child))
+          return false;
+        ++I;
+        continue;
+      }
+      // Gather declarations that follow the Objective-C implementation
+      // declarations but are lexically contained in the implementation.
+      LexicallyNestedDeclarations.clear();
+      for (++I; I != E; ++I) {
+        Decl *Sibling = *I;
+        if (!SM.isBeforeInTranslationUnit(Sibling->getLocStart(),
+                                          Child->getLocEnd()))
+          break;
+        if (!BaseType::canIgnoreChildDeclWhileTraversingDeclContext(Sibling))
+          LexicallyNestedDeclarations.push_back(Sibling);
+      }
+      if (!BaseType::getDerived().TraverseDecl(Child))
+        return false;
+    }
+    return true;
+  }
+
+private:
+  bool TraverseAdditionalLexicallyNestedDeclarations() {
+    // FIXME: Ideally the gathered declarations and the declarations in the
+    // @implementation should be mixed and sorted to get a true lexical order,
+    // but right now we only care about getting the correct lexical parent, so
+    // we can traverse the gathered nested declarations after the declarations
+    // in the decl context.
+    assert(!BaseType::getDerived().shouldTraversePostOrder() &&
+           "post-order traversal is not supported for lexically ordered "
+           "recursive ast visitor");
+    for (Decl *D : LexicallyNestedDeclarations) {
+      if (!BaseType::getDerived().TraverseDecl(D))
+        return false;
+    }
+    return true;
+  }
+
+  const SourceManager &SM;
+  llvm::SmallVector<Decl *, 8> LexicallyNestedDeclarations;
+};
+
+} // end namespace clang
+
+#endif // LLVM_CLANG_AST_LEXICALLY_ORDERED_RECURSIVEASTVISITOR_H
diff --git a/include/clang/AST/RecursiveASTVisitor.h b/include/clang/AST/RecursiveASTVisitor.h
index 097ed212cf..fbb3067317 100644
--- a/include/clang/AST/RecursiveASTVisitor.h
+++ b/include/clang/AST/RecursiveASTVisitor.h
@@ -497,6 +497,8 @@ public:
   bool Visit##CLASS##Decl(CLASS##Decl *D) { return true; }
 #include "clang/AST/DeclNodes.inc"
 
+  bool canIgnoreChildDeclWhileTraversingDeclContext(const Decl *Child);
+
 private:
   // These are helper methods used by more than one Traverse* method.
   bool TraverseTemplateParameterListHelper(TemplateParameterList *TPL);
@@ -1345,14 +1347,20 @@ DEF_TRAVERSE_TYPELOC(PipeType, { TRY_TO(TraverseTypeLoc(TL.getValueLoc())); })
 // than those.
 
 template <typename Derived>
+bool RecursiveASTVisitor<Derived>::canIgnoreChildDeclWhileTraversingDeclContext(
+    const Decl *Child) {
+  // BlockDecls and CapturedDecls are traversed through BlockExprs and
+  // CapturedStmts respectively.
+  return isa<BlockDecl>(Child) || isa<CapturedDecl>(Child);
+}
+
+template <typename Derived>
 bool RecursiveASTVisitor<Derived>::TraverseDeclContextHelper(DeclContext *DC) {
   if (!DC)
     return true;
 
   for (auto *Child : DC->decls()) {
-    // BlockDecls and CapturedDecls are traversed through BlockExprs and
-    // CapturedStmts respectively.
-    if (!isa<BlockDecl>(Child) && !isa<CapturedDecl>(Child))
+    if (!canIgnoreChildDeclWhileTraversingDeclContext(Child))
       TRY_TO(TraverseDecl(Child));
   }
 
diff --git a/include/clang/Basic/SourceLocation.h b/include/clang/Basic/SourceLocation.h
index 6afec33445..991cb78695 100644
--- a/include/clang/Basic/SourceLocation.h
+++ b/include/clang/Basic/SourceLocation.h
@@ -172,6 +172,11 @@ public:
     return getFromRawEncoding((unsigned)(uintptr_t)Encoding);
   }
 
+  static bool isPairOfFileLocations(SourceLocation Start, SourceLocation End) {
+    return Start.isValid() && Start.isFileID() && End.isValid() &&
+           End.isFileID();
+  }
+
   void print(raw_ostream &OS, const SourceManager &SM) const;
   std::string printToString(const SourceManager &SM) const;
   void dump(const SourceManager &SM) const;
diff --git a/include/clang/Basic/SourceManager.h b/include/clang/Basic/SourceManager.h
index ed3f8dfa86..40eb1c9f6c 100644
--- a/include/clang/Basic/SourceManager.h
+++ b/include/clang/Basic/SourceManager.h
@@ -1520,6 +1520,14 @@ public:
     return LHSLoaded;
   }
 
+  /// Return true if the Point is within Start and End.
+  bool isPointWithin(SourceLocation Location, SourceLocation Start,
+                     SourceLocation End) const {
+    return Location == Start || Location == End ||
+           (isBeforeInTranslationUnit(Start, Location) &&
+            isBeforeInTranslationUnit(Location, End));
+  }
+
   // Iterators over FileInfos.
   typedef llvm::DenseMap<const FileEntry*, SrcMgr::ContentCache*>
       ::const_iterator fileinfo_iterator;
diff --git a/include/clang/Tooling/Refactoring/ASTSelection.h b/include/clang/Tooling/Refactoring/ASTSelection.h
new file mode 100644
index 0000000000..462e217d73
--- /dev/null
+++ b/include/clang/Tooling/Refactoring/ASTSelection.h
@@ -0,0 +1,74 @@
+//===--- ASTSelection.h - Clang refactoring library -----------------------===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_CLANG_TOOLING_REFACTOR_AST_SELECTION_H
+#define LLVM_CLANG_TOOLING_REFACTOR_AST_SELECTION_H
+
+#include "clang/AST/ASTTypeTraits.h"
+#include "clang/Basic/LLVM.h"
+#include "clang/Basic/SourceLocation.h"
+#include <vector>
+
+namespace clang {
+
+class ASTContext;
+
+namespace tooling {
+
+enum class SourceSelectionKind {
+  /// A node that's not selected.
+  None,
+
+  /// A node that's considered to be selected because the whole selection range
+  /// is inside of its source range.
+  ContainsSelection,
+  /// A node that's considered to be selected because the start of the selection
+  /// range is inside its source range.
+  ContainsSelectionStart,
+  /// A node that's considered to be selected because the end of the selection
+  /// range is inside its source range.
+  ContainsSelectionEnd,
+
+  /// A node that's considered to be selected because the node is entirely in
+  /// the selection range.
+  InsideSelection,
+};
+
+/// Represents a selected AST node.
+///
+/// AST selection is represented using a tree of \c SelectedASTNode. The tree
+/// follows the top-down shape of the actual AST. Each selected node has
+/// a selection kind. The kind might be none as the node itself might not
+/// actually be selected, e.g. a statement in macro whose child is in a macro
+/// argument.
+struct SelectedASTNode {
+  ast_type_traits::DynTypedNode Node;
+  SourceSelectionKind SelectionKind;
+  std::vector<SelectedASTNode> Children;
+
+  SelectedASTNode(const ast_type_traits::DynTypedNode &Node,
+                  SourceSelectionKind SelectionKind)
+      : Node(Node), SelectionKind(SelectionKind) {}
+  SelectedASTNode(SelectedASTNode &&) = default;
+  SelectedASTNode &operator=(SelectedASTNode &&) = default;
+
+  void dump(llvm::raw_ostream &OS = llvm::errs()) const;
+};
+
+/// Traverses the given ASTContext and creates a tree of selected AST nodes.
+///
+/// \returns None if no nodes are selected in the AST, or a selected AST node
+/// that corresponds to the TranslationUnitDecl otherwise.
+Optional<SelectedASTNode> findSelectedASTNodes(const ASTContext &Context,
+                                               SourceRange SelectionRange);
+
+} // end namespace tooling
+} // end namespace clang
+
+#endif // LLVM_CLANG_TOOLING_REFACTOR_AST_SELECTION_H