Register affect range in lookahead

This change will prevent us from skipping potential parts of a node when parsing incrementally

Author: StevenWong12

Sources/SwiftParser/Lookahead.swift

@@ -27,22 +27,29 @@ extension Parser {
     /// i.e. how far it looked ahead.
     var tokensConsumed: Int = 0
 
+    var parseTransition: IncrementalParseTransition?
+    /// Initialize the offset to the start of source
+    var currentOffset: Int
+
     private init(
       lexemes: Lexer.LexemeSequence,
-      currentToken: Lexer.Lexeme
+      currentToken: Lexer.Lexeme,
+      parseTransition: IncrementalParseTransition? = nil
     ) {
       self.lexemes = lexemes
       self.currentToken = currentToken
+      self.parseTransition = parseTransition
+      self.currentOffset = lexemes.getOffsetToStart(self.currentToken)
     }
 
     fileprivate init(cloning other: Parser) {
-      self.init(lexemes: other.lexemes, currentToken: other.currentToken)
+      self.init(lexemes: other.lexemes, currentToken: other.currentToken, parseTransition: other.parseTransition)
     }
 
     /// Initiates a lookahead session from the current point in this
     /// lookahead session.
     func lookahead() -> Lookahead {
-      return Lookahead(lexemes: self.lexemes, currentToken: self.currentToken)
+      return Lookahead(lexemes: self.lexemes, currentToken: self.currentToken, parseTransition: self.parseTransition)
     }
   }
 
@@ -89,6 +96,7 @@ extension Parser.Lookahead {
 
   mutating func consumeAnyToken() {
     tokensConsumed += 1
+    registerAffectRange()
     self.currentToken = self.lexemes.advance()
   }
 
@@ -399,4 +407,10 @@ extension Parser.Lookahead {
       }
     }
   }
+
+  fileprivate mutating func registerAffectRange() {
+    self.parseTransition?.registerAffectRange(at: self.currentOffset, length: self.currentToken.byteLength)
+    currentOffset += self.currentToken.byteLength
+    self.parseTransition?.registerAffectRange(at: self.currentOffset, length: self.lexemes.peek().byteLength)
+  }
 }

Sources/SwiftParser/Parser.swift

