[SyntaxParse] Eagarly cache deferred range for deferred layout node

Author: rintaro

include/swift/Parse/ParsedRawSyntaxNode.h

@@ -52,6 +52,7 @@ class ParsedRawSyntaxNode {
   };
   struct DeferredLayoutNode {
     MutableArrayRef<ParsedRawSyntaxNode> Children;
+    CharSourceRange Range;
   };
   struct DeferredTokenNode {
     const ParsedTriviaPiece *TriviaPieces;
@@ -72,9 +73,9 @@ class ParsedRawSyntaxNode {
   /// Primary used for capturing a deferred missing token.
   bool IsMissing = false;
 
-  ParsedRawSyntaxNode(syntax::SyntaxKind k,
+  ParsedRawSyntaxNode(syntax::SyntaxKind k, CharSourceRange r,
                       MutableArrayRef<ParsedRawSyntaxNode> deferredNodes)
-    : DeferredLayout({deferredNodes}),
+    : DeferredLayout({deferredNodes, r}),
       SynKind(uint16_t(k)), TokKind(uint16_t(tok::unknown)),
       DK(DataKind::DeferredLayout) {
     assert(getKind() == k && "Syntax kind with too large value!");
@@ -211,14 +212,12 @@ class ParsedRawSyntaxNode {
     return copy;
   }
 
-  CharSourceRange getDeferredRange(bool includeTrivia) const {
+  CharSourceRange getDeferredRange() const {
     switch (DK) { 
     case DataKind::DeferredLayout:
-      return getDeferredLayoutRange(includeTrivia);
+      return getDeferredLayoutRange();
     case DataKind::DeferredToken:
-      return includeTrivia
-        ? getDeferredTokenRangeWithTrivia()
-        : getDeferredTokenRange();
+      return getDeferredTokenRangeWithTrivia();
     default:
       llvm_unreachable("node not deferred");
     }
@@ -243,20 +242,9 @@ class ParsedRawSyntaxNode {
 
   // Deferred Layout Data ====================================================//
 
-  CharSourceRange getDeferredLayoutRange(bool includeTrivia) const {
+  CharSourceRange getDeferredLayoutRange() const {
     assert(DK == DataKind::DeferredLayout);
-    auto HasValidRange = [includeTrivia](const ParsedRawSyntaxNode &Child) {
-      return !Child.isNull() && !Child.isMissing() &&
-        Child.getDeferredRange(includeTrivia).isValid();
-    };
-    auto first = llvm::find_if(getDeferredChildren(), HasValidRange);
-    if (first == getDeferredChildren().end())
-      return CharSourceRange();
-    auto last = llvm::find_if(llvm::reverse(getDeferredChildren()),
-                              HasValidRange);
-    auto firstRange = first->getDeferredRange(includeTrivia);
-    firstRange.widen(last->getDeferredRange(includeTrivia));
-    return firstRange;
+    return DeferredLayout.Range;
   }
   ArrayRef<ParsedRawSyntaxNode> getDeferredChildren() const {
     assert(DK == DataKind::DeferredLayout);

lib/Parse/ParseType.cpp

@@ -899,9 +899,8 @@ Parser::parseOldStyleProtocolComposition() {
       replacement = "Any";
     } else {
       auto extractText = [&](ParsedTypeSyntax &Type) -> StringRef {
-        auto SourceRange = Type.getRaw()
-          .getDeferredRange(/*includeTrivia=*/false);
-        return SourceMgr.extractText(SourceRange);
+        auto SourceRange = Type.getRaw().getDeferredRange();
+        return SourceMgr.extractText(SourceRange).trim();
       };
       auto Begin = Protocols.begin();
       replacement += extractText(*Begin);

lib/Parse/ParsedRawSyntaxNode.cpp

@@ -21,18 +21,31 @@ ParsedRawSyntaxNode
 ParsedRawSyntaxNode::makeDeferred(SyntaxKind k,
                                   MutableArrayRef<ParsedRawSyntaxNode> deferredNodes,
                                   SyntaxParsingContext &ctx) {
+  CharSourceRange range;
   if (deferredNodes.empty()) {
-    return ParsedRawSyntaxNode(k, {});
+    return ParsedRawSyntaxNode(k, range, {});
   }
   ParsedRawSyntaxNode *newPtr =
     ctx.getScratchAlloc().Allocate<ParsedRawSyntaxNode>(deferredNodes.size());
 
-  // uninitialized move;
   auto ptr = newPtr;
-  for (auto &node : deferredNodes)
-    :: new (static_cast<void *>(ptr++)) ParsedRawSyntaxNode(std::move(node));
+  for (auto &node : deferredNodes) {
+    // Cached range.
+    if (!node.isNull() && !node.isMissing()) {
+      auto nodeRange = node.getDeferredRange();
+      if (nodeRange.isValid()) {
+        if (range.isInvalid())
+          range = nodeRange;
+        else
+          range.widen(nodeRange);
+      }
+    }
 
-  return ParsedRawSyntaxNode(k, makeMutableArrayRef(newPtr, deferredNodes.size()));
+    // uninitialized move;
+    :: new (static_cast<void *>(ptr++)) ParsedRawSyntaxNode(std::move(node));
+  }
+  return ParsedRawSyntaxNode(k, range,
+                             makeMutableArrayRef(newPtr, deferredNodes.size()));
 }
 
 ParsedRawSyntaxNode

lib/Parse/ParsedRawSyntaxRecorder.cpp

@@ -138,7 +138,7 @@ void ParsedRawSyntaxRecorder::verifyElementRanges(ArrayRef<ParsedRawSyntaxNode>
       continue;
     CharSourceRange range = elem.isRecorded()
       ? elem.getRecordedRange()
-      : elem.getDeferredRange(/*includeTrivia=*/true);
+      : elem.getDeferredRange();
     if (range.isValid()) {
       assert((prevEndLoc.isInvalid() || range.getStart() == prevEndLoc)
              && "Non-contiguous child ranges?");