[Macros] Add parentContext(of:) operation to MacroExpansionContext.

Extend `MacroExpansionContext` with a new operation `parentContext(of:)` that
provides a "pruned" version of the innermost enclosing syntax node for a
given node. This allows us to establish contextual information for the
syntax nodes that are passed to a macro expansion, without exposing
information about unrelated parts of the source file.

Author: DougGregor

Sources/SwiftSyntaxMacros/BasicMacroExpansionContext.swift

@@ -65,6 +65,8 @@ public class BasicMacroExpansionContext {
   /// Used in conjunction with `expansionDiscriminator`.
   private var uniqueNames: [String: Int] = [:]
 
+  // The parent contexts of any nodes that have been detached.
+  private var parentContexts: [Syntax: Syntax] = [:]
 }
 
 extension BasicMacroExpansionContext {
@@ -86,6 +88,15 @@ extension BasicMacroExpansionContext {
       addDisconnected(detached, at: node.position, in: rootSourceFile)
     }
 
+    // Record macro parent contexts.
+    do {
+      var childNode = Syntax(detached)
+      for parentNode in node.allMacroParentContexts() {
+        parentContexts[childNode] = parentNode
+        childNode = parentNode
+      }
+    }
+
     return detached
   }
 }
@@ -190,4 +201,12 @@ extension BasicMacroExpansionContext: MacroExpansionContext {
     let converter = SourceLocationConverter(file: fileName, tree: rootSourceFile)
     return converter.location(for: rawPosition.advanced(by: offsetAdjustment))
   }
+
+  public func parentContext<Node: SyntaxProtocol>(of node: Node) -> Syntax? {
+    if let context = node.allMacroParentContexts().first {
+      return context
+    }
+
+    return parentContexts[Syntax(node)]
+  }
 }

Sources/SwiftSyntaxMacros/CMakeLists.txt

@@ -26,6 +26,7 @@ add_swift_host_library(SwiftSyntaxMacros
   MacroReplacement.swift
   MacroSystem.swift
   Syntax+MacroEvaluation.swift
+  Syntax+ParentContext.swift
 )
 
 target_link_libraries(SwiftSyntaxMacros PUBLIC

Sources/SwiftSyntaxMacros/MacroExpansionContext.swift

@@ -78,6 +78,19 @@ public protocol MacroExpansionContext: AnyObject {
     at position: PositionInSyntaxNode,
     filePathMode: SourceLocationFilePathMode
   ) -> AbstractSourceLocation?
+
+  /// Determine the parent context of the given syntax node.
+  ///
+  /// For a syntax node that is part of the syntax provided to a macro
+  /// expansion, find the innermost enclosing context node. A context
+  /// node is an entity such as a function declaration, type declaration,
+  /// or extension that can have other entities nested inside it.
+  /// The resulting context node will have any information about nested
+  /// entities removed from it, to prevent macro expansion operations from
+  /// seeing unrelated code within the program. For more information
+  /// about the identification and pruning of parent contexts, see
+  /// `SyntaxProtocol.asMacroParentContext`.
+  func parentContext<Node: SyntaxProtocol>(of node: Node) -> Syntax?
 }
 
 extension MacroExpansionContext {
@@ -220,6 +233,23 @@ extension MacroExpansionContext {
   }
 }
 
