Fix NoAccessLevelOnExtensionDeclaration to update members inside #if blocks.

Fixes #966.

Author: allevato

Sources/SwiftFormat/Rules/NoAccessLevelOnExtensionDeclaration.swift

@@ -21,88 +21,197 @@ import SwiftSyntax
 ///         `internal`, as that is the default access level) have the explicit access level removed.
 @_spi(Rules)
 public final class NoAccessLevelOnExtensionDeclaration: SyntaxFormatRule {
+  private enum State {
+    /// The rule is currently visiting top-level declarations.
+    case topLevel
+
+    /// The rule is currently inside an extension that has the given access level keyword.
+    case insideExtension(accessKeyword: Keyword)
+  }
+
+  /// Tracks the state of the rule to determine which action should be taken on visited
+  /// declarations.
+  private var state: State = .topLevel
+
+  /// Findings propagated up to the extension visitor from any members that were rewritten.
+  private var notesFromRewrittenMembers: [Finding.Note] = []
+
   public override func visit(_ node: ExtensionDeclSyntax) -> DeclSyntax {
     guard
+      // Skip nested extensions; these are semantic errors but they still parse successfully.
+      case .topLevel = state,
+      // Skip extensions that don't have an access level modifier.
       let accessKeyword = node.modifiers.accessLevelModifier,
       case .keyword(let keyword) = accessKeyword.name.tokenKind
     else {
       return DeclSyntax(node)
     }
 
-    var result = node
+    self.notesFromRewrittenMembers = []
+
+    let keywordToAdd: Keyword?
+    let message: Finding.Message
 
     switch keyword {
-    // Public, private, fileprivate, or package keywords need to be moved to members
     case .public, .private, .fileprivate, .package:
-      // The effective access level of the members of a `private` extension is `fileprivate`, so
-      // we have to update the keyword to ensure that the result is correct.
-      var accessKeywordToAdd = accessKeyword
-      let message: Finding.Message
+      // These access level modifiers need to be moved to members. Additionally, `private` is a
+      // special case, because the *effective* access level for a top-level private extension is
+      // `fileprivate`, so we need to preserve that when we apply it to the members.
       if keyword == .private {
-        accessKeywordToAdd.name.tokenKind = .keyword(.fileprivate)
+        keywordToAdd = .fileprivate
         message = .moveAccessKeywordAndMakeFileprivate(keyword: accessKeyword.name.text)
       } else {
+        keywordToAdd = keyword
         message = .moveAccessKeyword(keyword: accessKeyword.name.text)
       }
 
-      let (newMembers, notes) =
-        addMemberAccessKeyword(accessKeywordToAdd, toMembersIn: node.memberBlock)
-      diagnose(message, on: accessKeyword, notes: notes)
-
-      result.modifiers.remove(anyOf: [keyword])
-      result.extensionKeyword.leadingTrivia = accessKeyword.leadingTrivia
-      result.memberBlock.members = newMembers
-      return DeclSyntax(result)
-
-    // Internal keyword redundant, delete
     case .internal:
-      diagnose(.removeRedundantAccessKeyword, on: accessKeyword)
-
-      result.modifiers.remove(anyOf: [keyword])
-      result.extensionKeyword.leadingTrivia = accessKeyword.leadingTrivia
-      return DeclSyntax(result)
+      // If the access level keyword was `internal`, then it's redundant and we can just remove it.
+      // We don't need to modify the members at all in this case.
+      message = .removeRedundantAccessKeyword
+      keywordToAdd = nil
 
     default:
-      break
+      // For anything else, just return the extension and its members unchanged.
+      return DeclSyntax(node)
+    }
+
+    // We don't have to worry about maintaining a stack here; even though extensions can nest from
+    // a valid parse point of view, we ignore nested extensions because they're obviously wrong
+    // semantically (and would be an error later during compilation).
+    var result: ExtensionDeclSyntax
+    if let keywordToAdd {
+      // Visit the children in the new state to add the keyword to the extension members.
+      self.state = .insideExtension(accessKeyword: keywordToAdd)
+      defer { self.state = .topLevel }
+
+      result = super.visit(node).as(ExtensionDeclSyntax.self)!
+    } else {
+      // We don't need to visit the children in this case, and we don't need to update the state.
+      result = node
     }
 
+    // Finally, emit the finding (which includes notes from any rewritten members) and remove the
+    // access level keyword from the extension itself.
+    diagnose(message, on: accessKeyword, notes: self.notesFromRewrittenMembers)
+    result.modifiers.remove(anyOf: [keyword])
+    result.extensionKeyword.leadingTrivia = accessKeyword.leadingTrivia
     return DeclSyntax(result)
   }
 
-  // Adds given keyword to all members in declaration block
-  private func addMemberAccessKeyword(
-    _ modifier: DeclModifierSyntax,
-    toMembersIn memberBlock: MemberBlockSyntax
-  ) -> (MemberBlockItemListSyntax, [Finding.Note]) {
-    var newMembers: [MemberBlockItemSyntax] = []
-    var notes: [Finding.Note] = []
-
-    for memberItem in memberBlock.members {
-      let decl = memberItem.decl
-      guard
-        let modifiers = decl.asProtocol(WithModifiersSyntax.self)?.modifiers,
-        modifiers.accessLevelModifier == nil
-      else {
-        newMembers.append(memberItem)
-        continue
-      }
+  public override func visit(_ node: ActorDeclSyntax) -> DeclSyntax {
+    return applyingAccessModifierIfNone(to: node)
+  }
+
+  public override func visit(_ node: ClassDeclSyntax) -> DeclSyntax {
+    return applyingAccessModifierIfNone(to: node)
+  }
+
+  public override func visit(_ node: EnumDeclSyntax) -> DeclSyntax {
+    return applyingAccessModifierIfNone(to: node)
+  }
+
+  public override func visit(_ node: FunctionDeclSyntax) -> DeclSyntax {
+    return applyingAccessModifierIfNone(to: node)
+  }
+
+  public override func visit(_ node: InitializerDeclSyntax) -> DeclSyntax {
+    return applyingAccessModifierIfNone(to: node)
+  }
+
+  public override func visit(_ node: StructDeclSyntax) -> DeclSyntax {
+    return applyingAccessModifierIfNone(to: node)
+  }
+
+  public override func visit(_ node: SubscriptDeclSyntax) -> DeclSyntax {
+    return applyingAccessModifierIfNone(to: node)
+  }
 
-      // Create a note associated with each declaration that needs to have an access level modifier
-      // added to it.
-      notes.append(
-        Finding.Note(
-          message: .addModifierToExtensionMember(keyword: modifier.name.text),
-          location:
-            Finding.Location(decl.startLocation(converter: context.sourceLocationConverter))
-        )
+  public override func visit(_ node: TypeAliasDeclSyntax) -> DeclSyntax {
+    return applyingAccessModifierIfNone(to: node)
+  }
+
+  public override func visit(_ node: VariableDeclSyntax) -> DeclSyntax {
+    return applyingAccessModifierIfNone(to: node)
+  }
+
+  /// Adds `modifier` to `decl` if it doesn't already have an explicit access level modifier and
+  /// returns the new declaration.
+  ///
+  /// If `decl` already has an access level modifier, it is returned unchanged.
+  private func applyingAccessModifierIfNone(to decl: some DeclSyntaxProtocol) -> DeclSyntax {
+    // Only go further if we are applying an access level keyword and if the decl is one that
+    // allows modifiers but doesn't already have an access level modifier.
+    guard
+      case .insideExtension(let accessKeyword) = state,
+      let modifiers = decl.asProtocol(WithModifiersSyntax.self)?.modifiers,
+      modifiers.accessLevelModifier == nil
+    else {
+      return DeclSyntax(decl)
+    }
+
+    // Create a note associated with each declaration that needs to have an access level modifier
+    // added to it.
+    self.notesFromRewrittenMembers.append(
+      Finding.Note(
+        message: .addModifierToExtensionMember(keyword: TokenSyntax.keyword(accessKeyword).text),
+        location:
+          Finding.Location(decl.startLocation(converter: context.sourceLocationConverter))
       )
+    )
 
-      var newItem = memberItem
-      newItem.decl = applyingAccessModifierIfNone(modifier, to: decl)
-      newMembers.append(newItem)
+    switch Syntax(decl).as(SyntaxEnum.self) {
+    case .actorDecl(let actorDecl):
+      return applyingAccessModifierIfNone(accessKeyword, to: actorDecl, declKeywordKeyPath: \.actorKeyword)
+    case .classDecl(let classDecl):
+      return applyingAccessModifierIfNone(accessKeyword, to: classDecl, declKeywordKeyPath: \.classKeyword)
+    case .enumDecl(let enumDecl):
+      return applyingAccessModifierIfNone(accessKeyword, to: enumDecl, declKeywordKeyPath: \.enumKeyword)
+    case .initializerDecl(let initDecl):
+      return applyingAccessModifierIfNone(accessKeyword, to: initDecl, declKeywordKeyPath: \.initKeyword)
+    case .functionDecl(let funcDecl):
+      return applyingAccessModifierIfNone(accessKeyword, to: funcDecl, declKeywordKeyPath: \.funcKeyword)
+    case .structDecl(let structDecl):
+      return applyingAccessModifierIfNone(accessKeyword, to: structDecl, declKeywordKeyPath: \.structKeyword)
+    case .subscriptDecl(let subscriptDecl):
+      return applyingAccessModifierIfNone(accessKeyword, to: subscriptDecl, declKeywordKeyPath: \.subscriptKeyword)
+    case .typeAliasDecl(let typeAliasDecl):
+      return applyingAccessModifierIfNone(accessKeyword, to: typeAliasDecl, declKeywordKeyPath: \.typealiasKeyword)
+    case .variableDecl(let varDecl):
+      return applyingAccessModifierIfNone(accessKeyword, to: varDecl, declKeywordKeyPath: \.bindingSpecifier)
+    default:
+      return DeclSyntax(decl)
+    }
+  }
+
+  private func applyingAccessModifierIfNone<Decl: DeclSyntaxProtocol & WithModifiersSyntax>(
+    _ modifier: Keyword,
+    to decl: Decl,
+    declKeywordKeyPath: WritableKeyPath<Decl, TokenSyntax>
+  ) -> DeclSyntax {
+    // If there's already an access modifier among the modifier list, bail out.
+    guard decl.modifiers.accessLevelModifier == nil else { return DeclSyntax(decl) }
+
+    var result = decl
+    var modifier = DeclModifierSyntax(name: .keyword(modifier))
+    modifier.trailingTrivia = [.spaces(1)]
+
+    guard var firstModifier = decl.modifiers.first else {
+      // If there are no modifiers at all, add the one being requested, moving the leading trivia
+      // from the decl keyword to that modifier (to preserve leading comments, newlines, etc.).
+      modifier.leadingTrivia = decl[keyPath: declKeywordKeyPath].leadingTrivia
+      result[keyPath: declKeywordKeyPath].leadingTrivia = []
+      result.modifiers = .init([modifier])
+      return DeclSyntax(result)
     }
 
-    return (MemberBlockItemListSyntax(newMembers), notes)
+    // Otherwise, insert the modifier at the front of the modifier list, moving the (original) first
+    // modifier's leading trivia to the new one (to preserve leading comments, newlines, etc.).
+    modifier.leadingTrivia = firstModifier.leadingTrivia
+    firstModifier.leadingTrivia = []
+    result.modifiers[result.modifiers.startIndex] = firstModifier
+    result.modifiers.insert(modifier, at: result.modifiers.startIndex)
+    return DeclSyntax(result)
   }
 }
 
@@ -122,81 +231,3 @@ extension Finding.Message {
     "add '\(keyword)' access modifier to this declaration"
   }
 }
-
-/// Adds `modifier` to `decl` if it doesn't already have an explicit access level modifier and
-/// returns the new declaration.
-///
-/// If `decl` already has an access level modifier, it is returned unchanged.
-private func applyingAccessModifierIfNone(
-  _ modifier: DeclModifierSyntax,
-  to decl: DeclSyntax
-) -> DeclSyntax {
-  switch Syntax(decl).as(SyntaxEnum.self) {
-  case .actorDecl(let actorDecl):
-    return applyingAccessModifierIfNone(modifier, to: actorDecl, declKeywordKeyPath: \.actorKeyword)
-  case .classDecl(let classDecl):
-    return applyingAccessModifierIfNone(modifier, to: classDecl, declKeywordKeyPath: \.classKeyword)
-  case .enumDecl(let enumDecl):
-    return applyingAccessModifierIfNone(modifier, to: enumDecl, declKeywordKeyPath: \.enumKeyword)
-  case .initializerDecl(let initDecl):
-    return applyingAccessModifierIfNone(modifier, to: initDecl, declKeywordKeyPath: \.initKeyword)
-  case .functionDecl(let funcDecl):
-    return applyingAccessModifierIfNone(modifier, to: funcDecl, declKeywordKeyPath: \.funcKeyword)
-  case .structDecl(let structDecl):
-    return applyingAccessModifierIfNone(
-      modifier,
-      to: structDecl,
-      declKeywordKeyPath: \.structKeyword
-    )
-  case .subscriptDecl(let subscriptDecl):
-    return applyingAccessModifierIfNone(
-      modifier,
-      to: subscriptDecl,
-      declKeywordKeyPath: \.subscriptKeyword
-    )
-  case .typeAliasDecl(let typeAliasDecl):
-    return applyingAccessModifierIfNone(
-      modifier,
-      to: typeAliasDecl,
-      declKeywordKeyPath: \.typealiasKeyword
-    )
-  case .variableDecl(let varDecl):
-    return applyingAccessModifierIfNone(
-      modifier,
-      to: varDecl,
-      declKeywordKeyPath: \.bindingSpecifier
-    )
-  default:
-    return decl
-  }
-}
-
-private func applyingAccessModifierIfNone<Decl: DeclSyntaxProtocol & WithModifiersSyntax>(
-  _ modifier: DeclModifierSyntax,
-  to decl: Decl,
-  declKeywordKeyPath: WritableKeyPath<Decl, TokenSyntax>
-) -> DeclSyntax {
-  // If there's already an access modifier among the modifier list, bail out.
-  guard decl.modifiers.accessLevelModifier == nil else { return DeclSyntax(decl) }
-
-  var result = decl
-  var modifier = modifier
-  modifier.trailingTrivia = [.spaces(1)]
-
-  guard var firstModifier = decl.modifiers.first else {
-    // If there are no modifiers at all, add the one being requested, moving the leading trivia
-    // from the decl keyword to that modifier (to preserve leading comments, newlines, etc.).
-    modifier.leadingTrivia = decl[keyPath: declKeywordKeyPath].leadingTrivia
-    result[keyPath: declKeywordKeyPath].leadingTrivia = []
-    result.modifiers = .init([modifier])
-    return DeclSyntax(result)
-  }
-
-  // Otherwise, insert the modifier at the front of the modifier list, moving the (original) first
-  // modifier's leading trivia to the new one (to preserve leading comments, newlines, etc.).
-  modifier.leadingTrivia = firstModifier.leadingTrivia
-  firstModifier.leadingTrivia = []
-  result.modifiers[result.modifiers.startIndex] = firstModifier
-  result.modifiers.insert(modifier, at: result.modifiers.startIndex)
-  return DeclSyntax(result)
-}