[Syntax] Assert invariants on tree structure and fix a bug in mutations

Add checks for some structural invariants when building and mutating
the syntax trees.

Fix a bug failing the invariants after mutations: the parent of nodes
added into the tree was null.

4 files changed, 56 insertions, 9 deletions

diff --git a/clang/lib/Tooling/Syntax/BuildTree.cpp b/clang/lib/Tooling/Syntax/BuildTree.cpp
index 7357cab6f97..aa8844771d3 100644
--- a/clang/lib/Tooling/Syntax/BuildTree.cpp
+++ b/clang/lib/Tooling/Syntax/BuildTree.cpp
@@ -92,7 +92,9 @@ public:
     Pending.foldChildren(Arena, Tokens.drop_back(),
                          new (Arena.allocator()) syntax::TranslationUnit);
 
-    return cast<syntax::TranslationUnit>(std::move(Pending).finalize());
+    auto *TU = cast<syntax::TranslationUnit>(std::move(Pending).finalize());
+    TU->assertInvariantsRecursive();
+    return TU;
   }
 
   /// getRange() finds the syntax tokens corresponding to the passed source
diff --git a/clang/lib/Tooling/Syntax/Mutations.cpp b/clang/lib/Tooling/Syntax/Mutations.cpp
index 7278aff5f18..72458528202 100644
--- a/clang/lib/Tooling/Syntax/Mutations.cpp
+++ b/clang/lib/Tooling/Syntax/Mutations.cpp
@@ -29,30 +29,43 @@ class syntax::MutationsImpl {
 public:
   /// Add a new node with a specified role.
   static void addAfter(syntax::Node *Anchor, syntax::Node *New, NodeRole Role) {
+    assert(Anchor != nullptr);
     assert(New->Parent == nullptr);
     assert(New->NextSibling == nullptr);
     assert(!New->isDetached());
     assert(Role != NodeRole::Detached);
 
     New->Role = static_cast<unsigned>(Role);
-    Anchor->parent()->replaceChildRangeLowLevel(Anchor, Anchor, New);
+    auto *P = Anchor->parent();
+    P->replaceChildRangeLowLevel(Anchor, Anchor, New);
+
+    P->assertInvariants();
   }
 
   /// Replace the node, keeping the role.
   static void replace(syntax::Node *Old, syntax::Node *New) {
+    assert(Old != nullptr);
+    assert(Old->Parent != nullptr);
+    assert(Old->canModify());
     assert(New->Parent == nullptr);
     assert(New->NextSibling == nullptr);
     assert(New->isDetached());
 
     New->Role = Old->Role;
-    Old->parent()->replaceChildRangeLowLevel(findPrevious(Old),
-                                             Old->nextSibling(), New);
+    auto *P = Old->parent();
+    P->replaceChildRangeLowLevel(findPrevious(Old), Old->nextSibling(), New);
+
+    P->assertInvariants();
   }
 
   /// Completely remove the node from its parent.
   static void remove(syntax::Node *N) {
-    N->parent()->replaceChildRangeLowLevel(findPrevious(N), N->nextSibling(),
-                                           /*New=*/nullptr);
+    auto *P = N->parent();
+    P->replaceChildRangeLowLevel(findPrevious(N), N->nextSibling(),
+                                 /*New=*/nullptr);
+
+    P->assertInvariants();
+    N->assertInvariants();
   }
 
 private:
@@ -69,6 +82,7 @@ private:
 };
 
 void syntax::removeStatement(syntax::Arena &A, syntax::Statement *S) {
+  assert(S);
   assert(S->canModify());
 
   if (isa<CompoundStatement>(S->parent())) {
diff --git a/clang/lib/Tooling/Syntax/Synthesis.cpp b/clang/lib/Tooling/Syntax/Synthesis.cpp
index 73ae71f9a6c..cbd9579f4f0 100644
--- a/clang/lib/Tooling/Syntax/Synthesis.cpp
+++ b/clang/lib/Tooling/Syntax/Synthesis.cpp
@@ -26,7 +26,9 @@ clang::syntax::Leaf *syntax::createPunctuation(clang::syntax::Arena &A,
                     .second;
   assert(Tokens.size() == 1);
   assert(Tokens.front().kind() == K);
-  return new (A.allocator()) clang::syntax::Leaf(Tokens.begin());
+  auto *L = new (A.allocator()) clang::syntax::Leaf(Tokens.begin());
+  L->assertInvariants();
+  return L;
 }
 
 clang::syntax::EmptyStatement *
@@ -34,5 +36,6 @@ syntax::createEmptyStatement(clang::syntax::Arena &A) {
   auto *S = new (A.allocator()) clang::syntax::EmptyStatement;
   FactoryImpl::prependChildLowLevel(S, createPunctuation(A, clang::tok::semi),
                                     NodeRole::Unknown);
+  S->assertInvariants();
   return S;
 }
diff --git a/clang/lib/Tooling/Syntax/Tree.cpp b/clang/lib/Tooling/Syntax/Tree.cpp
index d7436ae9513..5282857df5a 100644
--- a/clang/lib/Tooling/Syntax/Tree.cpp
+++ b/clang/lib/Tooling/Syntax/Tree.cpp
@@ -11,6 +11,7 @@
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/Support/Casting.h"
+#include <cassert>
 
 using namespace clang;
 
@@ -91,8 +92,10 @@ void syntax::Tree::replaceChildRangeLowLevel(Node *BeforeBegin, Node *End,
 
   if (New) {
     auto *Last = New;
-    while (auto *Next = Last->nextSibling())
-      Last = Next;
+    for (auto *N = New; N != nullptr; N = N->nextSibling()) {
+      Last = N;
+      N->Parent = this;
+    }
     Last->NextSibling = End;
   }
 
@@ -189,6 +192,31 @@ std::string syntax::Node::dumpTokens(const Arena &A) const {
   return OS.str();
 }
 
+void syntax::Node::assertInvariants() const {
+#ifndef NDEBUG
+  if (isDetached())
+    assert(parent() == nullptr);
+  else
+    assert(parent() != nullptr);
+
+  auto *T = dyn_cast<Tree>(this);
+  if (!T)
+    return;
+  for (auto *C = T->firstChild(); C; C = C->nextSibling()) {
+    if (T->isOriginal())
+      assert(C->isOriginal());
+    assert(!C->isDetached());
+    assert(C->parent() == T);
+  }
+#endif
+}
+
+void syntax::Node::assertInvariantsRecursive() const {
+#ifndef NDEBUG
+  traverse(this, [&](const syntax::Node *N) { N->assertInvariants(); });
+#endif
+}
+
 syntax::Leaf *syntax::Tree::firstLeaf() {
   auto *T = this;
   while (auto *C = T->firstChild()) {