[Macros] Cache parsed syntax tree in compiler plugins

The compiler may send the same syntax to the plugins multiple times.
For example, 'memberAttribute' macro request contains parent nominal
decl syntax, and the compiler sends a request for each members. Parsing
it multiple times is a waste.

rdar://129624305

Author: rintaro

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()
     self.hostCapability = HostCapability()
   }
 

Sources/SwiftCompilerPluginMessageHandling/Macros.swift

@@ -56,7 +56,7 @@ extension CompilerPluginMessageHandler {
     expandingSyntax: PluginMessage.Syntax,
     lexicalContext: [PluginMessage.Syntax]?
   ) -> PluginToHostMessage {
-    let sourceManager = SourceManager()
+    let sourceManager = SourceManager(syntaxrRegistry: 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(syntaxrRegistry: 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: [Key: Syntax] = [:]
+
+  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
+  }
+
+  func clear() {
+    storage = [:]
+  }
+}
+
 /// 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(syntaxrRegistry: ParsedSyntaxRegistry) {
+    self.syntaxRegistry = syntaxrRegistry
+  }
+
   /// 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*/ })
     }