From 2c183afff78081d12057372a0399d199c35b7c4e Mon Sep 17 00:00:00 2001
From: Rob Amos <Robert.Amos@anz.com>
Date: Fri, 13 Dec 2024 22:26:16 +1100
Subject: [PATCH] Added correct detection of conformance scopes when inside a
 public extension

---
 Sources/VexilMacros/FlagContainerMacro.swift  | 24 ++++--
 .../EquatableFlagContainerMacroTests.swift    | 78 +++++++++++++++++++
 2 files changed, 96 insertions(+), 6 deletions(-)

diff --git a/Sources/VexilMacros/FlagContainerMacro.swift b/Sources/VexilMacros/FlagContainerMacro.swift
index abbc9e1..9b21dae 100644
--- a/Sources/VexilMacros/FlagContainerMacro.swift
+++ b/Sources/VexilMacros/FlagContainerMacro.swift
@@ -25,7 +25,13 @@ extension FlagContainerMacro: MemberMacro {
         providingMembersOf declaration: some DeclGroupSyntax,
         in context: some MacroExpansionContext
     ) throws -> [DeclSyntax] {
-        try [
+        // If the declaration doesn't have any scopes attached we might be inheriting scopes from a public extension
+        var scopes = declaration.modifiers.scopeSyntax
+        if scopes.isEmpty, let parent = context.lexicalContext.first?.as(ExtensionDeclSyntax.self) {
+            scopes = parent.modifiers.scopeSyntax
+        }
+
+        return try [
 
             // Properties
 
@@ -43,7 +49,7 @@ extension FlagContainerMacro: MemberMacro {
                     ExprSyntax("self._flagKeyPath = _flagKeyPath")
                     ExprSyntax("self._flagLookup = _flagLookup")
                 }
-                .with(\.modifiers, declaration.modifiers.scopeSyntax)
+                .with(\.modifiers, scopes)
             ),
 
         ]
@@ -79,10 +85,16 @@ extension FlagContainerMacro: ExtensionMacro {
 
         // Check that conformance doesn't already exist, or that we are inside a unit test.
         // The latter is a workaround for https://github.com/apple/swift-syntax/issues/2031
-        guard shouldGenerateConformance.flagContainer  else {
+        guard shouldGenerateConformance.flagContainer else {
             return []
         }
 
+        // If the declaration doesn't have any scopes attached we might be inheriting scopes from a public extension
+        var scopes = declaration.modifiers.scopeSyntax
+        if scopes.isEmpty, let parent = context.lexicalContext.first?.as(ExtensionDeclSyntax.self) {
+            scopes = parent.modifiers.scopeSyntax
+        }
+
         var decls = try [
             ExtensionDeclSyntax(
                 extendedType: type,
@@ -102,7 +114,7 @@ extension FlagContainerMacro: ExtensionMacro {
                     }
                     "visitor.endContainer(keyPath: _flagKeyPath)"
                 }
-                .with(\.modifiers, declaration.modifiers.scopeSyntax)
+                .with(\.modifiers, scopes)
 
                 // Flag Key Paths
 
@@ -135,7 +147,7 @@ extension FlagContainerMacro: ExtensionMacro {
                         "[:]"
                     }
                 }
-                .with(\.modifiers, declaration.modifiers.scopeSyntax)
+                .with(\.modifiers, scopes)
 
             },
         ]
@@ -161,7 +173,7 @@ extension FlagContainerMacro: ExtensionMacro {
                                 ExprSyntax("lhs.\(lastBinding) == rhs.\(lastBinding)")
                             }
                         }
-                        .with(\.modifiers, Array(declaration.modifiers.scopeSyntax) + [ DeclModifierSyntax(name: .keyword(.static)) ])
+                        .with(\.modifiers, Array(scopes) + [ DeclModifierSyntax(name: .keyword(.static)) ])
                     }
                 },
             ]
diff --git a/Tests/VexilMacroTests/EquatableFlagContainerMacroTests.swift b/Tests/VexilMacroTests/EquatableFlagContainerMacroTests.swift
index 98c0ccd..b566cf3 100644
--- a/Tests/VexilMacroTests/EquatableFlagContainerMacroTests.swift
+++ b/Tests/VexilMacroTests/EquatableFlagContainerMacroTests.swift
@@ -233,6 +233,84 @@ final class EquatableFlagContainerMacroTests: XCTestCase {
         )
     }
 
+    func testExpandsPublicExtension() throws {
+        assertMacroExpansion(
+            """
+            public extension SomeContainer {
+                @FlagContainer
+                struct TestFlags {
+                    @Flag(default: false, description: "Some Flag")
+                    var someFlag: Bool
+                }
+            }
+            """,
+            expandedSource:
+            """
+            public extension SomeContainer {
+                struct TestFlags {
+                    var someFlag: Bool {
+                        get {
+                            _flagLookup.value(for: _flagKeyPath.append(.automatic("some-flag"))) ?? false
+                        }
+                    }
+
+                    var $someFlag: FlagWigwag<Bool> {
+                        FlagWigwag(
+                            keyPath: _flagKeyPath.append(.automatic("some-flag")),
+                            name: nil,
+                            defaultValue: false,
+                            description: "Some Flag",
+                            displayOption: .default,
+                            lookup: _flagLookup
+                        )
+                    }
+
+                    fileprivate let _flagKeyPath: FlagKeyPath
+
+                    fileprivate let _flagLookup: any FlagLookup
+
+                    public init(_flagKeyPath: FlagKeyPath, _flagLookup: any FlagLookup) {
+                        self._flagKeyPath = _flagKeyPath
+                        self._flagLookup = _flagLookup
+                    }
+                }
+            }
+
+            extension SomeContainer.TestFlags: FlagContainer {
+                public func walk(visitor: any FlagVisitor) {
+                    visitor.beginContainer(keyPath: _flagKeyPath, containerType: SomeContainer.TestFlags.self)
+                    visitor.visitFlag(
+                        keyPath: _flagKeyPath.append(.automatic("some-flag")),
+                        value: { [self] in
+                            _flagLookup.value(for: _flagKeyPath.append(.automatic("some-flag")))
+                        },
+                        defaultValue: false,
+                        wigwag: { [self] in
+                            $someFlag
+                        }
+                    )
+                    visitor.endContainer(keyPath: _flagKeyPath)
+                }
+                public var _allFlagKeyPaths: [PartialKeyPath<SomeContainer.TestFlags>: FlagKeyPath] {
+                    [
+                        \\SomeContainer.TestFlags.someFlag: _flagKeyPath.append(.automatic("some-flag")),
+                    ]
+                }
+            }
+
+            extension SomeContainer.TestFlags: Equatable {
+                public static func ==(lhs: SomeContainer.TestFlags, rhs: SomeContainer.TestFlags) -> Bool {
+                    lhs.someFlag == rhs.someFlag
+                }
+            }
+            """,
+            macros: [
+                "FlagContainer": FlagContainerMacro.self,
+                "Flag": FlagMacro.self,
+            ]
+        )
+    }
+
     func testExpandsButAlreadyConforming() throws {
         assertMacroExpansion(
             """