Improve fixit trivia and modifier handling (#49)

* Do not delete modifiers when adding 'public' in fixit * Better trivia management and replace incorrect accessibility modifiers * Improve coverage
dfed · Jan 29, 2024 · 25a40e0 · 25a40e0
1 parent 7dfe063
commit 25a40e0
Show file tree

Hide file tree

Showing 3 changed files with 369 additions and 33 deletions.
diff --git a/Sources/SafeDICore/Models/ConcreteDeclSyntaxProtocol.swift b/Sources/SafeDICore/Models/ConcreteDeclSyntaxProtocol.swift
@@ -24,16 +24,44 @@ public protocol ConcreteDeclSyntaxProtocol: SyntaxProtocol {
     var attributes: AttributeListSyntax { get set }
     var modifiers: DeclModifierListSyntax { get set }
     var inheritanceClause: InheritanceClauseSyntax? { get set }
+    var keyword: TokenSyntax { get set }
     var name: TokenSyntax { get set }
     var declType: ConcreteDeclType { get }
 }
 
 extension ActorDeclSyntax: ConcreteDeclSyntaxProtocol {
+    public var keyword: SwiftSyntax.TokenSyntax {
+        get {
+            actorKeyword
+        }
+        set {
+            actorKeyword = newValue
+        }
+    }
+
     public var declType: ConcreteDeclType { .actorType }
 }
 extension ClassDeclSyntax: ConcreteDeclSyntaxProtocol {
+    public var keyword: SwiftSyntax.TokenSyntax {
+        get {
+            classKeyword
+        }
+        set {
+            classKeyword = newValue
+        }
+    }
+
     public var declType: ConcreteDeclType { .classType }
 }
 extension StructDeclSyntax: ConcreteDeclSyntaxProtocol {
+    public var keyword: SwiftSyntax.TokenSyntax {
+        get {
+            structKeyword
+        }
+        set {
+            structKeyword = newValue
+        }
+    }
+
     public var declType: ConcreteDeclType { .structType }
 }
diff --git a/Sources/SafeDICore/Visitors/InstantiableVisitor.swift b/Sources/SafeDICore/Visitors/InstantiableVisitor.swift
@@ -146,7 +146,7 @@ public final class InstantiableVisitor: SyntaxVisitor {
             return .skipChildren
         }
         guard let instantiableMacro = node.attributes.instantiableMacro else {
-            // Not an external instantiable type. We do not care.
+            // Not an instantiable type. We do not care.
             return .skipChildren
         }
 
@@ -160,20 +160,16 @@ public final class InstantiableVisitor: SyntaxVisitor {
         guard declarationType.isExtension else {
             return .skipChildren
         }
-        guard let instantiableType else {
-            // We're being called on code that will not compile.
-            // We are visiting a function but we haven't visited the extension yet.
-            // Just move on.
-            return .skipChildren
-        }
         guard node.name.text == Self.instantiateMethodName else {
             // We don't care about this method.
             return .skipChildren
         }
 
         if
             let returnClause = node.signature.returnClause,
-            returnClause.type.typeDescription != instantiableType {
+            returnClause.type.typeDescription != instantiableType,
+            let instantiableType
+        {
             var modifiedSignature = node.signature
             modifiedSignature.returnClause = ReturnClauseSyntax(
                 arrow: .arrowToken(
@@ -215,7 +211,7 @@ public final class InstantiableVisitor: SyntaxVisitor {
                     DeclModifierSyntax(
                         name: TokenSyntax(
                             TokenKind.keyword(.public),
-                            leadingTrivia: .newline,
+                            leadingTrivia: node.modifiers.first?.leadingTrivia ?? node.funcKeyword.leadingTrivia,
                             presence: .present
                         )
                     ),
@@ -228,6 +224,7 @@ public final class InstantiableVisitor: SyntaxVisitor {
                     )
                 )
             )
+            modifiedNode.funcKeyword.leadingTrivia = .spaces(0)
             diagnostics.append(Diagnostic(
                 node: node,
                 error: FixableInstantiableError.missingAttributes,
@@ -397,18 +394,27 @@ public final class InstantiableVisitor: SyntaxVisitor {
 
     private func processModifiers(_ modifiers: DeclModifierListSyntax, on node: some ConcreteDeclSyntaxProtocol) {
         if !node.modifiers.containsPublicOrOpen {
-            var modifiedNode = node
-            modifiedNode.modifiers = DeclModifierListSyntax(
-                arrayLiteral:
-                    DeclModifierSyntax(
-                        name: TokenSyntax(
-                            TokenKind.keyword(.public),
-                            leadingTrivia: .newline,
-                            trailingTrivia: .space,
-                            presence: .present
-                        )
-                    )
+            let publicModifier = DeclModifierSyntax(
+                name: TokenSyntax(
+                    TokenKind.keyword(.public),
+                    leadingTrivia: node.modifiers.first?.leadingTrivia ?? .newline,
+                    trailingTrivia: node.modifiers.first?.trailingTrivia ?? .space,
+                    presence: .present
+                )
             )
+            var modifiedNode = node
+            if var firstModifier = modifiedNode.modifiers.first {
+                firstModifier.name.leadingTrivia = .spaces(0)
+                modifiedNode.modifiers.replaceSubrange(
+                    modifiedNode.modifiers.startIndex..<modifiedNode.attributes.index(after: modifiedNode.attributes.startIndex),
+                    with: [publicModifier, firstModifier])
+                modifiedNode.modifiers = modifiedNode.modifiers.filter {
+                    $0.name.text != "internal" && $0.name.text != "fileprivate" && $0.name.text != "private"
+                }
+            } else {
+                modifiedNode.modifiers = [publicModifier]
+                modifiedNode.keyword.leadingTrivia = .spaces(0)
+            }
             diagnostics.append(Diagnostic(
                 node: node,
                 error: FixableInstantiableError.missingPublicOrOpenAttribute,