Fix newline parsing at trailing trivia

Sources/SwiftParser/Lexer/Cursor.swift

@@ -243,6 +243,9 @@ extension Lexer {
     /// If we have already lexed a token, the kind of the previously lexed token
     var previousTokenKind: RawTokenKind?
 
+    /// If we have already lexed a token, stores whether the previous lexeme‘s ending contains a newline.
+    var previousLexemeTrailingNewlinePresence: NewlinePresence?
+
     /// If the `previousTokenKind` is `.keyword`, the keyword kind. Otherwise
     /// `nil`.
     var previousKeyword: Keyword?
@@ -317,21 +320,25 @@ extension Lexer {
     /// If `tokenKind` is `.keyword`, the kind of keyword produced, otherwise
     /// `nil`.
     let keywordKind: Keyword?
+    /// Indicates whether the end of the lexed token text contains a newline.
+    let trailingNewlinePresence: Lexer.Cursor.NewlinePresence
 
     private init(
       _ tokenKind: RawTokenKind,
       flags: Lexer.Lexeme.Flags,
       error: Cursor.LexingDiagnostic?,
       stateTransition: StateTransition?,
       trailingTriviaLexingMode: Lexer.Cursor.TriviaLexingMode?,
-      keywordKind: Keyword?
+      keywordKind: Keyword?,
+      trailingNewlinePresence: Lexer.Cursor.NewlinePresence
     ) {
       self.tokenKind = tokenKind
       self.flags = flags
       self.error = error
       self.stateTransition = stateTransition
       self.trailingTriviaLexingMode = trailingTriviaLexingMode
       self.keywordKind = keywordKind
+      self.trailingNewlinePresence = trailingNewlinePresence
     }
 
     /// Create a lexer result. Note that keywords should use `Result.keyword`
@@ -341,7 +348,8 @@ extension Lexer {
       flags: Lexer.Lexeme.Flags = [],
       error: Cursor.LexingDiagnostic? = nil,
       stateTransition: StateTransition? = nil,
-      trailingTriviaLexingMode: Lexer.Cursor.TriviaLexingMode? = nil
+      trailingTriviaLexingMode: Lexer.Cursor.TriviaLexingMode? = nil,
+      trailingNewlinePresence: Lexer.Cursor.NewlinePresence = .absent
     ) {
       precondition(tokenKind != .keyword, "Use Result.keyword instead")
       self.init(
@@ -350,7 +358,8 @@ extension Lexer {
         error: error,
         stateTransition: stateTransition,
         trailingTriviaLexingMode: trailingTriviaLexingMode,
-        keywordKind: nil
+        keywordKind: nil,
+        trailingNewlinePresence: trailingNewlinePresence
       )
     }
 
@@ -362,7 +371,8 @@ extension Lexer {
         error: nil,
         stateTransition: nil,
         trailingTriviaLexingMode: nil,
-        keywordKind: kind
+        keywordKind: kind,
+        trailingNewlinePresence: .absent
       )
     }
   }
@@ -430,6 +440,16 @@ extension Lexer.Cursor {
       result = lexInRegexLiteral(lexemes.pointee[index...], existingPtr: lexemes)
     }
 
+    var flags = result.flags
+    if newlineInLeadingTrivia == .present {
+      flags.insert(.isAtStartOfLine)
+    }
+    if let previousLexemeTrailingNewlinePresence, previousLexemeTrailingNewlinePresence == .present {
+      flags.insert(.isAtStartOfLine)
+    }
+
+    self.previousLexemeTrailingNewlinePresence = result.trailingNewlinePresence
+
     if let stateTransition = result.stateTransition {
       self.stateStack.perform(stateTransition: stateTransition, stateAllocator: stateAllocator)
     }
@@ -438,18 +458,14 @@ extension Lexer.Cursor {
     let trailingTriviaStart = self
     if let trailingTriviaMode = result.trailingTriviaLexingMode ?? currentState.trailingTriviaLexingMode(cursor: self) {
       let triviaResult = self.lexTrivia(mode: trailingTriviaMode)
+      self.previousLexemeTrailingNewlinePresence = triviaResult.newlinePresence
       diagnostic = TokenDiagnostic(combining: diagnostic, triviaResult.error?.tokenDiagnostic(tokenStart: cursor))
     }
 
     if self.currentState.shouldPopStateWhenReachingNewlineInTrailingTrivia && self.is(at: "\r", "\n") {
       self.stateStack.perform(stateTransition: .pop, stateAllocator: stateAllocator)
     }
 
-    var flags = result.flags
-    if newlineInLeadingTrivia == .present {
-      flags.insert(.isAtStartOfLine)
-    }
-
     diagnostic = TokenDiagnostic(combining: diagnostic, result.error?.tokenDiagnostic(tokenStart: cursor))
 
     let lexeme = Lexer.Lexeme(
@@ -1889,7 +1905,7 @@ extension Lexer.Cursor {
           if character == UInt8(ascii: "\r") {
             _ = self.advance(matching: "\n")
           }
-          return Lexer.Result(.stringSegment, error: error)
+          return Lexer.Result(.stringSegment, error: error, trailingNewlinePresence: .present)
         } else {
           // Single line literals cannot span multiple lines.
           // Terminate the string here and go back to normal lexing (instead of `afterStringLiteral`)

Tests/SwiftParserTest/LexerTests.swift

@@ -1182,9 +1182,9 @@ public class LexerTests: XCTestCase {
       """#,
       lexemes: [
         LexemeSpec(.multilineStringQuote, leading: "  ", text: #"""""#, trailing: "\n"),
-        LexemeSpec(.stringSegment, text: "  line 1\n"),
-        LexemeSpec(.stringSegment, text: "  line 2\n"),
-        LexemeSpec(.stringSegment, text: "  "),
+        LexemeSpec(.stringSegment, text: "  line 1\n", flags: .isAtStartOfLine),
+        LexemeSpec(.stringSegment, text: "  line 2\n", flags: .isAtStartOfLine),
+        LexemeSpec(.stringSegment, text: "  ", flags: .isAtStartOfLine),
         LexemeSpec(.multilineStringQuote, text: #"""""#),
       ]
     )
@@ -1198,9 +1198,9 @@ public class LexerTests: XCTestCase {
       """#,
       lexemes: [
         LexemeSpec(.multilineStringQuote, leading: "  ", text: #"""""#, trailing: "\n"),
-        LexemeSpec(.stringSegment, text: "  line 1 ", trailing: "\\\n"),
-        LexemeSpec(.stringSegment, text: "  line 2\n"),
-        LexemeSpec(.stringSegment, text: "  "),
+        LexemeSpec(.stringSegment, text: "  line 1 ", trailing: "\\\n", flags: .isAtStartOfLine),
+        LexemeSpec(.stringSegment, text: "  line 2\n", flags: .isAtStartOfLine),
+        LexemeSpec(.stringSegment, text: "  ", flags: .isAtStartOfLine),
         LexemeSpec(.multilineStringQuote, text: #"""""#),
       ]
     )

Tests/SwiftParserTest/StatementTests.swift

@@ -763,4 +763,23 @@ final class StatementTests: XCTestCase {
         """
     )
   }
+
+  func testTrailingTriviaIncludesNewline() {
+    assertParse(
+      """
+      let a = 2/*
+      */let b = 3
+      """
+    )
+
+    assertParse(
+      """
+      let a = 2/*
+
+
+
+      */let b = 3
+      """
+    )
+  }
 }