Merge pull request #3372 from rintaro/arena-auto-compaction

[SwiftParser] Compact incremental-parse arenas past a threshold

Author: rintaro

Sources/SwiftParser/IncrementalParseTransition.swift

@@ -52,10 +52,28 @@ public typealias ReusedNodeCallback = (_ node: Syntax) -> Void
 /// Keeps track of a previously parsed syntax tree and the source edits that
 /// occurred since it was created.
 public final class IncrementalParseTransition {
+  /// The default value for `compactArenaThreshold`.
+  ///
+  /// Bounds the retained arena count to roughly this many while keeping
+  /// compaction (a full reparse) infrequent enough not to disrupt editing
+  /// latency on large files.
+  public static let defaultCompactArenaThreshold = 64
+
   fileprivate let previousIncrementalParseResult: IncrementalParseResult
   fileprivate let edits: ConcurrentEdits
   fileprivate let reusedNodeCallback: ReusedNodeCallback?
 
+  /// When the previous tree retains at least this many arenas, the next
+  /// incremental parse is replaced by a full reparse that collapses the result
+  /// back into a single arena.
+  ///
+  /// Incremental parsing reuses subtrees from the previous tree by reference,
+  /// which keeps each reused subtree's origin arena alive. Over a long editing
+  /// session this accumulates arenas (and their memory) without bound. This
+  /// threshold caps the number of retained arenas, trading an occasional full
+  /// reparse for bounded memory. `nil` disables compaction.
+  fileprivate let compactArenaThreshold: Int?
+
   /// - Parameters:
   ///   - previousTree: The previous tree to do lookups on.
   ///   - edits: The edits that have occurred since the last parse that resulted
@@ -71,21 +89,43 @@ public final class IncrementalParseTransition {
     self.previousIncrementalParseResult = IncrementalParseResult(tree: previousTree, lookaheadRanges: lookaheadRanges)
     self.edits = edits
     self.reusedNodeCallback = reusedNodeCallback
+    self.compactArenaThreshold = Self.defaultCompactArenaThreshold
   }
 
   /// - Parameters:
   ///   - previousIncrementalParseResult: The previous incremental parse result to do lookups on.
   ///   - edits: The edits that have occurred since the last parse that resulted
   ///            in the new source that is about to be parsed.
   ///   - reusedNodeCallback: Optional closure to accept information about the re-used node. For each node that gets re-used `reusedNodeCallback` is called.
+  ///   - compactArenaThreshold: If the previous tree retains at least this many
+  ///     arenas, the next incremental parse is replaced by a full reparse to
+  ///     bound memory. Defaults to ``defaultCompactArenaThreshold``; pass `nil`
+  ///     to disable compaction.
   public init(
     previousIncrementalParseResult: IncrementalParseResult,
     edits: ConcurrentEdits,
-    reusedNodeCallback: ReusedNodeCallback? = nil
+    reusedNodeCallback: ReusedNodeCallback? = nil,
+    compactArenaThreshold: Int? = IncrementalParseTransition.defaultCompactArenaThreshold
   ) {
     self.previousIncrementalParseResult = previousIncrementalParseResult
     self.edits = edits
     self.reusedNodeCallback = reusedNodeCallback
+    self.compactArenaThreshold = compactArenaThreshold
+  }
+
+  /// Whether the next incremental parse against this transition should instead
+  /// be a full reparse to collapse accumulated arenas.
+  ///
+  /// An incremental parse retains at most `previous + 1` arenas (its own new
+  /// arena plus a subset of the previous tree's arenas), so this can be decided
+  /// from the previous tree alone, before parsing.
+  var shouldCompact: Bool {
+    // When compaction is disabled, avoid computing `retainedArenaCount` at all;
+    // it walks the retained child-arena set (O(arenas)), which is exactly the
+    // unbounded quantity we would otherwise be paying every parse. When enabled,
+    // the count is bounded by the threshold, so the walk is O(threshold).
+    guard let compactArenaThreshold else { return false }
+    return previousIncrementalParseResult.tree.raw.arena.retainedArenaCount + 1 > compactArenaThreshold
   }
 }
 

Sources/SwiftParser/ParseSourceFile.swift

