WIP: port incremental parse to CodeBlockItem

Author: StevenWong12

Sources/SwiftParser/IncrementalParseTransition.swift

@@ -26,29 +26,40 @@ public protocol IncrementalParseReusedNodeDelegate {
   ///   - range: The source region of the currently parsed source.
   ///   - previousNode: The node from the previous tree that is associated with
   ///                   the skipped source region.
-  func parserReusedNode(range: ByteSourceRange, previousNode: Syntax)
+  func parserReusedNode(previousNode: Syntax)
 }
 
 /// An implementation of `IncrementalParseReusedNodeDelegate` that just collects
 /// the range and re-used node into an array.
 public final class IncrementalParseReusedNodeCollector:
   IncrementalParseReusedNodeDelegate
 {
-  public var rangeAndNodes: [(ByteSourceRange, Syntax)] = []
+  public var nodes: [Syntax] = []
 
   public init() {}
 
-  public func parserReusedNode(range: ByteSourceRange, previousNode: Syntax) {
-    rangeAndNodes.append((range, previousNode))
+  public func parserReusedNode(previousNode: Syntax) {
+    nodes.append(previousNode)
+  }
+}
+
+public final class IncrementalParseNodeAffectRangeCollector {
+  fileprivate var nodeAffectRange: [RawSyntax.ID: Int] = [:]
+
+  public init() {}
+
+  @_spi(RawSyntax)
+  public func registerNodeForIncrementalParse(node: RawSyntax, length: Int) {
+    self.nodeAffectRange[node.id] = length
   }
 }
 
 /// 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?
 
   /// - Parameters:
   ///   - previousTree: The previous tree to do lookups on.
@@ -59,7 +70,7 @@ public final class IncrementalParseTransition {
   public init(
     previousTree: SourceFileSyntax,
     edits: ConcurrentEdits,
-    reusedNodeDelegate: IncrementalParseReusedNodeDelegate? = nil
+    reusedNodeDelegate: IncrementalParseReusedNodeDelegate?
   ) {
     self.previousTree = previousTree
     self.edits = edits
@@ -71,23 +82,31 @@ public final class IncrementalParseTransition {
 /// updated source that was already parsed during a previous parse invocation.
 public struct IncrementalParseLookup {
   fileprivate let transition: IncrementalParseTransition
+
   fileprivate var cursor: SyntaxCursor
 
+  fileprivate let nodeAffectRangeCollector: IncrementalParseNodeAffectRangeCollector
+
   /// Create a new ``IncrementalParseLookup`` that can look nodes up based on the
   /// given ``IncrementalParseTransition``.
-  public init(transition: IncrementalParseTransition) {
+  public init(transition: IncrementalParseTransition, nodeAffectRangeCollector: IncrementalParseNodeAffectRangeCollector) {
     self.transition = transition
     self.cursor = .init(root: Syntax(transition.previousTree))
+    self.nodeAffectRangeCollector = nodeAffectRangeCollector
   }
 
-  fileprivate var edits: ConcurrentEdits {
+  fileprivate var edits: ConcurrentEdits? {
     return transition.edits
   }
 
   fileprivate var reusedDelegate: IncrementalParseReusedNodeDelegate? {
     return transition.reusedDelegate
   }
 
+  fileprivate var nodeAffectRange: [RawSyntax.ID: Int] {
+    return nodeAffectRangeCollector.nodeAffectRange
+  }
+
   /// Does a lookup to see if the current source `offset` should be associated
   /// with a known ``Syntax`` node and its region skipped during parsing.
   ///
@@ -109,7 +128,6 @@ 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),
         previousNode: node
       )
     }
@@ -134,6 +152,8 @@ public struct IncrementalParseLookup {
     prevPosition: AbsolutePosition,
     kind: SyntaxKind
   ) -> Bool {
+    guard let edits = edits else { return false }
+
     let node = cursor.node
     if node.position != prevPosition {
       return false
@@ -148,24 +168,13 @@ public struct IncrementalParseLookup {
       return true
     }
 
-    // Node can also not be reused if an edit has been made in the next token's
-    // text, e.g. because `private struct Foo {}` parses as a CodeBlockItem with
-    // a StructDecl inside and `private struc Foo {}` parses as two
-    // CodeBlockItems one for `private` and one for `struc Foo {}`
-    var nextLeafNodeLength: SourceLength = .zero
-    if let nextSibling = cursor.nextSibling {
-      // Fast path check: if next sibling is before all the edits then we can
-      // re-use the node.
-      if !edits.edits.isEmpty && edits.edits.first!.range.offset > nextSibling.endPosition.utf8Offset {
-        return true
-      }
-      if let nextToken = nextSibling.firstToken(viewMode: .sourceAccurate) {
-        nextLeafNodeLength = nextToken.leadingTriviaLength + nextToken.contentLength
-      }
+    guard let nodeAffectRangeLength = self.nodeAffectRange[node.raw.id] else {
+      return false
     }
+
     let nodeAffectRange = ByteSourceRange(
       offset: node.position.utf8Offset,
-      length: (node.totalLength + nextLeafNodeLength).utf8Length
+      length: nodeAffectRangeLength
     )
 
     for edit in edits.edits {
@@ -184,6 +193,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 {

Sources/SwiftParser/Lexer/LexemeSequence.swift

@@ -32,17 +32,28 @@ extension Lexer {
     /// usually less than 0.1% of the memory allocated by the syntax arena.
     var lexerStateAllocator = BumpPtrAllocator(slabSize: 256)
 
+    var offsetToNextTokenEnd: Int {
+      self.getOffsetToStart(self.nextToken) + self.nextToken.byteLength
+    }
+
+    var lookaheadTracker: UnsafeMutablePointer<LookaheadTracker>
+
     fileprivate init(sourceBufferStart: Lexer.Cursor, cursor: Lexer.Cursor) {
       self.sourceBufferStart = sourceBufferStart
       self.cursor = cursor
       self.nextToken = self.cursor.nextToken(sourceBufferStart: self.sourceBufferStart, stateAllocator: lexerStateAllocator)
+      self.lookaheadTracker = .allocate(capacity: 1)
     }
 
     @_spi(Testing)
     public mutating func next() -> Lexer.Lexeme? {
       return self.advance()
     }
 
+    public func recordFurthestOffset() {
+      self.lookaheadTracker.pointee.furthestOffset = self.offsetToNextTokenEnd
+    }
+
     mutating func advance() -> Lexer.Lexeme {
       defer {
         self.nextToken = self.cursor.nextToken(sourceBufferStart: self.sourceBufferStart, stateAllocator: lexerStateAllocator)

Sources/SwiftParser/Lookahead.swift

@@ -33,6 +33,7 @@ extension Parser {
     ) {
       self.lexemes = lexemes
       self.currentToken = currentToken
+      self.lexemes.recordFurthestOffset()
     }
 
     fileprivate init(cloning other: Parser) {
@@ -90,6 +91,7 @@ extension Parser.Lookahead {
   mutating func consumeAnyToken() {
     tokensConsumed += 1
     self.currentToken = self.lexemes.advance()
+    self.lexemes.recordFurthestOffset()
   }
 
   mutating func consumeAnyToken(remapping: RawTokenKind) {

Sources/SwiftParser/Parser.swift

@@ -103,6 +103,8 @@ public struct Parser {
 
   let parseTransition: IncrementalParseTransition?
 
+  var parseNodeAffectRange: IncrementalParseNodeAffectRangeCollector?
+
   /// A default maximum nesting level that is used if the client didn't
   /// explicitly specify one. Debug builds of the parser comume a lot more stack
   /// space and thus have a lower default maximum nesting level.
@@ -113,7 +115,12 @@ public struct Parser {
   #endif
 
   /// Initializes a ``Parser`` from the given string.
-  public init(_ input: String, maximumNestingLevel: Int? = nil, parseTransition: IncrementalParseTransition? = nil) {
+  public init(
+    _ input: String,
+    maximumNestingLevel: Int? = nil,
+    parseNodeAffectRange: IncrementalParseNodeAffectRangeCollector? = nil,
+    parseTransition: IncrementalParseTransition? = nil
+  ) {
     self.maximumNestingLevel = maximumNestingLevel ?? Self.defaultMaximumNestingLevel
 
     self.arena = ParsingSyntaxArena(
@@ -129,6 +136,7 @@ public struct Parser {
     self.parseTransition = parseTransition
     self.lexemes = Lexer.tokenize(interned)
     self.currentToken = self.lexemes.advance()
+    self.parseNodeAffectRange = parseNodeAffectRange
   }
 
   /// Initializes a ``Parser`` from the given input buffer.
@@ -145,7 +153,13 @@ 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,
+    parseNodeAffectRange: IncrementalParseNodeAffectRangeCollector? = nil,
+    parseTransition: IncrementalParseTransition? = nil,
+    arena: ParsingSyntaxArena? = nil
+  ) {
     self.maximumNestingLevel = maximumNestingLevel ?? Self.defaultMaximumNestingLevel
 
     var sourceBuffer: UnsafeBufferPointer<UInt8>
@@ -163,6 +177,7 @@ public struct Parser {
     self.parseTransition = parseTransition
     self.lexemes = Lexer.tokenize(sourceBuffer)
     self.currentToken = self.lexemes.advance()
+    self.parseNodeAffectRange = parseNodeAffectRange
   }
 
   mutating func missingToken(_ kind: RawTokenKind, text: SyntaxText? = nil) -> RawTokenSyntax {
@@ -241,6 +256,7 @@ public struct Parser {
 extension Parser {
   /// Retrieves the token following the current token without consuming it.
   func peek() -> Lexer.Lexeme {
+    lexemes.recordFurthestOffset()
     return self.lexemes.peek()
   }
 }
@@ -637,9 +653,12 @@ extension Parser {
 // MARK: Incremental Parsing
 extension Parser {
   mutating func loadCurrentSyntaxNodeFromCache(for kind: SyntaxKind) -> Syntax? {
-    guard let parseTransition = self.parseTransition else { return nil }
-
-    var lookUpHelper = IncrementalParseLookup(transition: parseTransition)
+    guard let parseTransition = self.parseTransition,
+      let parseNodeAffectRange = self.parseNodeAffectRange
+    else {
+      return nil
+    }
+    var lookUpHelper = IncrementalParseLookup(transition: parseTransition, nodeAffectRangeCollector: parseNodeAffectRange)
     let currentOffset = self.lexemes.getOffsetToStart(self.currentToken)
     if let node = lookUpHelper.lookUp(currentOffset, kind: kind) {
       self.lexemes.advance(by: node.byteSize, currentToken: &self.currentToken)
@@ -648,4 +667,31 @@ extension Parser {
 
     return nil
   }
+
+  func registerNodeForIncrementalParse(node: RawSyntax, startToken: Lexer.Lexeme) {
+    if let parseNodeAffectRange {
+      parseNodeAffectRange.registerNodeForIncrementalParse(
+        node: node,
+        length: max(lookaheadFurthestOffset - self.lexemes.getOffsetToStart(startToken), node.byteLength + currentToken.byteLength)
+      )
+    }
+  }
+
+  public var lookaheadFurthestOffset: Int {
+    return lexemes.lookaheadTracker.pointee.furthestOffset
+  }
+}
+
+/// Record the furthest offset to `sourceBufferStart` that is reached by ``Parser.Lookahead`` or ``Parser.peek()`` in ``Lexer.LexemeSequence``
+public struct LookaheadTracker {
+  private var _furthestOffset: Int
+
+  public var furthestOffset: Int {
+    get { _furthestOffset }
+    set { _furthestOffset = max(newValue, furthestOffset) }
+  }
+
+  init(furthestOffset: Int = 0) {
+    self._furthestOffset = furthestOffset
+  }
 }

Sources/SwiftParser/TopLevel.swift

@@ -151,7 +151,9 @@ extension Parser {
   ///     statement → compiler-control-statement
   ///     statements → statement statements?
   mutating func parseCodeBlockItem(isAtTopLevel: Bool, allowInitDecl: Bool) -> RawCodeBlockItemSyntax? {
+    let startToken = self.currentToken
     if let syntax = self.loadCurrentSyntaxNodeFromCache(for: .codeBlockItem) {
+      self.registerNodeForIncrementalParse(node: syntax.raw, startToken: startToken)
       return RawCodeBlockItemSyntax(syntax.raw)
     }
 
@@ -187,12 +189,17 @@ extension Parser {
     if item.raw.isEmpty && semi == nil && trailingSemis.isEmpty {
       return nil
     }
-    return RawCodeBlockItemSyntax(
+
+    let result = RawCodeBlockItemSyntax(
       item: item,
       semicolon: semi,
       RawUnexpectedNodesSyntax(trailingSemis, arena: self.arena),
       arena: self.arena
     )
+
+    self.registerNodeForIncrementalParse(node: result.raw, startToken: startToken)
+
+    return result
   }
 
   private mutating func parseStatementItem() -> RawCodeBlockItemSyntax.Item {

Sources/SwiftParser/generated/Parser+Entry.swift

@@ -19,9 +19,10 @@ extension Parser {
   /// `Parser.init` for more details.
   public static func parse(
     source: String,
+    parseNodeAffectRange: IncrementalParseNodeAffectRangeCollector? = nil,
     parseTransition: IncrementalParseTransition? = nil
   ) -> SourceFileSyntax {
-    var parser = Parser(source, parseTransition: parseTransition)
+    var parser = Parser(source, parseNodeAffectRange: parseNodeAffectRange, parseTransition: parseTransition)
     return SourceFileSyntax.parse(from: &parser)
   }
 
@@ -30,9 +31,10 @@ extension Parser {
   public static func parse(
     source: UnsafeBufferPointer<UInt8>,
     maximumNestingLevel: Int? = nil,
+    parseNodeAffectRange: IncrementalParseNodeAffectRangeCollector? = nil,
     parseTransition: IncrementalParseTransition? = nil
   ) -> SourceFileSyntax {
-    var parser = Parser(source, maximumNestingLevel: maximumNestingLevel, parseTransition: parseTransition)
+    var parser = Parser(source, maximumNestingLevel: maximumNestingLevel, parseNodeAffectRange: parseNodeAffectRange, parseTransition: parseTransition)
     return SourceFileSyntax.parse(from: &parser)
   }
 }

Sources/SwiftSyntax/Raw/RawSyntax.swift

@@ -918,6 +918,12 @@ extension RawSyntax {
   }
 }
 
+extension RawSyntax: Identifiable {
+  public var id: UnsafeRawPointer {
+    return UnsafeRawPointer(pointer)
+  }
+}
+
 #if DEBUG
 /// See `SyntaxMemoryLayout`.
 var RawSyntaxDataMemoryLayouts: [String: SyntaxMemoryLayout.Value] = [

Sources/SwiftSyntax/Syntax.swift

@@ -448,6 +448,11 @@ public extension SyntaxProtocol {
     return ByteSourceRange(offset: position.utf8Offset, length: byteSize)
   }
 
+  /// The  byte source range of this node exluding leading and trailing trivia.
+  var byteRangeAfterTrimmingTrivia: ByteSourceRange {
+    return ByteSourceRange(offset: positionAfterSkippingLeadingTrivia.utf8Offset, length: byteSizeAfterTrimmingTrivia)
+  }
+
   /// The length this node takes up spelled out in the source, excluding its
   /// leading or trailing trivia.
   var contentLength: SourceLength {