Merge pull request #2014 from ahoppen/ahoppen/index-at

Add a function to get the `SyntaxChildrenIndex` of the n-th element in a `SyntaxCollection`

Author: ahoppen

Sources/SwiftSyntax/SyntaxChildren.swift

@@ -189,7 +189,10 @@ struct RawSyntaxChildren: BidirectionalCollection {
   }
 
   func index(after index: SyntaxChildrenIndex) -> SyntaxChildrenIndex {
-    let node = parent.layoutView!.children[Int(index.data!.indexInParent)]
+    guard let indexData = index.data else {
+      preconditionFailure("Cannot get the index after the end index")
+    }
+    let node = parent.layoutView!.children[Int(indexData.indexInParent)]
     return self.index(index, advancedBy: node)
   }
 

Sources/SwiftSyntax/SyntaxCollection.swift

@@ -105,6 +105,28 @@ extension SyntaxCollection {
     return node.indexInParent
   }
 
+  /// Returns the index of the n-th element in this collection.
+  ///
+  /// The behavior is undefined if `offset` is greater than the number of
+  /// elements in this collection.
+  ///
+  /// - Complexity: O(`offset`) because all previous element need to be iterated
+  ///   to find the byte offset of the `offset`-th node within the source file.
+  ///
+  /// - Note: Because getting the `n`-th element in a ``SyntaxCollection`` is
+  ///   not O(1), ``SyntaxCollection`` doesn’t provide a subscript to retrieve
+  ///   the `n`-th element. Such a subscript would mask the complexity of
+  ///   getting the `n`-th element.
+  public func index(at offset: Int) -> SyntaxChildrenIndex {
+    if offset > self.count / 2 {
+      // If we are closer to the end of the collection, it's more efficient to
+      // calculate the index starting from the back.
+      return self.index(self.endIndex, offsetBy: -(self.count - offset))
+    } else {
+      return self.index(self.startIndex, offsetBy: offset)
+    }
+  }
+
   /// Creates a new collection by appending the provided syntax element
   /// to the children.
   ///

Sources/_SwiftSyntaxTestSupport/SyntaxCollection+AtIndex.swift.swift