@@ -135,6 +135,9 @@ extension Parser {
     source: String,
     parseTransition: IncrementalParseTransition?
   ) -> IncrementalParseResult {
+    // Drop the transition (forcing a full reparse) when the previous tree is
+    // fragmented enough that this parse would reach its compaction threshold.
+    let parseTransition = (parseTransition?.shouldCompact ?? false) ? nil : parseTransition
     return withParser(
       source: source,
       maximumNestingLevel: nil,
@@ -153,6 +156,7 @@ extension Parser {
     maximumNestingLevel: Int? = nil,
     parseTransition: IncrementalParseTransition?
   ) -> IncrementalParseResult {
+    let parseTransition = (parseTransition?.shouldCompact ?? false) ? nil : parseTransition
     return withParser(
       source: source,
       maximumNestingLevel: maximumNestingLevel,

Sources/SwiftSyntax/Raw/RawSyntaxArena.swift

@@ -177,6 +177,19 @@ public class RawSyntaxArena {
   func contains(text: SyntaxText) -> Bool {
     return (text.isEmpty || allocator.contains(address: text.baseAddress!))
   }
+
+  /// Number of distinct arenas (this one plus retained child arenas) that the
+  /// tree rooted in this arena keeps alive.
+  public var retainedArenaCount: Int {
+    var seen: Set<RawSyntaxArenaRef> = [RawSyntaxArenaRef(self)]
+    var stack: [RawSyntaxArena] = [self]
+    while let arena = stack.popLast() {
+      for childRef in arena.childRefs where seen.insert(childRef).inserted {
+        stack.append(childRef.value)
+      }
+    }
+    return seen.count
+  }
 }
 
 /// RawSyntaxArena for parsing.
@@ -229,6 +242,11 @@ public struct RetainedRawSyntaxArena: @unchecked Sendable {
   fileprivate func arenaRef() -> RawSyntaxArenaRef {
     return RawSyntaxArenaRef(arena)
   }
+
+  /// Number of arenas (self + retained children) kept alive by this tree.
+  public var retainedArenaCount: Int {
+    arena.retainedArenaCount
+  }
 }
 
 /// Unsafely unowned reference to ``RawSyntaxArena``. The user is responsible to
@@ -245,7 +263,7 @@ struct RawSyntaxArenaRef: Hashable, @unchecked Sendable {
   }
 
   /// Returns the ``RawSyntaxArena``
-  private var value: RawSyntaxArena {
+  fileprivate var value: RawSyntaxArena {
     self._value.takeUnretainedValue()
   }
 

Tests/SwiftParserTest/IncrementalParsingTests.swift

@@ -11,7 +11,7 @@
 //===----------------------------------------------------------------------===//
 
 import SwiftParser
-import SwiftSyntax
+@_spi(RawSyntax) import SwiftSyntax
 import XCTest
 import _SwiftSyntaxTestSupport
 
@@ -466,4 +466,115 @@ class IncrementalParsingTests: ParserTestCase {
       ]
     )
   }
+
+  // MARK: - Arena compaction
+
+  /// A source file of `count` distinct top-level functions. Editing one function
+  /// reuses the others from earlier arenas, so editing distinct functions in
+  /// turn accumulates retained arenas.
+  private func manyFunctions(_ count: Int) -> String {
+    return (0..<count).map { "func f\($0)() -> Int { return \($0) }\n" }.joined()
+  }
+
+  /// Insert `text` at the start of `needle` in `source`, returning the new
+  /// source and the corresponding edit.
+  private func insert(
+    _ text: String,
+    before needle: String,
+    in source: String
+  ) -> (newSource: String, edit: SourceEdit) {
+    let range = source.range(of: needle)!
+    let offset = source.utf8.distance(from: source.utf8.startIndex, to: range.lowerBound.samePosition(in: source.utf8)!)
+    var utf8 = Array(source.utf8)
+    utf8.insert(contentsOf: Array(text.utf8), at: offset)
+    let edit = SourceEdit(
+      range: AbsolutePosition(utf8Offset: offset)..<AbsolutePosition(utf8Offset: offset),
+      replacement: text
+    )
+    return (String(decoding: utf8, as: UTF8.self), edit)
+  }
+
+  /// Edits distinct functions in turn and returns the peak retained arena count
+  /// observed across the edit sequence. Asserts every intermediate tree
+  /// round-trips.
+  private func peakRetainedArenas(functionCount: Int, edits: Int, compactArenaThreshold: Int?) throws -> Int {
+    var source = manyFunctions(functionCount)
+    var result = Parser.parseIncrementally(source: source, parseTransition: nil)
+    var peak = result.tree.raw.arena.retainedArenaCount
+    for i in 0..<edits {
+      let (newSource, edit) = insert("  ", before: "return \(i % functionCount) ", in: source)
+      let transition = IncrementalParseTransition(
+        previousIncrementalParseResult: result,
+        edits: try ConcurrentEdits(concurrent: [edit]),
+        compactArenaThreshold: compactArenaThreshold
+      )
+      result = Parser.parseIncrementally(source: newSource, parseTransition: transition)
+      XCTAssertEqual(result.tree.description, newSource, "tree did not round-trip after edit \(i)")
+      peak = max(peak, result.tree.raw.arena.retainedArenaCount)
+      source = newSource
+    }
+    return peak
+  }
+
+  /// With compaction disabled, editing distinct functions accumulates arenas
+  /// well beyond a small bound — this validates that the compaction test below
+  /// is actually exercising arena accumulation.
+  public func testArenaCountAccumulatesWithoutCompaction() throws {
+    let peak = try peakRetainedArenas(functionCount: 16, edits: 16, compactArenaThreshold: nil)
+    // Each edit reuses the other functions from earlier arenas, so the retained
+    // count grows by roughly one per edit (peak here is near `1 + edits` ≈ 17).
+    // This is only a control asserting that accumulation happens at all, so use
+    // a loose lower bound: comfortably above the compaction-bounded regime (the
+    // test below caps at 4) yet well under the true peak, so it stays robust to
+    // small variations in how nodes are pinned across arenas.
+    XCTAssertGreaterThan(peak, 8, "expected arenas to accumulate without compaction")
+  }
+
+  /// With a low threshold, compaction keeps the retained arena count bounded by
+  /// the threshold across a long edit sequence, while every tree still
+  /// round-trips.
+  public func testArenaCompactionBoundsRetainedArenaCount() throws {
+    let threshold = 4
+    let peak = try peakRetainedArenas(functionCount: 16, edits: 32, compactArenaThreshold: threshold)
+    XCTAssertLessThanOrEqual(peak, threshold, "compaction should keep retained arenas at or below the threshold")
+  }
+
+  /// A tree produced by a compaction (full reparse) can still be used as the
+  /// basis for a subsequent incremental parse that reuses nodes.
+  public func testIncrementalParseAfterCompaction() throws {
+    let source = manyFunctions(16)
+    var result = Parser.parseIncrementally(source: source, parseTransition: nil)
+
+    // Force a compaction on the next parse with threshold 1 (any incremental
+    // parse's previous tree retains >= 1 arena, so `previous + 1 > 1` always
+    // holds and forces a compaction).
+    let (afterFirst, edit1) = insert("  ", before: "return 0 ", in: source)
+    result = Parser.parseIncrementally(
+      source: afterFirst,
+      parseTransition: IncrementalParseTransition(
+        previousIncrementalParseResult: result,
+        edits: try ConcurrentEdits(concurrent: [edit1]),
+        compactArenaThreshold: 1
+      )
+    )
+    XCTAssertEqual(result.tree.raw.arena.retainedArenaCount, 1, "compaction should collapse to a single arena")
+    XCTAssertEqual(result.tree.description, afterFirst)
+
+    // The compacted tree drives a normal incremental parse that reuses nodes.
+    var reusedCodeBlockItems = 0
+    let (afterSecond, edit2) = insert("  ", before: "return 5 ", in: afterFirst)
+    let result2 = Parser.parseIncrementally(
+      source: afterSecond,
+      parseTransition: IncrementalParseTransition(
+        previousIncrementalParseResult: result,
+        edits: try ConcurrentEdits(concurrent: [edit2]),
+        reusedNodeCallback: { node in
+          if node.kind == .codeBlockItem { reusedCodeBlockItems += 1 }
+        },
+        compactArenaThreshold: nil
+      )
+    )
+    XCTAssertEqual(result2.tree.description, afterSecond)
+    XCTAssertGreaterThan(reusedCodeBlockItems, 0, "incremental parse after compaction should reuse nodes")
+  }
 }