+extension MacroExpansionContext {
+  /// Determine the parent context of the given syntax node.
+  ///
+  /// For a syntax node that is part of the syntax provided to a macro
+  /// expansion, find the innermost enclosing context node. A context
+  /// node is an entity such as a function declaration, type declaration,
+  /// or extension that can have other entities nested inside it.
+  /// The resulting context node will have any information about nested
+  /// entities removed from it, to prevent macro expansion operations from
+  /// seeing unrelated code within the program. For more information
+  /// about the identification and pruning of parent contexts, see
+  /// `SyntaxProtocol.asMacroParentContext`.
+  public func parentContext<Node: SyntaxProtocol>(of node: Node) -> Syntax? {
+    return nil
+  }
+}
+
 /// Describe the position within a syntax node that can be used to compute
 /// source locations.
 public enum PositionInSyntaxNode {

Sources/SwiftSyntaxMacros/Syntax+ParentContext.swift

@@ -0,0 +1,63 @@
+//===----------------------------------------------------------------------===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2014 - 2023 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
+//
+//===----------------------------------------------------------------------===//
+
+import SwiftSyntax
+import SwiftSyntaxBuilder
+
+extension SyntaxProtocol {
+  /// If this is a
+  public func asMacroParentContext() -> Syntax? {
+    switch Syntax(self).asProtocol(SyntaxProtocol.self) {
+    // Functions
+    case var function as HasTrailingOptionalCodeBlock:
+      function.body = nil
+      return Syntax(function).detach()
+
+    // Nominal types and extensions.
+    case var typeOrExtension as HasTrailingMemberDeclBlock:
+      typeOrExtension.members = MemberDeclBlockSyntax(members: MemberDeclListSyntax())
+      return Syntax(typeOrExtension).detach()
+
+    case var subscriptDecl as SubscriptDeclSyntax:
+      subscriptDecl.accessor = nil
+      return Syntax(subscriptDecl).detach()
+
+    case is EnumCaseElementSyntax:
+      return Syntax(self).detach()
+
+    case var patternBinding as PatternBindingSyntax:
+      patternBinding.accessor = nil
+      patternBinding.initializer = nil
+      return Syntax(patternBinding).detach()
+
+    default:
+      return nil
+    }
+  }
+
+  /// Return an array of enclosing parent contexts for the purpose of macros,
+  /// from the innermost enclosing parent context (first in the array) to the
+  /// outermost.
+  public func allMacroParentContexts() -> [Syntax] {
+    var parentContexts: [Syntax] = []
+    var currentNode = Syntax(self)
+    while let parentNode = currentNode.parent {
+      if let parentContext = parentNode.asMacroParentContext() {
+        parentContexts.append(parentContext)
+      }
+
+      currentNode = parentNode
+    }
+
+    return parentContexts
+  }
+}

Tests/SwiftSyntaxMacrosTest/MacroSystemTests.swift

@@ -14,7 +14,7 @@ import SwiftDiagnostics
 import SwiftParser
 import SwiftSyntax
 import SwiftSyntaxBuilder
-import SwiftSyntaxMacros
+@_spi(Testing) import SwiftSyntaxMacros
 import _SwiftSyntaxTestSupport
 import XCTest
 
@@ -145,6 +145,125 @@ public struct FileIDMacro: ExpressionMacro {
   }
 }
 
