Merge pull request #16264 from nkcsgexi/cache-absolute-pos-c-plus

libSyntax: cache absolute positions on SyntaxData.

Author: nkcsgexi

include/swift/Syntax/RawSyntax.h

@@ -458,6 +458,10 @@ class RawSyntax final
   llvm::Optional<AbsolutePosition>
   accumulateAbsolutePosition(AbsolutePosition &Pos) const;
 
+  /// Advance the provided AbsolutePosition by the first trivia of this node.
+  /// Return true if we found this trivia; otherwise false.
+  bool accumulateLeadingTrivia(AbsolutePosition &Pos) const;
+
   /// Print this piece of syntax recursively.
   void print(llvm::raw_ostream &OS, SyntaxPrintOptions Opts) const;
 

include/swift/Syntax/Syntax.h

@@ -190,7 +190,9 @@ class Syntax {
 
   /// Get the absolute position of this raw syntax: its offset, line,
   /// and column.
-  AbsolutePosition getAbsolutePosition() const;
+  AbsolutePosition getAbsolutePosition() const {
+    return Data->getAbsolutePosition();
+  }
 
   // TODO: hasSameStructureAs ?
 };

include/swift/Syntax/SyntaxData.h

@@ -78,6 +78,9 @@ class SyntaxData final
   /// If there is no parent, this is 0.
   const CursorIndex IndexInParent;
 
+  /// Cache the absolute position of this node.
+  Optional<AbsolutePosition> PositionCache;
+
   size_t numTrailingObjects(OverloadToken<AtomicCache<SyntaxData>>) const {
     return Raw->getNumChildren();
   }
@@ -129,6 +132,13 @@ class SyntaxData final
     return {getTrailingObjects<AtomicCache<SyntaxData>>(), getNumChildren()};
   }
 
+  /// Get the node immediately before this current node. Return 0 if we cannot
+  /// find such node.
+  RC<SyntaxData> getPreviousNode() const;
+
+  /// Get the absolute position without skipping the leading trivia of this node.
+  AbsolutePosition getAbsolutePositionWithLeadingTrivia() const;
+
 public:
   ~SyntaxData() {
     for (auto &I : getChildren())
@@ -235,6 +245,10 @@ class SyntaxData final
     });
   }
 
+  /// Calculate the absolute position of this node, use cache if the cache
+  /// is populated.
+  AbsolutePosition getAbsolutePosition() const;
+
   /// Returns true if the data node represents type syntax.
   bool isType() const;
 

lib/Syntax/RawSyntax.cpp

@@ -189,6 +189,24 @@ RawSyntax::accumulateAbsolutePosition(AbsolutePosition &Pos) const {
   return Ret;
 }
 
+bool RawSyntax::accumulateLeadingTrivia(AbsolutePosition &Pos) const {
+ if (isToken()) {
+    if (!isMissing()) {
+      for (auto &Leader: getLeadingTrivia())
+        Leader.accumulateAbsolutePosition(Pos);
+      return true;
+    }
+  } else {
+    for (auto &Child: getLayout()) {
+      if (!Child)
+        continue;
+      if (Child->accumulateLeadingTrivia(Pos))
+        return true;
+    }
+  }
+  return false;
+}
+
 void RawSyntax::print(llvm::raw_ostream &OS, SyntaxPrintOptions Opts) const {
   if (isMissing())
     return;

lib/Syntax/Syntax.cpp

@@ -96,59 +96,3 @@ llvm::Optional<Syntax> Syntax::getChild(const size_t N) const {
     return llvm::None;
   return Syntax {Root, ChildData.get()};
 }
-
-AbsolutePosition Syntax::getAbsolutePosition() const {
-  assert(getRoot().is<SourceFileSyntax>() &&
-         "Absolute position can only be calculated for nodes which has "
-         "SourceFile root");
-  AbsolutePosition Pos;
-
-  /// This visitor collects all of the nodes before this node to calculate its
-  /// offset from the begenning of the file.
-  class Visitor: public SyntaxVisitor {
-    AbsolutePosition &Pos;
-    const SyntaxData *Target;
-    bool Found = false;
-
-  public:
-    Visitor(AbsolutePosition &Pos, const SyntaxData *Target)
-        : Pos(Pos), Target(Target) {}
-    ~Visitor() { assert(Found); }
-    void visitPre(Syntax Node) override {
-      // Check if this node is the target;
-      Found |= Node.getDataPointer() == Target;
-    }
-    void visit(TokenSyntax Node) override {
-      // Ignore missing node and ignore the nodes after this node.
-      if (Found || Node.isMissing())
-        return;
-      // Collect all the offsets.
-      Node.getRaw()->accumulateAbsolutePosition(Pos);
-    }
-  } Calculator(Pos, getDataPointer());
-
-  /// This visitor visit the first token node of this node to accumulate its
-  /// leading trivia. Therefore, the calculated absolute location will point
-  /// to the actual token start.
-  class FirstTokenFinder: public SyntaxVisitor {
-    AbsolutePosition &Pos;
-    bool Found = false;
-
-  public:
-    FirstTokenFinder(AbsolutePosition &Pos): Pos(Pos) {}
-    void visit(TokenSyntax Node) override {
-      if (Found || Node.isMissing())
-        return;
-      Found = true;
-      for (auto &Leader : Node.getRaw()->getLeadingTrivia())
-        Leader.accumulateAbsolutePosition(Pos);
-    }
-  } FTFinder(Pos);
-
-  // Visit the root to get all the nodes before this node.
-  getRoot().accept(Calculator);
-
-  // Visit this node to accumulate the leading trivia of its first token.
-  const_cast<Syntax*>(this)->accept(FTFinder);
-  return Pos;
-}

lib/Syntax/SyntaxData.cpp

@@ -50,3 +50,38 @@ bool SyntaxData::isUnknown() const {
 void SyntaxData::dump(llvm::raw_ostream &OS) const {
   Raw->dump(OS, 0);
 }
+
+RC<SyntaxData> SyntaxData::getPreviousNode() const {
+  if (size_t N = getIndexInParent()) {
+    if (hasParent()) {
+      for (size_t I = N - 1; ; I--) {
+        if (auto C = getParent()->getChild(I)) {
+          return C;
+        }
+        if (I == 0)
+          break;
+      }
+    }
+  }
+  return hasParent() ? Parent->getPreviousNode() : nullptr;
+}
+
+AbsolutePosition SyntaxData::getAbsolutePositionWithLeadingTrivia() const {
+  if (PositionCache.hasValue())
+    return *PositionCache;
+  if (auto P = getPreviousNode()) {
+    auto Result = P->getAbsolutePositionWithLeadingTrivia();
+    P->getRaw()->accumulateAbsolutePosition(Result);
+    // FIXME: avoid using const_cast.
+    const_cast<SyntaxData*>(this)->PositionCache = Result;
+  } else {
+    const_cast<SyntaxData*>(this)->PositionCache = AbsolutePosition();
+  }
+  return *PositionCache;
+}
+
+AbsolutePosition SyntaxData::getAbsolutePosition() const {
+  auto Result = getAbsolutePositionWithLeadingTrivia();
+  getRaw()->accumulateLeadingTrivia(Result);
+  return Result;
+}