From 740199b17708556f538a2cfd4cc8f3abb9989ae8 Mon Sep 17 00:00:00 2001
From: Doug Gregor <dgregor@apple.com>
Date: Thu, 5 Sep 2024 21:47:48 -0700
Subject: [PATCH] Adopt Rintaro's improved implementation of
 comment/#sourceLocation removal

---
 .../SwiftIfConfig/ActiveSyntaxRewriter.swift  | 38 +++++++++----------
 1 file changed, 17 insertions(+), 21 deletions(-)

diff --git a/Sources/SwiftIfConfig/ActiveSyntaxRewriter.swift b/Sources/SwiftIfConfig/ActiveSyntaxRewriter.swift
index f4ac3b4c695..0293098c8f3 100644
--- a/Sources/SwiftIfConfig/ActiveSyntaxRewriter.swift
+++ b/Sources/SwiftIfConfig/ActiveSyntaxRewriter.swift
@@ -443,33 +443,29 @@ extension IfConfigDeclSyntax {
   }
 }
 
+/// Syntax visitor that writes out a syntax tree without comments or #sourceLocations.
+fileprivate class DescriptionWithoutCommentsAndSourceLocationsVisitor: SyntaxVisitor {
+  var result: String = ""
+  override func visit(_ token: TokenSyntax) -> SyntaxVisitorContinueKind {
+    token.leadingTrivia.writeWithoutComments(to: &result)
+    token.text.write(to: &result)
+    token.trailingTrivia.writeWithoutComments(to: &result)
+    return .skipChildren
+  }
+  override func visit(_ node: PoundSourceLocationSyntax) -> SyntaxVisitorContinueKind {
+    return .skipChildren
+  }
+}
+
 extension SyntaxProtocol {
   // Produce the source code for this syntax node with all of the comments
   // and #sourceLocations removed. Each comment will be replaced with either
   // a newline or a space, depending on whether the comment involved a newline.
   @_spi(Compiler)
   public var descriptionWithoutCommentsAndSourceLocations: String {
-    var result = ""
-    var skipUntilRParen = false
-    for token in tokens(viewMode: .sourceAccurate) {
-      // Skip #sourceLocation(...).
-      if token.tokenKind == .poundSourceLocation {
-        skipUntilRParen = true
-        continue
-      }
-
-      if skipUntilRParen {
-        if token.tokenKind == .rightParen {
-          skipUntilRParen = false
-        }
-        continue
-      }
-
-      token.leadingTrivia.writeWithoutComments(to: &result)
-      token.text.write(to: &result)
-      token.trailingTrivia.writeWithoutComments(to: &result)
-    }
-    return result
+    let visitor = DescriptionWithoutCommentsAndSourceLocationsVisitor(viewMode: .sourceAccurate)
+    visitor.walk(self)
+    return visitor.result
   }
 }