+extension PatternBindingSyntax {
+  /// When the variable is declaring a single binding, produce the name of
+  /// that binding.
+  fileprivate var singleBindingName: String? {
+    if let identifierPattern = pattern.as(IdentifierPatternSyntax.self) {
+      return identifierPattern.identifier.text
+    }
+
+    return nil
+  }
+}
+
+extension SyntaxProtocol {
+  /// Form a function name.
+  private func formFunctionName(
+    _ baseName: String, _ parameters: ParameterClauseSyntax?,
+    isSubscript: Bool = false
+  ) -> String {
+    let argumentNames: [String] = parameters?.parameterList.map { param in
+      let argumentLabelText = param.argumentName?.text ?? "_"
+      return argumentLabelText + ":"
+    } ?? []
+
+    return "\(baseName)(\(argumentNames.joined(separator: "")))"
+  }
+
+  /// Form the #function name for the given node.
+  fileprivate func functionName<Context: MacroExpansionContext>(
+    in context: Context
+  ) -> String? {
+    // Declarations with parameters.
+    // FIXME: Can we abstract over these?
+    if let function = self.as(FunctionDeclSyntax.self) {
+      return formFunctionName(
+        function.identifier.text, function.signature.input
+      )
+    }
+
+    if let initializer = self.as(InitializerDeclSyntax.self) {
+      return formFunctionName("init", initializer.signature.input)
+    }
+
+    if let subscriptDecl = self.as(SubscriptDeclSyntax.self) {
+      return formFunctionName(
+        "subscript", subscriptDecl.indices, isSubscript: true
+      )
+    }
+
+    if let enumCase = self.as(EnumCaseElementSyntax.self) {
+      guard let associatedValue = enumCase.associatedValue else {
+        return enumCase.identifier.text
+      }
+
+      let argumentNames = associatedValue.parameterList.map { param in
+        guard let firstName = param.firstName else {
+          return "_:"
+        }
+
+        return firstName.text + ":"
+      }.joined()
+
+      return "\(enumCase.identifier.text)(\(argumentNames))"
+    }
+
+    // Accessors use their enclosing context, i.e., a subscript or pattern
+    // binding.
+    if self.is(AccessorDeclSyntax.self) {
+      guard let parentContext = context.parentContext(of: self) else {
+        return nil
+      }
+
+      return parentContext.functionName(in: context)
+    }
+
+    // All declarations with identifiers.
+    if let identified = self.asProtocol(IdentifiedDeclSyntax.self) {
+      return identified.identifier.text
+    }
+
+    // Extensions
+    if let extensionDecl = self.as(ExtensionDeclSyntax.self) {
+      // FIXME: It would be nice to be able to switch on type syntax...
+      let extendedType = extensionDecl.extendedType
+      if let simple = extendedType.as(SimpleTypeIdentifierSyntax.self) {
+        return simple.name.text
+      }
+
+      if let member = extendedType.as(MemberTypeIdentifierSyntax.self) {
+        return member.name.text
+      }
+    }
+
+    // Pattern bindings.
+    if let patternBinding = self.as(PatternBindingSyntax.self),
+       let singleVarName = patternBinding.singleBindingName {
+      return singleVarName
+    }
+
+    return nil
+  }
+}
+
+public struct FunctionMacro: ExpressionMacro {
+  public static func expansion<
+    Node: FreestandingMacroExpansionSyntax,
+    Context: MacroExpansionContext
+  >(
+    of node: Node,
+    in context: Context
+  ) -> ExprSyntax {
+    guard let parentContext = context.parentContext(of: node),
+          let name = parentContext.functionName(in: context) else {
+      return #""<unknown>""#
+    }
+
+    return ExprSyntax("\(literal: name)").with(\.leadingTrivia, node.leadingTrivia)
+  }
+}
+
 /// Macro whose only purpose is to ensure that we cannot see "out" of the
 /// macro expansion syntax node we were given.
 struct CheckContextIndependenceMacro: ExpressionMacro {
@@ -707,6 +826,7 @@ public let testMacros: [String: Macro.Type] = [
   "colorLiteral": ColorLiteralMacro.self,
   "column": ColumnMacro.self,
   "fileID": FileIDMacro.self,
+  "function": FunctionMacro.self,
   "imageLiteral": ImageLiteralMacro.self,
   "stringify": StringifyMacro.self,
   "myError": ErrorMacro.self,
@@ -769,6 +889,70 @@ final class MacroSystemTests: XCTestCase {
     )
   }
 
+  func testPoundFunction() {
+    assertMacroExpansion(
+      macros: testMacros,
+      """
+      func f(a: Int, _: Double, c: Int) {
+        print(#function)
+      }
+
+      struct X {
+        var computed: String {
+          get {
+            #function
+          }
+        }
+
+        init(from: String) {
+          #function
+        }
+
+        subscript(a: Int) -> String {
+          #function
+        }
+
+        subscript(a a: Int) -> String {
+          #function
+        }
+      }
+
+      extension A {
+        static var staticProp: String = #function
+      }
+      """,
+      """
+      func f(a: Int, _: Double, c: Int) {
+        print("f(a:_:c:)")
+      }
+
+      struct X {
+        var computed: String {
+          get {
+            "computed"
+          }
+        }
+
+        init(from: String) {
+          "init(from:)"
+        }
+
+        subscript(a: Int) -> String {
+          "subscript(_:)"
+        }
+
+        subscript(a a: Int) -> String {
+          "subscript(a:)"
+        }
+      }
+
+      extension A {
+        static var staticProp: String = "staticProp"
+      }
+      """
+    )
+  }
+
   func testContextUniqueLocalNames() {
     let context = BasicMacroExpansionContext()