Add suggested changes

Author: MAJKFL

Sources/SwiftLexicalScopes/LexicalScopes.swift

@@ -13,22 +13,19 @@
 import Foundation
 import SwiftSyntax
 
-public struct LexicalScopes {
-
+public enum LexicalScopes {
   /// Given syntax node position, returns all available labeled statements.
-  public static func lookupLabeledStmts(at syntax: SyntaxProtocol)
-    -> [LabeledStmtSyntax]
-  {
+  public static func lookupLabeledStmts(at syntax: SyntaxProtocol) -> [LabeledStmtSyntax] {
     guard let scope = syntax.scope else { return [] }
     return scope.lookupLabeledStmts(at: syntax)
   }
 
   /// Given syntax node position, returns the current switch case and it's fallthrough destination.
-  public static func lookupFallthroughSourceAndDest(at syntax: SyntaxProtocol) -> (
-    SwitchCaseSyntax?, SwitchCaseSyntax?
-  ) {
+  public static func lookupFallthroughSourceAndDest(
+    at syntax: SyntaxProtocol
+  ) -> (source: SwitchCaseSyntax?, destination: SwitchCaseSyntax?) {
     guard let scope = syntax.scope else { return (nil, nil) }
-    return scope.lookupFallthroughSourceAndDest(at: syntax)
+    return scope.lookupFallthroughSourceAndDestination(at: syntax)
   }
 
   /// Given syntax node position, returns the closest ancestor catch node.

Sources/SwiftLexicalScopes/Scope.swift

@@ -25,20 +25,21 @@ extension SyntaxProtocol {
   }
 }
 
+/// Provide common functionality for specialized scope implementatations.
 protocol Scope {
+  /// The parent of this scope.
   var parent: Scope? { get }
 
+  /// Syntax node that introduces this protocol.
   var sourceSyntax: SyntaxProtocol { get }
 
+  /// Returns the declaration `name` refers to at a particular syntax node location.
   func getDeclarationFor(name: String, at syntax: SyntaxProtocol) -> Syntax?
 }
 
 extension Scope {
-  var parent: Scope? {
-    getParentScope(forSyntax: sourceSyntax)
-  }
-
-  private func getParentScope(forSyntax syntax: SyntaxProtocol?) -> Scope? {
+  /// Recursively walks up syntax tree and finds the closest scope other than this scope.
+  func getParentScope(forSyntax syntax: SyntaxProtocol?) -> Scope? {
     if let lookedUpScope = syntax?.scope, lookedUpScope.sourceSyntax.id == syntax?.id {
       return getParentScope(forSyntax: sourceSyntax.parent)
     } else {
@@ -48,42 +49,35 @@ extension Scope {
 
   // MARK: - lookupLabeledStmts
 
+  /// Given syntax node position, returns all available labeled statements.
   func lookupLabeledStmts(at syntax: SyntaxProtocol) -> [LabeledStmtSyntax] {
-    var result = [LabeledStmtSyntax]()
-    lookupLabeledStmtsHelper(at: syntax.parent, accumulator: &result)
-    return result
+    return lookupLabeledStmtsHelper(at: syntax.parent, accumulator: [])
   }
 
-  private func lookupLabeledStmtsHelper(at syntax: Syntax?, accumulator: inout [LabeledStmtSyntax]) {
-    guard let syntax, !syntax.is(MemberBlockSyntax.self) else { return }
+  /// Helper method to recursively collect labeled statements from the syntax node's parents.
+  private func lookupLabeledStmtsHelper(at syntax: Syntax?, accumulator: [LabeledStmtSyntax]) -> [LabeledStmtSyntax] {
+    guard let syntax, !syntax.is(MemberBlockSyntax.self) else { return accumulator }
     if let labeledStmtSyntax = syntax.as(LabeledStmtSyntax.self) {
-      accumulator.append(labeledStmtSyntax)
-      lookupLabeledStmtsHelper(at: labeledStmtSyntax.parent, accumulator: &accumulator)
+      return lookupLabeledStmtsHelper(at: labeledStmtSyntax.parent, accumulator: accumulator + [labeledStmtSyntax])
     } else {
-      lookupLabeledStmtsHelper(at: syntax.parent, accumulator: &accumulator)
+      return lookupLabeledStmtsHelper(at: syntax.parent, accumulator: accumulator)
     }
   }
 
   // MARK: - lookupFallthroughSourceAndDest
 
-  func lookupFallthroughSourceAndDest(at syntax: SyntaxProtocol) -> (SwitchCaseSyntax?, SwitchCaseSyntax?) {
-    guard let originalSwitchCase = lookupClosestSwitchCaseSyntaxAncestor(at: syntax) else { return (nil, nil) }
+  /// Given syntax node position, returns the current switch case and it's fallthrough destination.
+  func lookupFallthroughSourceAndDestination(at syntax: SyntaxProtocol) -> (SwitchCaseSyntax?, SwitchCaseSyntax?) {
+    guard let originalSwitchCase = syntax.ancestorOrSelf(mapping: { $0.as(SwitchCaseSyntax.self) }) else {
+      return (nil, nil)
+    }
 
     let nextSwitchCase = lookupNextSwitchCase(at: originalSwitchCase)
 
     return (originalSwitchCase, nextSwitchCase)
   }
 
-  private func lookupClosestSwitchCaseSyntaxAncestor(at syntax: SyntaxProtocol?) -> SwitchCaseSyntax? {
-    guard let syntax else { return nil }
-
-    if let switchCaseSyntax = syntax.as(SwitchCaseSyntax.self) {
-      return switchCaseSyntax
-    } else {
-      return lookupClosestSwitchCaseSyntaxAncestor(at: syntax.parent)
-    }
-  }
-
+  /// Given a switch case, returns the case that follows according to the parent.
   private func lookupNextSwitchCase(at switchCaseSyntax: SwitchCaseSyntax) -> SwitchCaseSyntax? {
     guard let switchCaseListSyntax = switchCaseSyntax.parent?.as(SwitchCaseListSyntax.self) else { return nil }
 
@@ -104,30 +98,32 @@ extension Scope {
 
   // MARK: - lookupCatchNode
 
+  /// Given syntax node position, returns the closest ancestor catch node.
   func lookupCatchNode(at syntax: Syntax) -> Syntax? {
     return lookupCatchNodeHelper(at: syntax, traversedCatchClause: false)
   }
 
+  /// Given syntax node location, finds where an error could be caught. If set to `true`, `traverseCatchClause`lookup will skip the next do statement.
   private func lookupCatchNodeHelper(at syntax: Syntax?, traversedCatchClause: Bool) -> Syntax? {
     guard let syntax else { return nil }
 
-    switch syntax.syntaxNodeType {
-    case is DoStmtSyntax.Type:
+    switch syntax.as(SyntaxEnum.self) {
+    case .doStmt:
       if traversedCatchClause {
         return lookupCatchNodeHelper(at: syntax.parent, traversedCatchClause: false)
       } else {
         return syntax
       }
-    case is CatchClauseSyntax.Type:
+    case .catchClause:
       return lookupCatchNodeHelper(at: syntax.parent, traversedCatchClause: true)
-    case is TryExprSyntax.Type:
-      if syntax.as(TryExprSyntax.self)!.questionOrExclamationMark != nil {
+    case .tryExpr(let tryExpr):
+      if tryExpr.questionOrExclamationMark != nil {
         return syntax
       } else {
         return lookupCatchNodeHelper(at: syntax.parent, traversedCatchClause: traversedCatchClause)
       }
-    case is FunctionDeclSyntax.Type:
-      if syntax.as(FunctionDeclSyntax.self)!.signature.effectSpecifiers?.throwsClause != nil {
+    case .functionDecl(let functionDecl):
+      if functionDecl.signature.effectSpecifiers?.throwsClause != nil {
         return syntax
       } else {
         return lookupCatchNodeHelper(at: syntax.parent, traversedCatchClause: traversedCatchClause)

Tests/SwiftLexicalScopesTest/SimpleQueryTests.swift

@@ -18,10 +18,6 @@ final class testSimpleQueries: XCTestCase {
   func testLabeledStmtLookupThreeNestedScopes() {
     assertLexicalScopeQuery(
       source: """
-        while true {
-          5️⃣break
-        }
-
         1️⃣one: for i in 1..<10 {
           while true {
             2️⃣two: do {
@@ -34,7 +30,21 @@ final class testSimpleQueries: XCTestCase {
       methodUnderTest: { argument in
         LexicalScopes.lookupLabeledStmts(at: argument)
       },
-      expected: ["3️⃣": ["2️⃣", "1️⃣"], "4️⃣": ["1️⃣"], "5️⃣": []]
+      expected: ["3️⃣": ["2️⃣", "1️⃣"], "4️⃣": ["1️⃣"]]
+    )
+  }
+
+  func testNoLabeledStatements() {
+    assertLexicalScopeQuery(
+      source: """
+        while true {
+          1️⃣break
+        }
+        """,
+      methodUnderTest: { argument in
+        LexicalScopes.lookupLabeledStmts(at: argument)
+      },
+      expected: ["1️⃣": []]
     )
   }
 
@@ -63,21 +73,21 @@ final class testSimpleQueries: XCTestCase {
     assertLexicalScopeQuery(
       source: """
         func foo() {
-            7️⃣print(0)
+          7️⃣print(0)
         }
 
         switch a {
         1️⃣case 1:
-            2️⃣print(1)
+          2️⃣print(1)
         3️⃣case 2:
-            4️⃣print(2)
+          4️⃣print(2)
         5️⃣default:
-            6️⃣print(3)
+          6️⃣print(3)
         }
         """,
       methodUnderTest: { argument in
         let result = LexicalScopes.lookupFallthroughSourceAndDest(at: argument)
-        return [result.0, result.1]
+        return [result.source, result.destination]
       },
       expected: ["2️⃣": ["1️⃣", "3️⃣"], "4️⃣": ["3️⃣", "5️⃣"], "6️⃣": ["5️⃣", nil], "7️⃣": [nil, nil]]
     )
@@ -112,9 +122,9 @@ final class testSimpleQueries: XCTestCase {
         1️⃣func foo() rethrows {
           2️⃣do {
             3️⃣do {
-                try 4️⃣f()
+              try 4️⃣f()
             } catch {
-                try 5️⃣f()
+              try 5️⃣f()
             }
           } catch {
             6️⃣try! 7️⃣f()