Detect single-line previews like #Preview("x") { foo() }

PrefireExecutable/Sources/PrefireCore/Previews/PreviewLoader.swift

@@ -1,72 +1,86 @@
 import Foundation
+import SwiftParser
+import SwiftSyntax
 
 /// Simplify and improve the process of locating and parsing generated preview code for Swift projects.
 enum PreviewLoader {
     enum Constants {
-        static let previewMarker = "#Preview"
+        static let previewMacro = "Preview"
         static let prefireDisableMarker = ".prefireIgnored()"
         static let prefireEnabledMarker = ".prefireEnabled()"
-        static let openingBrace: Character = "{"
-        static let closingBrace: Character = "}"
     }
 
     /// Extract the preview body using the passed content
     ///
     /// - Parameters:
     ///   - content: File content
     ///   - defaultEnabled: Whether automatic view inclusion should be allowed. Default value is true.
-    /// - Returns: An array representing the results of the macro preview without the initial `#Preview` and final `}`.
+    /// - Returns: An array representing the results of the macro preview, each starting with `#Preview`
+    ///            and ending with the closing `}` of the preview closure followed by a newline.
     static func previewBodies(from content: String, defaultEnabled: Bool) -> [String]? {
-        let lines = content.split(separator: "\n", omittingEmptySubsequences: false)
+        // Locate `#Preview` macros with SwiftSyntax rather than a line-based brace scanner.
+        // The scanner used to miscount braces that appear inside string literals, comments or
+        // raw strings (e.g. `#Preview("{braced}")`), truncating the collected body. SwiftSyntax
+        // understands the grammar, so the macro's real source range is always correct.
+        let sourceFile = Parser.parse(source: content)
+        let collector = PreviewMacroCollector()
+        collector.walk(sourceFile)
+
+        let sourceBytes = Array(content.utf8)
         var previewBodies: [String] = []
 
-        var currentBody: String = ""
-        var previewWasFound = false
-        var viewMustBeLoaded = defaultEnabled
-        var braceBalance: Int? = nil
+        for preview in collector.previews {
+            guard preview.startOffset <= preview.endOffset, preview.endOffset <= sourceBytes.count else { continue }
+            let body = String(decoding: sourceBytes[preview.startOffset..<preview.endOffset], as: UTF8.self) + "\n"
 
-        for line in lines {
-            if line.hasPrefix(Constants.previewMarker) {
-                previewWasFound = true
-                currentBody = ""
-                braceBalance = nil
+            let viewMustBeLoaded: Bool
+            if defaultEnabled {
+                viewMustBeLoaded = !body.contains(Constants.prefireDisableMarker)
+            } else {
+                viewMustBeLoaded = body.contains(Constants.prefireEnabledMarker)
             }
 
-            guard previewWasFound else { continue }
-
-            let braceChange = line.reduce(0) { (count, char) in
-                if char == Constants.openingBrace {
-                    return count + 1
-                } else if char == Constants.closingBrace {
-                    return count - 1
-                }
-                return count
+            if viewMustBeLoaded {
+                previewBodies.append(body)
             }
+        }
 
-            if braceChange != 0 || braceBalance != nil {
-                braceBalance = (braceBalance ?? 0) + braceChange
-            }
+        return previewBodies.isEmpty ? nil : previewBodies
+    }
+}
 
-            if defaultEnabled {
-                if line.contains(Constants.prefireDisableMarker) {
-                    viewMustBeLoaded = false
-                }
-            } else if !viewMustBeLoaded {
-                viewMustBeLoaded = line.contains(Constants.prefireEnabledMarker)
-            }
+/// Collects the source ranges of every `#Preview` macro in a file, in source order.
+private final class PreviewMacroCollector: SyntaxVisitor {
+    struct PreviewRange {
+        let startOffset: Int
+        let endOffset: Int
+    }
 
-            currentBody.append(String(line) + "\n")
+    private(set) var previews: [PreviewRange] = []
 
-            if braceBalance == 0 {
-                if !currentBody.isEmpty, viewMustBeLoaded {
-                    previewBodies.append(currentBody)
-                }
+    init() {
+        super.init(viewMode: .sourceAccurate)
+    }
 
-                previewWasFound = false
-                viewMustBeLoaded = defaultEnabled
-            }
-        }
+    override func visit(_ node: MacroExpansionDeclSyntax) -> SyntaxVisitorContinueKind {
+        record(macroName: node.macroName.text, pound: node.pound, node: node)
+        return .visitChildren
+    }
 
-        return previewBodies.isEmpty ? nil : previewBodies
+    override func visit(_ node: MacroExpansionExprSyntax) -> SyntaxVisitorContinueKind {
+        record(macroName: node.macroName.text, pound: node.pound, node: node)
+        return .visitChildren
+    }
+
+    /// Records the macro's range, starting at the `#` token so that leading attributes
+    /// (e.g. `@available(...)`) are excluded, matching the shape the downstream parser expects.
+    private func record(macroName: String, pound: TokenSyntax, node: some SyntaxProtocol) {
+        guard macroName == PreviewLoader.Constants.previewMacro else { return }
+        previews.append(
+            PreviewRange(
+                startOffset: pound.positionAfterSkippingLeadingTrivia.utf8Offset,
+                endOffset: node.endPositionBeforeTrailingTrivia.utf8Offset
+            )
+        )
     }
 }

PrefireExecutable/Sources/PrefireCore/Previews/RawPreviewModel.swift

@@ -29,8 +29,14 @@ extension RawPreviewModel {
     init?(from macroBody: String, filename: String) {
         guard !macroBody.isEmpty else { return nil }
 
-        var lines = macroBody.split(separator: "\n", omittingEmptySubsequences: false).dropLast(2)
-        let firstLine = lines.removeFirst()
+        // Only the first line is needed here (for displayName and traits); the body and
+        // properties are parsed from the full `macroBody` via SwiftSyntax below. Taking the
+        // first line directly avoids crashing on single-line previews like
+        // `#Preview("x") { previewFoo() }`, where `dropLast(2)` leaves an empty collection
+        // and `removeFirst()` would trap.
+        guard let firstLine = macroBody.split(separator: "\n", omittingEmptySubsequences: false).first else {
+            return nil
+        }
 
         // Define displayName by splitting the first line by "
         let parts = firstLine.split(separator: "\"")

PrefireExecutable/Tests/PrefireTests/PreviewLoaderTests.swift

@@ -49,6 +49,50 @@ class PreviewLoaderTests: XCTestCase {
         XCTAssertEqual(previews?.count, 1)
         XCTAssertEqual(previews?[0], previewRepresentations[3])
     }
+
+    func test_loadSingleLinePreviewBody() {
+        let content = "#Preview(\"SingleLine\") { previewFoo() }\n"
+        let previews = PreviewLoader.previewBodies(from: content, defaultEnabled: true)
+
+        XCTAssertEqual(previews?.count, 1)
+        XCTAssertEqual(previews?[0], content)
+    }
+
+    func test_loadSingleLineAndMultiLinePreviewBodies() {
+        let content = """
+        #Preview("SingleLine") { previewFoo() }
+        #Preview("MultiLine") {
+            Text("TestView")
+        }
+
+        """
+        let previews = PreviewLoader.previewBodies(from: content, defaultEnabled: true)
+
+        XCTAssertEqual(previews?.count, 2)
+        XCTAssertEqual(previews?[0], "#Preview(\"SingleLine\") { previewFoo() }\n")
+        XCTAssertEqual(previews?[1], "#Preview(\"MultiLine\") {\n    Text(\"TestView\")\n}\n")
+    }
+
+    func test_loadSingleLinePreviewBodyDefaultDisabled() {
+        let content = "#Preview(\"SingleLine\") { previewFoo() }\n"
+        let previews = PreviewLoader.previewBodies(from: content, defaultEnabled: false)
+
+        XCTAssertNil(previews)
+    }
+
+    func test_loadPreviewWithBraceInDisplayName() {
+        let content = """
+        #Preview("{braced}")
+        {
+            Text("TestView")
+        }
+
+        """
+        let previews = PreviewLoader.previewBodies(from: content, defaultEnabled: true)
+
+        XCTAssertEqual(previews?.count, 1)
+        XCTAssertEqual(previews?[0], "#Preview(\"{braced}\")\n{\n    Text(\"TestView\")\n}\n")
+    }
 }
 
 // MARK: - Previews

PrefireExecutable/Tests/PrefireTests/RawPreviewModelTests.swift

@@ -201,4 +201,14 @@ class RawPreviewModelTests: XCTestCase {
         XCTAssertEqual(rawPreviewModel?.displayName, "Test")
         XCTAssertEqual(rawPreviewModel?.traits, [".device"])
     }
+
+    func test_initWithSingleLineBody() {
+        let previewBody = "#Preview(\"SingleLine\", traits: .sizeThatFitsLayout) { Text(\"TestView\") }\n"
+        let rawPreviewModel = RawPreviewModel(from: previewBody, filename: "Test")
+
+        XCTAssertEqual(rawPreviewModel?.body, "Text(\"TestView\")")
+        XCTAssertEqual(rawPreviewModel?.properties, nil)
+        XCTAssertEqual(rawPreviewModel?.displayName, "SingleLine")
+        XCTAssertEqual(rawPreviewModel?.traits, [".sizeThatFitsLayout"])
+    }
 }