[clang-diff] Enable postorder traversal

include/clang/Tooling/ASTDiff/ASTDiff.h

public:
   StringRef getFilename() const;
 
   int getSize() const;
   NodeRef getRoot() const;
   NodeId getRootId() const;
   using PreorderIterator = const Node *;
   PreorderIterator begin() const;
   PreorderIterator end() const;
+  using PostorderIterator = NodeRefIterator;
+  PostorderIterator postorder_begin() const;
+  PostorderIterator postorder_end() const;
+  llvm::iterator_range<PostorderIterator> postorder() const;
 
   NodeRef getNode(NodeId Id) const;
 
   class Impl;
   std::unique_ptr<Impl> TreeImpl;
 };
 
 /// Represents a Clang AST node, alongside some additional information.

lib/Tooling/ASTDiff/ASTDiff.cpp

Impl(SyntaxTree *Parent,
       : Impl(Parent, dyn_cast<Decl>(Node), AST) {}
 
   SyntaxTree *Parent;
   ASTUnit &AST;
   PrintingPolicy TypePP;
   /// Nodes in preorder.
   std::vector<Node> Nodes;
   NodeList Leaves;
-  // Maps preorder indices to postorder ones.
-  std::vector<int> PostorderIds;
+  std::vector<int> PreorderToPostorderId;
   NodeList NodesBfs;
+  NodeList NodesPostorder;
   std::map<NodeId, SourceRange> TemplateArgumentLocations;
 
   int getSize() const { return Nodes.size(); }
   NodeRef getRoot() const { return getNode(getRootId()); }
   NodeId getRootId() const { return 0; }
   PreorderIterator begin() const { return &getRoot(); }
   PreorderIterator end() const { return begin() + getSize(); }
+  PostorderIterator postorder_begin() const { return NodesPostorder.begin(); }
+  PostorderIterator postorder_end() const { return NodesPostorder.end(); }
+  llvm::iterator_range<PostorderIterator> postorder() const {
+    return {postorder_begin(), postorder_end()};
+  }
 
   NodeRef getNode(NodeId Id) const { return Nodes[Id]; }
   Node &getMutableNode(NodeRef N) { return Nodes[N.getId()]; }
   Node &getMutableNode(NodeId Id) { return getMutableNode(getNode(Id)); }
 
 private:
   void initTree();
   void setLeftMostDescendants();
 };
 
 NodeRef NodeRefIterator::operator*() const { return Tree->getNode(*IdPointer); }
 
 NodeRefIterator &NodeRefIterator::operator++() { return ++IdPointer, *this; }
 NodeRefIterator &NodeRefIterator::operator+(int Offset) {
   return IdPointer += Offset, *this;
 }
+NodeRefIterator::difference_type NodeRefIterator::
+operator-(const NodeRefIterator &Other) const {
+  assert(Tree == Other.Tree &&
+         "Cannot subtract two iterators of different trees.");
+  return IdPointer - Other.IdPointer;
+}
 
 bool NodeRefIterator::operator!=(const NodeRefIterator &Other) const {
   assert(Tree == Other.Tree &&
          "Cannot compare two iterators of different trees.");
   return IdPointer != Other.IdPointer;
 }
 
+bool NodeRefIterator::operator==(const NodeRefIterator &Other) const {
+  return !(*this != Other);
+}
+
 static bool isSpecializedNodeExcluded(const Decl *D) { return D->isImplicit(); }
 static bool isSpecializedNodeExcluded(const Stmt *S) { return false; }
 static bool isSpecializedNodeExcluded(CXXCtorInitializer *I) {
   return !I->isWritten();
 }
 static bool isSpecializedNodeExcluded(const TemplateArgumentLoc *S) {
   return false;
 }
bool TraverseTemplateName(TemplateName Template) {
     PostTraverse(SavedState);
     return true;
   }
 };
 } // end anonymous namespace
 
 SyntaxTree::Impl::Impl(SyntaxTree *Parent, ASTUnit &AST)
     : Parent(Parent), AST(AST), TypePP(AST.getLangOpts()), Leaves(*this),