@@ -111,7 +111,7 @@ fileprivate func assertRefactorPlaceholderCall(
 ) throws {
   var parser = Parser(expr)
   let call = try XCTUnwrap(ExprSyntax.parse(from: &parser).as(FunctionCallExprSyntax.self), file: file, line: line)
-  let arg = try XCTUnwrap(call.arguments[placeholder].as(LabeledExprSyntax.self), file: file, line: line)
+  let arg = try XCTUnwrap(call.arguments[call.arguments.index(at: placeholder)].as(LabeledExprSyntax.self), file: file, line: line)
   let token: TokenSyntax = try XCTUnwrap(arg.expression.as(EditorPlaceholderExprSyntax.self), file: file, line: line).placeholder
 
   try assertRefactor(token, context: (), provider: ExpandEditorPlaceholders.self, expected: [SourceEdit.replace(call, with: expected)], file: file, line: line)
@@ -126,7 +126,7 @@ fileprivate func assertRefactorPlaceholderToken(
 ) throws {
   var parser = Parser(expr)
   let call = try XCTUnwrap(ExprSyntax.parse(from: &parser).as(FunctionCallExprSyntax.self), file: file, line: line)
-  let arg = try XCTUnwrap(call.arguments[placeholder].as(LabeledExprSyntax.self), file: file, line: line)
+  let arg = try XCTUnwrap(call.arguments[call.arguments.index(at: placeholder)].as(LabeledExprSyntax.self), file: file, line: line)
   let token: TokenSyntax = try XCTUnwrap(arg.expression.as(EditorPlaceholderExprSyntax.self), file: file, line: line).placeholder
 
   try assertRefactor(token, context: (), provider: ExpandEditorPlaceholders.self, expected: [SourceEdit.replace(token, with: expected)], file: file, line: line)

Tests/SwiftRefactorTest/ExpandEditorPlaceholdersTests.swift

@@ -145,7 +145,7 @@ public struct ErrorMacro: DeclarationMacro {
       let stringLiteral = firstElement.expression
         .as(StringLiteralExprSyntax.self),
       stringLiteral.segments.count == 1,
-      case let .stringSegment(messageString) = stringLiteral.segments[0]
+      case let .stringSegment(messageString) = stringLiteral.segments.first
     else {
       throw MacroExpansionErrorMessage("#error macro requires a string literal")
     }
@@ -170,7 +170,7 @@ struct DefineBitwidthNumberedStructsMacro: DeclarationMacro {
       let stringLiteral = firstElement.expression
         .as(StringLiteralExprSyntax.self),
       stringLiteral.segments.count == 1,
-      case let .stringSegment(prefix) = stringLiteral.segments[0]
+      case let .stringSegment(prefix) = stringLiteral.segments.first
     else {
       throw MacroExpansionErrorMessage(
         "#bitwidthNumberedStructs macro requires a string literal"

Tests/SwiftSyntaxMacroExpansionTest/MacroSystemTests.swift

@@ -33,9 +33,10 @@ public class AbsolutePositionTests: XCTestCase {
       statements: CodeBlockItemListSyntax(l),
       endOfFileToken: .endOfFileToken()
     )
-    _ = root.statements[idx].position
-    _ = root.statements[idx].totalLength.utf8Length
-    _ = root.statements[idx].positionAfterSkippingLeadingTrivia
+    let statements = root.statements
+    _ = statements[statements.index(at: idx)].position
+    _ = statements[statements.index(at: idx)].totalLength.utf8Length
+    _ = statements[statements.index(at: idx)].positionAfterSkippingLeadingTrivia
   }
 
   static let leadingTrivia = Trivia(pieces: [
@@ -76,7 +77,7 @@ public class AbsolutePositionTests: XCTestCase {
     let root = self.createSourceFile(idx + 1)
     XCTAssertEqual(3, root.leadingTrivia.count)
     XCTAssertEqual(0, root.trailingTrivia.count)
-    let state = root.statements[idx]
+    let state = root.statements[root.statements.index(at: idx)]
     XCTAssertEqual(3, state.leadingTrivia.count)
     XCTAssertEqual(2, state.trailingTrivia.count)
     XCTAssertEqual(
@@ -137,7 +138,7 @@ public class AbsolutePositionTests: XCTestCase {
     let filePath = "/tmp/test.swift"
     let root = self.createSourceFile(2)
     let converter = SourceLocationConverter(fileName: filePath, tree: root)
-    let secondReturnStmt = root.statements[1]
+    let secondReturnStmt = root.statements[root.statements.index(at: 1)]
     let startLoc = secondReturnStmt.startLocation(converter: converter)
     XCTAssertEqual(startLoc.line, 8)
     XCTAssertEqual(startLoc.column, 1)

Tests/SwiftSyntaxTest/AbsolutePositionTests.swift

@@ -90,7 +90,7 @@ public class SyntaxCollectionsTests: XCTestCase {
     assertSyntaxCollectionManipulation(
       initialElements: [0, 1, 2],
       transformation: {
-        let index = $0.elements.index(after: $0.elements.startIndex)
+        let index = $0.elements.index(at: 1)
         $0.elements.insert(intElement(5), at: index)
       },
       expectedElements: [0, 5, 1, 2]
@@ -119,7 +119,7 @@ public class SyntaxCollectionsTests: XCTestCase {
     assertSyntaxCollectionManipulation(
       initialElements: [0, 1, 2],
       transformation: {
-        let index = $0.elements.index(after: $0.elements.startIndex)
+        let index = $0.elements.index(at: 1)
         $0.elements.insert(contentsOf: [intElement(5), intElement(6)], at: index)
       },
       expectedElements: [0, 5, 6, 1, 2]
@@ -148,7 +148,7 @@ public class SyntaxCollectionsTests: XCTestCase {
     assertSyntaxCollectionManipulation(
       initialElements: [0, 1, 2],
       transformation: {
-        let index = $0.elements.index(after: $0.elements.startIndex)
+        let index = $0.elements.index(at: 1)
         $0.elements.remove(at: index)
       },
       expectedElements: [0, 2]
@@ -157,7 +157,7 @@ public class SyntaxCollectionsTests: XCTestCase {
     assertSyntaxCollectionManipulation(
       initialElements: [0, 1, 2],
       transformation: {
-        let index = $0.elements.index($0.elements.startIndex, offsetBy: 2)
+        let index = $0.elements.index(at: 2)
         $0.elements.remove(at: index)
       },
       expectedElements: [0, 1]
@@ -169,7 +169,7 @@ public class SyntaxCollectionsTests: XCTestCase {
       initialElements: [0, 1, 2],
       transformation: {
         let startIndex = $0.elements.startIndex
-        let endIndex = $0.elements.index(after: $0.elements.startIndex)
+        let endIndex = $0.elements.index(at: 1)
         $0.elements.removeSubrange(startIndex..<endIndex)
       },
       expectedElements: [1, 2]
@@ -179,7 +179,7 @@ public class SyntaxCollectionsTests: XCTestCase {
       initialElements: [0, 1, 2],
       transformation: {
         let startIndex = $0.elements.startIndex
-        let endIndex = $0.elements.index($0.elements.startIndex, offsetBy: 2)
+        let endIndex = $0.elements.index(at: 2)
         $0.elements.removeSubrange(startIndex..<endIndex)
       },
       expectedElements: [2]