[Macros] Cache parsed syntax tree in compiler plugins

Sources/SwiftCompilerPluginMessageHandling/CMakeLists.txt

@@ -9,6 +9,7 @@
 add_swift_syntax_library(SwiftCompilerPluginMessageHandling
   CompilerPluginMessageHandler.swift
   Diagnostics.swift
+  LRUCache.swift
   Macros.swift
   PluginMacroExpansionContext.swift
   PluginMessageCompatibility.swift

Sources/SwiftCompilerPluginMessageHandling/CompilerPluginMessageHandler.swift

@@ -112,11 +112,15 @@ public class CompilerPluginMessageHandler<Provider: PluginProvider> {
   /// Object to provide actual plugin functions.
   let provider: Provider
 
+  /// Syntax registry shared between multiple requests.
+  let syntaxRegistry: ParsedSyntaxRegistry
+
   /// Plugin host capability
   var hostCapability: HostCapability
 
   public init(provider: Provider) {
     self.provider = provider
+    self.syntaxRegistry = ParsedSyntaxRegistry(cacheCapacity: 16)
     self.hostCapability = HostCapability()
   }
 

Sources/SwiftCompilerPluginMessageHandling/LRUCache.swift

@@ -0,0 +1,116 @@
+//===----------------------------------------------------------------------===//
+//
+// This source file is part of the Swift open source project
+//
+// Copyright (c) 2024 Apple Inc. and the Swift project authors
+// Licensed under Apache License v2.0 with Runtime Library Exception
+//
+// See http://swift.org/LICENSE.txt for license information
+// See http://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
+//
+//===----------------------------------------------------------------------===//
+
+/// Simple LRU cache.
+@_spi(Testing)
+public class LRUCache<Key: Hashable, Value> {
+  private class _Node {
+    unowned var prev: _Node? = nil
+    unowned var next: _Node? = nil
+
+    let key: Key
+    var value: Value
+
+    init(key: Key, value: Value) {
+      self.key = key
+      self.value = value
+    }
+  }
+
+  private var table: [Key: _Node]
+
+  // Double linked list
+  private unowned var head: _Node?
+  private unowned var tail: _Node?
+
+  public let capacity: Int
+
+  public init(capacity: Int) {
+    self.table = [:]
+    self.head = nil
+    self.tail = nil
+    self.capacity = capacity
+  }
+
+  public var count: Int {
+    return table.count
+  }
+
+  public subscript(key: Key) -> Value? {
+    get {
+      guard let node = table[key] else {
+        return nil
+      }
+      moveToHead(node: node)
+      return node.value
+    }
+
+    set {
+      switch (table[key], newValue) {
+      case let (nil, newValue?):  // create.
+        self.ensureCapacityForNewValue()
+        let node = _Node(key: key, value: newValue)
+        addToHead(node: node)
+        table[key] = node
+
+      case let (node?, newValue?):  // update.
+        moveToHead(node: node)
+        node.value = newValue
+
+      case let (node?, nil):  // delete.
+        remove(node: node)
+        table[key] = nil
+
+      case (nil, nil):  // no-op.
+        break
+      }
+    }
+  }
+
+  private func ensureCapacityForNewValue() {
+    while self.table.count >= self.capacity, let tail = self.tail {
+      remove(node: tail)
+      table[tail.key] = nil
+    }
+  }
+
+  private func moveToHead(node: _Node) {
+    if node === self.head {
+      return
+    }
+    remove(node: node)
+    addToHead(node: node)
+  }
+
+  private func addToHead(node: _Node) {
+    node.next = self.head
+    node.next?.prev = node
+    node.prev = nil
+    self.head = node
+    if self.tail == nil {
+      self.tail = node
+    }
+  }
+
+  private func remove(node: _Node) {
+    node.next?.prev = node.prev
+    node.prev?.next = node.next
+    if node === self.head {
+      self.head = node.next
+    }
+    if node === self.tail {
+      self.tail = node.prev
+    }
+    node.prev = nil
+    node.next = nil
+  }
+}

Sources/SwiftCompilerPluginMessageHandling/Macros.swift

@@ -56,7 +56,7 @@ extension CompilerPluginMessageHandler {
     expandingSyntax: PluginMessage.Syntax,
     lexicalContext: [PluginMessage.Syntax]?
   ) -> PluginToHostMessage {
-    let sourceManager = SourceManager()
+    let sourceManager = SourceManager(syntaxRegistry: syntaxRegistry)
     let syntax = sourceManager.add(expandingSyntax, foldingWith: .standardOperators)
 
     let context = PluginMacroExpansionContext(
@@ -120,7 +120,7 @@ extension CompilerPluginMessageHandler {
     conformanceListSyntax: PluginMessage.Syntax?,
     lexicalContext: [PluginMessage.Syntax]?
   ) -> PluginToHostMessage {
-    let sourceManager = SourceManager()
+    let sourceManager = SourceManager(syntaxRegistry: syntaxRegistry)
     let attributeNode = sourceManager.add(
       attributeSyntax,
       foldingWith: .standardOperators

Sources/SwiftCompilerPluginMessageHandling/PluginMacroExpansionContext.swift

@@ -24,6 +24,49 @@ import SwiftSyntax
 import SwiftSyntaxMacros
 #endif
 
+/// Caching parser for PluginMessage.Syntax
+class ParsedSyntaxRegistry {
+  struct Key: Hashable {
+    let source: String
+    let kind: PluginMessage.Syntax.Kind
+  }
+
+  private var storage: LRUCache<Key, Syntax>
+
+  init(cacheCapacity: Int) {
+    self.storage = LRUCache(capacity: cacheCapacity)
+  }
+
+  private func parse(source: String, kind: PluginMessage.Syntax.Kind) -> Syntax {
+    var parser = Parser(source)
+    switch kind {
+    case .declaration:
+      return Syntax(DeclSyntax.parse(from: &parser))
+    case .statement:
+      return Syntax(StmtSyntax.parse(from: &parser))
+    case .expression:
+      return Syntax(ExprSyntax.parse(from: &parser))
+    case .type:
+      return Syntax(TypeSyntax.parse(from: &parser))
+    case .pattern:
+      return Syntax(PatternSyntax.parse(from: &parser))
+    case .attribute:
+      return Syntax(AttributeSyntax.parse(from: &parser))
+    }
+  }
+
+  func get(source: String, kind: PluginMessage.Syntax.Kind) -> Syntax {
+    let key = Key(source: source, kind: kind)
+    if let cached = storage[key] {
+      return cached
+    }
+
+    let node = parse(source: source, kind: kind)
+    storage[key] = node
+    return node
+  }
+}
+
 /// Manages known source code combined with their filename/fileID. This can be
 /// used to get line/column from a syntax node in the managed source code.
 class SourceManager {
@@ -67,32 +110,24 @@ class SourceManager {
     var endUTF8Offset: Int
   }
 
+  /// Caching syntax parser.
+  private let syntaxRegistry: ParsedSyntaxRegistry
+
   /// Syntax added by `add(_:)` method. Keyed by the `id` of the node.
   private var knownSourceSyntax: [Syntax.ID: KnownSourceSyntax] = [:]
 
+  init(syntaxRegistry: ParsedSyntaxRegistry) {
+    self.syntaxRegistry = syntaxRegistry
+  }
+
   /// Convert syntax information to a ``Syntax`` node. The location informations
   /// are cached in the source manager to provide `location(of:)` et al.
   func add(
     _ syntaxInfo: PluginMessage.Syntax,
     foldingWith operatorTable: OperatorTable? = nil
   ) -> Syntax {
 
-    var node: Syntax
-    var parser = Parser(syntaxInfo.source)
-    switch syntaxInfo.kind {
-    case .declaration:
-      node = Syntax(DeclSyntax.parse(from: &parser))
-    case .statement:
-      node = Syntax(StmtSyntax.parse(from: &parser))
-    case .expression:
-      node = Syntax(ExprSyntax.parse(from: &parser))
-    case .type:
-      node = Syntax(TypeSyntax.parse(from: &parser))
-    case .pattern:
-      node = Syntax(PatternSyntax.parse(from: &parser))
-    case .attribute:
-      node = Syntax(AttributeSyntax.parse(from: &parser))
-    }
+    var node = syntaxRegistry.get(source: syntaxInfo.source, kind: syntaxInfo.kind)
     if let operatorTable {
       node = operatorTable.foldAll(node, errorHandler: { _ in /*ignore*/ })
     }

Tests/SwiftCompilerPluginTest/LRUCacheTests.swift

@@ -0,0 +1,37 @@
+//===----------------------------------------------------------------------===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2024 Apple Inc. and the Swift project authors
+// Licensed under Apache License v2.0 with Runtime Library Exception
+//
+// See https://swift.org/LICENSE.txt for license information
+// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
+//
+//===----------------------------------------------------------------------===//
+
+@_spi(Testing) import SwiftCompilerPluginMessageHandling
+import XCTest
+
+final class LRUCacheTests: XCTestCase {
+  func testBasic() {
+    let cache = LRUCache<String, Int>(capacity: 2)
+    cache["foo"] = 0
+    cache["bar"] = 1
+    XCTAssertEqual(cache["foo"], 0)
+    cache["baz"] = 2
+    XCTAssertEqual(cache["foo"], 0)
+    XCTAssertEqual(cache["bar"], nil)
+    XCTAssertEqual(cache["baz"], 2)
+    XCTAssertEqual(cache.count, 2)
+
+    cache["qux"] = nil
+    cache["baz"] = nil
+    cache["foo"] = 10
+    XCTAssertEqual(cache["foo"], 10)
+    XCTAssertEqual(cache["bar"], nil)
+    XCTAssertEqual(cache["baz"], nil)
+    XCTAssertEqual(cache["qux"], nil)
+    XCTAssertEqual(cache.count, 1)
+  }
+}