-      NodesBfs(*this) {
+      NodesBfs(*this), NodesPostorder(*this) {
   TypePP.AnonymousTagLocations = false;
 }
 
 SyntaxTree::Impl::Impl(SyntaxTree *Parent, Decl *N, ASTUnit &AST)
     : Impl(Parent, AST) {
   PreorderVisitor PreorderWalker(*this);
   PreorderWalker.TraverseDecl(N);
   initTree();
 }
 
 SyntaxTree::Impl::Impl(SyntaxTree *Parent, Stmt *N, ASTUnit &AST)
     : Impl(Parent, AST) {
   PreorderVisitor PreorderWalker(*this);
   PreorderWalker.TraverseStmt(N);
   initTree();
 }
 
-static std::vector<NodeId> getSubtreePostorder(SyntaxTree::Impl &Tree,
-                                               NodeId Root) {
-  std::vector<NodeId> Postorder;
-  std::function<void(NodeId)> Traverse = [&](NodeId Id) {
-    NodeRef N = Tree.getNode(Id);
-    for (NodeId Child : N.Children)
+static void getSubtreePostorder(NodeList &Ids, NodeRef Root) {
+  std::function<void(NodeRef)> Traverse = [&](NodeRef N) {
+    for (NodeRef Child : N)
       Traverse(Child);
-    Postorder.push_back(Id);
+    Ids.push_back(N.getId());
   };
   Traverse(Root);
-  return Postorder;
 }
 
 static void getSubtreeBfs(NodeList &Ids, NodeRef Root) {
   size_t Expanded = 0;
   Ids.push_back(Root.getId());
   while (Expanded < Ids.size())
     for (NodeRef Child : Ids[Expanded++])
       Ids.push_back(Child.getId());
 }
 
 void SyntaxTree::Impl::initTree() {
   setLeftMostDescendants();
   int PostorderId = 0;
-  PostorderIds.resize(getSize());
+  PreorderToPostorderId.resize(getSize());
   std::function<void(NodeRef)> PostorderTraverse = [&](NodeRef N) {
     for (NodeRef Child : N)
       PostorderTraverse(Child);
-    PostorderIds[N.getId()] = PostorderId;
+    PreorderToPostorderId[N.getId()] = PostorderId;
     ++PostorderId;
   };
   PostorderTraverse(getRoot());
   getSubtreeBfs(NodesBfs, getRoot());
+  getSubtreePostorder(NodesPostorder, getRoot());
 }
 
 void SyntaxTree::Impl::setLeftMostDescendants() {
   for (NodeRef Leaf : Leaves) {
     getMutableNode(Leaf).LeftMostDescendant = Leaf.getId();
     const Node *Parent, *Cur = &Leaf;
     while ((Parent = Cur->getParent()) && &Parent->getChild(0) == Cur) {
       Cur = Parent;
struct SNodeId {
   SNodeId &operator--() { return --Id, *this; }
   SNodeId operator+(int Other) const { return SNodeId(Id + Other); }
 };
 
 class Subtree {
 private:
   /// The parent tree.
   SyntaxTree::Impl &Tree;
-  /// Maps SNodeIds to original ids.
-  std::vector<NodeId> RootIds;
-  /// Maps subtree nodes to their leftmost descendants wtihin the subtree.
+  NodeRef Root;
+  /// Maps subtree nodes to their leftmost descendants wihin the subtree.
   std::vector<SNodeId> LeftMostDescendants;
 
 public:
   std::vector<SNodeId> KeyRoots;
 
-  Subtree(SyntaxTree::Impl &Tree, NodeId SubtreeRoot) : Tree(Tree) {
-    RootIds = getSubtreePostorder(Tree, SubtreeRoot);
+  Subtree(NodeRef Root) : Tree(Root.Tree), Root(Root) {
     int NumLeaves = setLeftMostDescendants();
     computeKeyRoots(NumLeaves);
   }
-  int getSize() const { return RootIds.size(); }
+  int getSize() const { return getNumberOfDescendants(Root); }
   NodeId getIdInRoot(SNodeId Id) const {
-    assert(Id > 0 && Id <= getSize() && "Invalid subtree node index.");
-    return RootIds[Id - 1];
+    return Tree.NodesPostorder[getPostorderIdInRoot(Id)].getId();
   }
   NodeRef getNode(SNodeId Id) const { return Tree.getNode(getIdInRoot(Id)); }
   SNodeId getLeftMostDescendant(SNodeId Id) const {
     assert(Id > 0 && Id <= getSize() && "Invalid subtree node index.");
     return LeftMostDescendants[Id - 1];
   }
-  /// Returns the postorder index of the leftmost descendant in the subtree.
-  NodeId getPostorderOffset() const {
-    return Tree.PostorderIds[getIdInRoot(SNodeId(1))];
+  NodeId getPostorderIdInRoot(SNodeId Id = SNodeId(1)) const {
+    assert(Id > 0 && Id <= getSize() && "Invalid subtree node index.");
+    return Id - 1 + Tree.PreorderToPostorderId[Root.LeftMostDescendant];
   }
 
 private:
   /// Returns the number of leafs in the subtree.
   int setLeftMostDescendants() {
     int NumLeaves = 0;
     LeftMostDescendants.resize(getSize());
     for (int I = 0; I < getSize(); ++I) {
       SNodeId SI(I + 1);
       NodeRef N = getNode(SI);
       NumLeaves += N.isLeaf();
-      assert(I == Tree.PostorderIds[getIdInRoot(SI)] - getPostorderOffset() &&
+      assert(I == Tree.PreorderToPostorderId[getIdInRoot(SI)] -
+                      getPostorderIdInRoot() &&
              "Postorder traversal in subtree should correspond to traversal in "
              "the root tree by a constant offset.");
-      LeftMostDescendants[I] = SNodeId(Tree.PostorderIds[N.LeftMostDescendant] -
-                                       getPostorderOffset());
+      LeftMostDescendants[I] =
+          SNodeId(Tree.PreorderToPostorderId[N.LeftMostDescendant] -
+                  getPostorderIdInRoot());
     }
     return NumLeaves;
   }
   void computeKeyRoots(int Leaves) {
     KeyRoots.resize(Leaves);
     std::unordered_set<int> Visited;
     int K = Leaves - 1;
     for (SNodeId I(getSize()); I > 0; --I) {
       SNodeId LeftDesc = getLeftMostDescendant(I);
       if (Visited.count(LeftDesc))
         continue;
       assert(K >= 0 && "K should be non-negative");
       KeyRoots[K] = I;
       Visited.insert(LeftDesc);
       --K;
     }
   }
 };
 
 /// Implementation of Zhang and Shasha's Algorithm for tree edit distance.
 /// Computes an optimal mapping between two trees using only insertion,
 /// deletion and update as edit actions (similar to the Levenshtein distance).
 class ZhangShashaMatcher {
   const ASTDiff::Impl &DiffImpl;
-  Subtree S1;
-  Subtree S2;
+  Subtree S1, S2;
   std::unique_ptr<std::unique_ptr<double[]>[]> TreeDist, ForestDist;
 
 public:
-  ZhangShashaMatcher(const ASTDiff::Impl &DiffImpl, SyntaxTree::Impl &T1,
-                     SyntaxTree::Impl &T2, NodeId Id1, NodeId Id2)
-      : DiffImpl(DiffImpl), S1(T1, Id1), S2(T2, Id2) {
+  ZhangShashaMatcher(const ASTDiff::Impl &DiffImpl, NodeRef N1, NodeRef N2)
+      : DiffImpl(DiffImpl), S1(N1), S2(N2) {
     TreeDist = llvm::make_unique<std::unique_ptr<double[]>[]>(
         size_t(S1.getSize()) + 1);
     ForestDist = llvm::make_unique<std::unique_ptr<double[]>[]>(
         size_t(S1.getSize()) + 1);
     for (int I = 0, E = S1.getSize() + 1; I < E; ++I) {
       TreeDist[I] = llvm::make_unique<double[]>(size_t(S2.getSize()) + 1);
       ForestDist[I] = llvm::make_unique<double[]>(size_t(S2.getSize()) + 1);
     }
bool ASTDiff::Impl::haveSameParents(NodeRef N1, NodeRef N2) const {
   const Node *P2 = N2.getParent();
   return (!P1 && !P2) || (P1 && P2 && getDst(*P1) == P2);
 }
 
 void ASTDiff::Impl::addOptimalMapping(NodeRef N1, NodeRef N2) {
   if (std::max(getNumberOfDescendants(N1), getNumberOfDescendants(N2)) >
       Options.MaxSize)
     return;
-  ZhangShashaMatcher Matcher(*this, T1, T2, N1.getId(), N2.getId());
+  ZhangShashaMatcher Matcher(*this, N1, N2);
   std::vector<std::pair<NodeId, NodeId>> R = Matcher.getMatchingNodes();
   for (const auto Tuple : R) {
     NodeRef N1 = T1.getNode(Tuple.first);
     NodeRef N2 = T2.getNode(Tuple.second);
     if (!getDst(N1) && !getSrc(N2))
       link(N1, N2);
   }
 }
if (Similarity >= Options.MinSimilarity && Similarity > HighestSimilarity) {
       HighestSimilarity = Similarity;
       Candidate = &N2;
     }
   }
   return Candidate;
 }
 
 void ASTDiff::Impl::matchBottomUp() {
-  std::vector<NodeId> Postorder = getSubtreePostorder(T1, T1.getRootId());
-  for (NodeId Id1 : Postorder) {
+  for (NodeRef N1 : T1.postorder()) {
+    NodeId Id1 = N1.getId();
     if (Id1 == T1.getRootId() && !getDst(T1.getRoot()) &&
         !getSrc(T2.getRoot())) {
       if (isMatchingPossible(T1.getRoot(), T2.getRoot())) {
         link(T1.getRoot(), T2.getRoot());
         addOptimalMapping(T1.getRoot(), T2.getRoot());
       }
       break;
     }
-    NodeRef N1 = T1.getNode(Id1);
     bool Matched = getDst(N1);
     bool MatchedChildren = false;
     for (NodeRef Child : N1) {
       if (getDst(Child)) {
         MatchedChildren = true;
         break;
       }
     }

 int SyntaxTree::getSize() const { return TreeImpl->getSize(); }
 NodeRef SyntaxTree::getRoot() const { return TreeImpl->getRoot(); }
 NodeId SyntaxTree::getRootId() const { return TreeImpl->getRootId(); }
 SyntaxTree::PreorderIterator SyntaxTree::begin() const {
   return TreeImpl->begin();
 }
 SyntaxTree::PreorderIterator SyntaxTree::end() const { return TreeImpl->end(); }
 
+SyntaxTree::PostorderIterator SyntaxTree::postorder_begin() const {
+  return TreeImpl->postorder_begin();
+}
+SyntaxTree::PostorderIterator SyntaxTree::postorder_end() const {
+  return TreeImpl->postorder_end();
+}
+llvm::iterator_range<SyntaxTree::PostorderIterator>
+SyntaxTree::postorder() const {
+  return TreeImpl->postorder();
+}
+
 } // end namespace diff
 } // end namespace clang