@@ -145,7 +145,12 @@ public struct Parser {
   ///            arena is created automatically, and `input` copied into the
   ///            arena. If non-`nil`, `input` must be within its registered
   ///            source buffer or allocator.
-  public init(_ input: UnsafeBufferPointer<UInt8>, maximumNestingLevel: Int? = nil, arena: ParsingSyntaxArena? = nil, parseTransition: IncrementalParseTransition? = nil) {
+  public init(
+    _ input: UnsafeBufferPointer<UInt8>,
+    maximumNestingLevel: Int? = nil,
+    arena: ParsingSyntaxArena? = nil,
+    parseTransition: IncrementalParseTransition? = nil
+  ) {
     self.maximumNestingLevel = maximumNestingLevel ?? Self.defaultMaximumNestingLevel
 
     var sourceBuffer: UnsafeBufferPointer<UInt8>
@@ -629,7 +634,11 @@ extension Parser {
 // MARK: Incremental Parsing
 extension Parser {
   mutating func loadCurrentSyntaxNodeFromCache(for kind: SyntaxKind) -> Syntax? {
-    guard let parseTransition = self.parseTransition else { return nil }
+    guard let parseTransition = self.parseTransition,
+      parseTransition.isValidTransition()
+    else {
+      return nil
+    }
 
     var lookUpHelper = IncrementalParseLookup(transition: parseTransition)
     let currentOffset = self.lexemes.getOffsetToStart(self.currentToken)

Sources/SwiftSyntax/IncrementalParseTransition.swift

@@ -45,9 +45,14 @@ public final class IncrementalParseReusedNodeCollector:
 /// Keeps track of a previously parsed syntax tree and the source edits that
 /// occurred since it was created.
 public final class IncrementalParseTransition {
-  fileprivate let previousTree: SourceFileSyntax
-  fileprivate let edits: ConcurrentEdits
-  fileprivate let reusedDelegate: IncrementalParseReusedNodeDelegate?
+  fileprivate var previousTree: SourceFileSyntax?
+  fileprivate var edits: ConcurrentEdits?
+  fileprivate var reusedDelegate: IncrementalParseReusedNodeDelegate?
+
+  fileprivate var previousLookaheadRange: [Int: Int] = [:]
+  /// Keep track of how far we would look when calling ``Lookahead``
+  /// Key is offset to buffer start and value is the length of we lookahead
+  fileprivate var cursorLookaheadRange: [Int: Int] = [:]
 
   /// - Parameters:
   ///   - previousTree: The previous tree to do lookups on.
@@ -56,14 +61,30 @@ public final class IncrementalParseTransition {
   ///   - reusedNodeDelegate: Optional delegate to accept information about the
   ///                         reused regions and nodes.
   public init(
-    previousTree: SourceFileSyntax,
-    edits: ConcurrentEdits,
+    previousTree: SourceFileSyntax? = nil,
+    edits: ConcurrentEdits? = nil,
     reusedNodeDelegate: IncrementalParseReusedNodeDelegate? = nil
   ) {
     self.previousTree = previousTree
     self.edits = edits
     self.reusedDelegate = reusedNodeDelegate
   }
+
+  public func setupTransition(tree: SourceFileSyntax, edits: ConcurrentEdits, delegate: IncrementalParseReusedNodeDelegate? = nil) {
+    self.previousTree = tree
+    self.edits = edits
+    self.reusedDelegate = delegate
+    self.previousLookaheadRange = cursorLookaheadRange
+    self.cursorLookaheadRange = [:]
+  }
+
+  public func registerAffectRange(at offset: Int, length: Int) {
+    self.cursorLookaheadRange[offset] = length
+  }
+
+  public func isValidTransition() -> Bool {
+    return previousTree != nil
+  }
 }
 
 fileprivate extension Sequence where Element: Comparable {
@@ -204,16 +225,14 @@ public struct ConcurrentEdits {
 /// updated source that was already parsed during a previous parse invocation.
 public struct IncrementalParseLookup {
   fileprivate let transition: IncrementalParseTransition
-  fileprivate var cursor: SyntaxCursor
 
   /// Create a new ``IncrementalParseLookup`` that can look nodes up based on the
   /// given ``IncrementalParseTransition``.
   public init(transition: IncrementalParseTransition) {
     self.transition = transition
-    self.cursor = .init(root: transition.previousTree.data)
   }
 
-  fileprivate var edits: ConcurrentEdits {
+  fileprivate var edits: ConcurrentEdits? {
     return transition.edits
   }
 
@@ -242,7 +261,7 @@ public struct IncrementalParseLookup {
     let node = cursorLookup(prevPosition: prevPosition, kind: kind)
     if let delegate = reusedDelegate, let node {
       delegate.parserReusedNode(
-        range: ByteSourceRange(offset: newOffset, length: node.byteSizeAfterTrimmingTrivia),
+        range: ByteSourceRange(offset: node.positionAfterSkippingLeadingTrivia.utf8Offset, length: node.byteSizeAfterTrimmingTrivia),
         previousNode: node
       )
     }
@@ -253,20 +272,28 @@ public struct IncrementalParseLookup {
     prevPosition: AbsolutePosition,
     kind: SyntaxKind
   ) -> Syntax? {
+    guard let data = transition.previousTree?.data else {
+      return nil
+    }
+
+    var cursor = SyntaxCursor(root: data)
     guard !cursor.finished else { return nil }
 
     while true {
-      if nodeAtCursorCanBeReused(prevPosition: prevPosition, kind: kind) {
+      if nodeAtCursorCanBeReused(cursor, prevPosition: prevPosition, kind: kind) {
         return cursor.asSyntaxNode
       }
       guard cursor.advanceToNextNode(at: prevPosition) else { return nil }
     }
   }
 
   fileprivate func nodeAtCursorCanBeReused(
+    _ cursor: SyntaxCursor,
     prevPosition: AbsolutePosition,
     kind: SyntaxKind
   ) -> Bool {
+    guard let edits = edits else { return false }
+
     let node = cursor.node
     if node.position != prevPosition {
       return false
@@ -298,7 +325,7 @@ public struct IncrementalParseLookup {
     }
     let nodeAffectRange = ByteSourceRange(
       offset: node.position.utf8Offset,
-      length: (node.raw.totalLength + nextLeafNodeLength).utf8Length
+      length: max(mergeLookaheadRange(at: node.position.utf8Offset, length: node.raw.totalLength.utf8Length), (node.raw.totalLength + nextLeafNodeLength).utf8Length)
     )
 
     for edit in edits.edits {
@@ -317,6 +344,8 @@ public struct IncrementalParseLookup {
   }
 
   fileprivate func translateToPreEditOffset(_ postEditOffset: Int) -> Int? {
+    guard let edits = edits else { return nil }
+
     var offset = postEditOffset
     for edit in edits.edits {
       if edit.range.offset > offset {
@@ -332,6 +361,21 @@ public struct IncrementalParseLookup {
     }
     return offset
   }
+
+  fileprivate func mergeLookaheadRange(at start: Int, length: Int) -> Int {
+    var totalLength = start + length
+
+    let targetRanges = transition.previousLookaheadRange.filter { $0.key >= totalLength }.sorted(by: { $0.key < $1.key })
+
+    for targetRange in targetRanges {
+      if targetRange.key != totalLength {
+        break
+      }
+      totalLength += targetRange.value
+    }
+
+    return totalLength
+  }
 }
 
 /// Functions as an iterator that walks the tree looking for nodes with a

Sources/_SwiftSyntaxTestSupport/IncrementalParseTestUtils.swift

@@ -30,10 +30,11 @@ public func assertIncrementalParse(
   let originalString = String(originalSource)
   let editedString = String(editedSource)
 
-  let originalTree = Parser.parse(source: originalString)
+  let transition = IncrementalParseTransition()
+  let originalTree = Parser.parse(source: originalString, parseTransition: transition)
 
   let reusedNodesCollector = IncrementalParseReusedNodeCollector()
-  let transition = IncrementalParseTransition(previousTree: originalTree, edits: concurrentEdits, reusedNodeDelegate: reusedNodesCollector)
+  transition.setupTransition(tree: originalTree, edits: concurrentEdits, delegate: reusedNodesCollector)
 
   let newTree = Parser.parse(source: editedString)
   let incrementallyParsedNewTree = Parser.parse(source: editedString, parseTransition: transition)

Tests/SwiftParserTest/IncrementalParsingTests.swift

@@ -25,8 +25,7 @@ public class IncrementalParsingTests: XCTestCase {
     )
   }
 
-  public func testReusedNode() throws {
-    try XCTSkipIf(true, "Swift parser does not handle node reuse yet")
+  public func testReusedNode() {
     assertIncrementalParse(
       """
       struct A⏩️⏸️A⏪️ {}
@@ -37,9 +36,8 @@ public class IncrementalParsingTests: XCTestCase {
       ]
     )
   }
-  
+
   public func testTrailingClosure() {
-    XCTExpectFailure("WIP: Add lookahead for incremental parse")
     assertIncrementalParse(
       """
       foo() {}
@@ -49,4 +47,19 @@ public class IncrementalParsingTests: XCTestCase {
       """
     )
   }
+
+  public func testMultiFunctionCall() {
+    assertIncrementalParse(
+      """
+      foo() {}
+      foo1() {}
+      foo2() {}
+      ⏩️⏸️foo3() {}⏪️
+      """,
+      reusedNodes: [
+        ReusedNodeSpec("foo() {}", kind: .codeBlockItem),
+        ReusedNodeSpec("foo1() {}", kind: .codeBlockItem),
+      ]
+    )
+  }
 }