Merge pull request #27024 from rintaro/syntaxparse-endloc-composition

[SyntaxParse] Fix the end location for CompositionTypeRepr

Author: rintaro

include/swift/Syntax/Syntax.h

@@ -39,6 +39,7 @@ namespace syntax {
 
 struct SyntaxVisitor;
 class SourceFileSyntax;
+class TokenSyntax;
 
 template <typename SyntaxNode>
 SyntaxNode make(RC<RawSyntax> Raw) {
@@ -168,6 +169,15 @@ class Syntax {
   /// Returns true if the node is "present" in the source.
   bool isPresent() const;
 
+
+  /// Returns the first non-missing token in this syntax. Returns None if there
+  /// is no non-missing token.
+  Optional<TokenSyntax> getFirstToken();
+
+  /// Returns the last non-missing token in this syntax. Returns None if there
+  /// is no non-missing token.
+  Optional<TokenSyntax> getLastToken();
+
   /// Print the syntax node with full fidelity to the given output stream.
   void print(llvm::raw_ostream &OS, SyntaxPrintOptions Opts = SyntaxPrintOptions()) const;
 

include/swift/Syntax/SyntaxData.h

@@ -149,6 +149,10 @@ class SyntaxData final
   /// node does not contain non-missing tokens.
   RC<SyntaxData> getFirstToken() const;
 
+  /// Get the last non-missing token node in this tree. Return nullptr if this
+  /// node does not contain non-missing tokens.
+  RC<SyntaxData> getLastToken() const;
+
   ~SyntaxData() {
     for (auto &I : getChildren())
       I.~AtomicCache<SyntaxData>();

lib/Parse/ASTGen.cpp

@@ -306,7 +306,7 @@ TypeRepr *ASTGen::generate(CompositionTypeSyntax Type, SourceLoc &Loc) {
 
   auto FirstTypeLoc = advanceLocBegin(Loc, FirstElem);
   auto FirstAmpersandLoc = advanceLocBegin(Loc, *FirstElem.getAmpersand());
-  auto LastTypeLoc = advanceLocBegin(Loc, LastElem);
+  auto LastTypeLoc = advanceLocBegin(Loc, *LastElem.getLastToken());
   return CompositionTypeRepr::create(Context, ElemTypes, FirstTypeLoc,
                                      {FirstAmpersandLoc, LastTypeLoc});
 }

lib/Syntax/Syntax.cpp

@@ -13,6 +13,7 @@
 #include "swift/Syntax/Syntax.h"
 #include "swift/Syntax/SyntaxData.h"
 #include "swift/Syntax/SyntaxVisitor.h"
+#include "swift/Syntax/TokenSyntax.h"
 
 using namespace swift;
 using namespace swift::syntax;
@@ -96,3 +97,15 @@ llvm::Optional<Syntax> Syntax::getChild(const size_t N) const {
     return llvm::None;
   return Syntax {Root, ChildData.get()};
 }
+
+Optional<TokenSyntax> Syntax::getFirstToken() {
+  if (auto tok = getData().getFirstToken())
+    return TokenSyntax(Root, tok.get());
+  return None;
+}
+
+Optional<TokenSyntax> Syntax::getLastToken() {
+  if (auto tok = getData().getLastToken())
+    return TokenSyntax(Root, tok.get());
+  return None;
+}

lib/Syntax/SyntaxData.cpp

@@ -105,6 +105,26 @@ RC<SyntaxData> SyntaxData::getFirstToken() const {
   return nullptr;
 }
 
+RC<SyntaxData> SyntaxData::getLastToken() const {
+  if (getRaw()->isToken()) {
+    // Get a reference counted version of this
+    assert(hasParent() && "The syntax tree should not conisist only of the root");
+    return getParent()->getChild(getIndexInParent());
+  }
+  for (size_t I = getNumChildren(); I != 0; --I) {
+    if (auto Child = getChild(I - 1)) {
+      if (Child->getRaw()->isMissing())
+        continue;
+      if (Child->getRaw()->isToken()) {
+        return Child;
+      } else if (auto Token = Child->getLastToken()) {
+        return Token;
+      }
+    }
+  }
+  return nullptr;
+}
+
 AbsolutePosition SyntaxData::getAbsolutePositionBeforeLeadingTrivia() const {
   if (PositionCache.hasValue())
     return *PositionCache;

test/Parse/composition_type_range.swift

@@ -0,0 +1,4 @@
+// RUN: not %target-swift-frontend -dump-ast %s 2>&1 | %FileCheck %s
+
+typealias A = B & protocol<C, D>
+// CHECK: (typealias range=[{{.+}}.swift:[[@LINE-1]]:1 - line:[[@LINE-1]]:32] "A" {{.*}})