[Macros] Store parsed syntax tres in a LRU cache

Author: rintaro

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

@@ -120,7 +120,7 @@ public class CompilerPluginMessageHandler<Provider: PluginProvider> {
 
   public init(provider: Provider) {
     self.provider = provider
-    self.syntaxRegistry = ParsedSyntaxRegistry()
+    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/PluginMacroExpansionContext.swift

@@ -31,7 +31,11 @@ class ParsedSyntaxRegistry {
     let kind: PluginMessage.Syntax.Kind
   }
 
-  private var storage: [Key: Syntax] = [:]
+  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)
@@ -61,10 +65,6 @@ class ParsedSyntaxRegistry {
     storage[key] = node
     return node
   }
-
-  func clear() {
-    storage = [:]
-  }
 }
 
 /// Manages known source code combined with their filename/fileID. This can be

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)
+  }
+}