From 013c07f697886649b068cd97127e528b4fe7c03e Mon Sep 17 00:00:00 2001
From: Ilya Biryukov <ibiryukov@google.com>
Date: Tue, 14 Jan 2020 12:34:00 +0100
Subject: [PATCH] [Syntax] Unset IsOriginal flag on nodes removed from the tree

And add a corresponding test.
Only nodes inside the TranslationUnit subtree can be marked as original,
computeReplacements() relies on this.
---
 clang/lib/Tooling/Syntax/Tree.cpp           | 27 +++++++++++++++------
 clang/unittests/Tooling/Syntax/TreeTest.cpp |  3 +++
 2 files changed, 22 insertions(+), 8 deletions(-)

diff --git a/clang/lib/Tooling/Syntax/Tree.cpp b/clang/lib/Tooling/Syntax/Tree.cpp
index 5282857df5a9e8..9f028c0be3b4cb 100644
--- a/clang/lib/Tooling/Syntax/Tree.cpp
+++ b/clang/lib/Tooling/Syntax/Tree.cpp
@@ -15,6 +15,23 @@
 
 using namespace clang;
 
+namespace {
+static void traverse(const syntax::Node *N,
+                     llvm::function_ref<void(const syntax::Node *)> Visit) {
+  if (auto *T = dyn_cast<syntax::Tree>(N)) {
+    for (auto *C = T->firstChild(); C; C = C->nextSibling())
+      traverse(C, Visit);
+  }
+  Visit(N);
+}
+static void traverse(syntax::Node *N,
+                     llvm::function_ref<void(syntax::Node *)> Visit) {
+  traverse(static_cast<const syntax::Node *>(N), [&](const syntax::Node *N) {
+    Visit(const_cast<syntax::Node *>(N));
+  });
+};
+} // namespace
+
 syntax::Arena::Arena(SourceManager &SourceMgr, const LangOptions &LangOpts,
                      TokenBuffer Tokens)
     : SourceMgr(SourceMgr), LangOpts(LangOpts), Tokens(std::move(Tokens)) {}
@@ -80,6 +97,8 @@ void syntax::Tree::replaceChildRangeLowLevel(Node *BeforeBegin, Node *End,
     N->Role = static_cast<unsigned>(NodeRole::Detached);
     N->Parent = nullptr;
     N->NextSibling = nullptr;
+    if (N->Original)
+      traverse(N, [&](Node *C) { C->Original = false; });
 
     N = Next;
   }
@@ -105,14 +124,6 @@ void syntax::Tree::replaceChildRangeLowLevel(Node *BeforeBegin, Node *End,
 }
 
 namespace {
-static void traverse(const syntax::Node *N,
-                     llvm::function_ref<void(const syntax::Node *)> Visit) {
-  if (auto *T = dyn_cast<syntax::Tree>(N)) {
-    for (auto *C = T->firstChild(); C; C = C->nextSibling())
-      traverse(C, Visit);
-  }
-  Visit(N);
-}
 static void dumpTokens(llvm::raw_ostream &OS, ArrayRef<syntax::Token> Tokens,
                        const SourceManager &SM) {
   assert(!Tokens.empty());
diff --git a/clang/unittests/Tooling/Syntax/TreeTest.cpp b/clang/unittests/Tooling/Syntax/TreeTest.cpp
index c457c78881241c..42d77d89353642 100644
--- a/clang/unittests/Tooling/Syntax/TreeTest.cpp
+++ b/clang/unittests/Tooling/Syntax/TreeTest.cpp
@@ -895,6 +895,9 @@ TEST_F(SyntaxTreeTest, Mutations) {
     auto *S = cast<syntax::Statement>(nodeByRange(Input.range(), TU));
     ASSERT_TRUE(S->canModify()) << "cannot remove a statement";
     syntax::removeStatement(*Arena, S);
+    EXPECT_TRUE(S->isDetached());
+    EXPECT_FALSE(S->isOriginal())
+        << "node removed from tree cannot be marked as original";
   };
 
   std::vector<std::pair<std::string /*Input*/, std::string /*Expected*/>>