Improve the IndexStoreFinder architecture and testability

README.md

@@ -4,7 +4,7 @@
 </p>
 
 # 🔎 SwiftFindRefs
-A Swift Package Manager CLI that locates every file in your Xcode DerivedData index referencing a chosen symbol. It resolves the correct IndexStore path automatically, queries Apple’s IndexStoreDB, and prints a deduplicated list of source files.
+A Swift Package Manager CLI that locates every file in your Xcode DerivedData index referencing a chosen symbol. It resolves the correct IndexStore path automatically, queries Apple’s IndexStoreDB, and prints a deduplicated list of source files. It uses Swift concurrency to scan multiple files in parallel, keeping discovery fast even for large workspaces.
 
 ## 🚀 Common use case 
 When working with multiple modules and moving models between them, finding all references to add missing imports is tedious. Using this CLI to feed file lists to AI agents dramatically improves refactoring results.

Sources/SwiftFindRefs/IndexStore/IndexStore+Providing.swift

@@ -0,0 +1,39 @@
+import IndexStore
+
+// MARK: - IndexStore Library Conformance
+
+extension UnitDependency: UnitDependencyProviding {}
+
+extension UnitReader: UnitReaderProviding {
+    func forEachDependency(_ callback: (UnitDependencyProviding) -> Void) {
+        forEach { dependency in
+            callback(dependency)
+        }
+    }
+}
+
+extension IndexStore: IndexStoreProviding {
+    func forEachUnit(_ callback: (UnitReaderProviding) -> Void) {
+        for unit in units {
+            callback(unit)
+        }
+    }
+
+    func recordReader(for recordName: String) throws -> RecordReaderProviding? {
+        try? RecordReader(indexStore: self, recordName: recordName)
+    }
+}
+
+extension RecordReader: RecordReaderProviding {
+    func forEachOccurrence(_ callback: (SymbolOccurrenceProviding) -> Void) {
+        forEach { occurrence in
+            callback(occurrence)
+        }
+    }
+}
+
+extension SymbolOccurrence: SymbolOccurrenceProviding {
+    var symbolMatching: SymbolMatching {
+        symbol
+    }
+}

Sources/SwiftFindRefs/IndexStore/IndexStoreFinder.swift

@@ -1,91 +1,78 @@
 import Foundation
-import IndexStore
+@preconcurrency import IndexStore
 
 struct IndexStoreFinder {
     let indexStorePath: String
 
     func fileReferences(of symbolName: String, symbolType: String?) throws -> [String] {
         let store = try IndexStore(path: indexStorePath)
-
-        // Pre-compute SymbolKind enum to avoid string comparison in hot loop
-        let expectedSymbolKind: SymbolKind? = symbolType.flatMap { parseSymbolKind($0) }
+        return try fileReferences(of: symbolName, symbolType: symbolType, from: store)
+    }
 
-        // Collect all record dependencies with their source paths in a single pass
-        var recordToSource: [String: String] = [:]
-        var allRecordNames = Set<String>()
+    func fileReferences(
+        of symbolName: String,
+        symbolType: String?,
+        from store: some IndexStoreProviding & Sendable
+    ) throws -> [String] {
+        let query = SymbolQuery(name: symbolName, kindString: symbolType)
+        let index = RecordIndex.build(from: store)
 
-        for unitReader in store.units {
-            // Skip system frameworks (SDK headers, etc.)
-            guard !unitReader.isSystem else { continue }
-
-            unitReader.forEach { dependency in
-                guard dependency.kind == .record else { return }
-                let recordName = dependency.name
-                allRecordNames.insert(recordName)
-                // Only store non-empty paths, prefer keeping existing paths
-                let filePath = dependency.filePath
-                if !filePath.isEmpty && recordToSource[recordName] == nil {
-                    recordToSource[recordName] = filePath
-                }
-            }
-        }
+        return searchRecordsInParallel(store: store, index: index, query: query)
+    }
 
-        // Convert to array for parallel processing
-        let recordNames = Array(allRecordNames)
-        let lock = NSLock()
-        var referencedFiles = Set<String>()
+    private func searchRecordsInParallel(
+        store: some IndexStoreProviding & Sendable,
+        index: RecordIndex,
+        query: SymbolQuery
+    ) -> [String] {
+        let referencedFiles = ThreadSafeSet<String>()
 
-        // Process records in parallel across all CPU cores
-        DispatchQueue.concurrentPerform(iterations: recordNames.count) { index in
-            let recordName = recordNames[index]
-            guard let recordReader = try? RecordReader(indexStore: store, recordName: recordName) else {
-                return
-            }
+        DispatchQueue.concurrentPerform(iterations: index.recordNames.count) { i in
+            let recordName = index.recordNames[i]
 
-            var foundInRecord = false
-            recordReader.forEach { (occurrence: SymbolOccurrence) in
-                guard !foundInRecord else { return }
-                guard occurrence.symbol.name == symbolName else { return }
-                if let expectedKind = expectedSymbolKind, occurrence.symbol.kind != expectedKind {
-                    return
-                }
-                foundInRecord = true
-            }
-
-            if foundInRecord {
-                let filename = recordToSource[recordName] ?? recordName
-                lock.lock()
+            if recordContainsSymbol(store: store, recordName: recordName, query: query) {
+                let filename = index.sourcePath(for: recordName)
                 referencedFiles.insert(filename)
-                lock.unlock()
             }
         }
 
-        return referencedFiles.sorted()
+        return referencedFiles.values().sorted()
     }
-
-    private func parseSymbolKind(_ type: String) -> SymbolKind? {
-        switch type.lowercased() {
-        case "class": return .class
-        case "struct": return .struct
-        case "enum": return .enum
-        case "protocol": return .protocol
-        case "function": return .function
-        case "variable": return .variable
-        case "typealias": return .typealias
-        case "instancemethod": return .instanceMethod
-        case "staticmethod": return .staticMethod
-        case "classmethod": return .classMethod
-        case "instanceproperty": return .instanceProperty
-        case "staticproperty": return .staticProperty
-        case "classproperty": return .classProperty
-        case "constructor": return .constructor
-        case "destructor": return .destructor
-        case "field": return .field
-        case "enumconstant": return .enumConstant
-        case "parameter": return .parameter
-        case "module": return .module
-        case "extension": return .extension
-        default: return nil
+
+    private func recordContainsSymbol(
+        store: some IndexStoreProviding,
+        recordName: String,
+        query: SymbolQuery
+    ) -> Bool {
+        guard let recordReader = try? store.recordReader(for: recordName) else {
+            return false
         }
+
+        var found = false
+        recordReader.forEachOccurrence { occurrence in
+            guard !found else { return }
+            if query.matches(occurrence.symbolMatching) {
+                found = true
+            }
+        }
+        return found
+    }
+}
+
+private final class ThreadSafeSet<Element: Hashable & Sendable>: @unchecked Sendable {
+    private let lock = NSLock()
+    private var storage = Set<Element>()
+
+    func insert(_ element: Element) {
+        lock.lock()
+        storage.insert(element)
+        lock.unlock()
+    }
+
+    func values() -> [Element] {
+        lock.lock()
+        let snapshot = Array(storage)
+        lock.unlock()
+        return snapshot
     }
 }

Sources/SwiftFindRefs/IndexStore/IndexStoreProviding.swift

@@ -0,0 +1,30 @@
+import IndexStore
+
+/// Protocol for unit dependency abstraction, enabling testability
+protocol UnitDependencyProviding {
+    var kind: DependencyKind { get }
+    var name: String { get }
+    var filePath: String { get }
+}
+
+/// Protocol for unit reader abstraction, enabling testability
+protocol UnitReaderProviding {
+    var isSystem: Bool { get }
+    func forEachDependency(_ callback: (UnitDependencyProviding) -> Void)
+}
+
+/// Protocol for index store abstraction, enabling testability
+protocol IndexStoreProviding {
+    func forEachUnit(_ callback: (UnitReaderProviding) -> Void)
+    func recordReader(for recordName: String) throws -> RecordReaderProviding?
+}
+
+/// Protocol for record reader abstraction, enabling testability
+protocol RecordReaderProviding {
+    func forEachOccurrence(_ callback: (SymbolOccurrenceProviding) -> Void)
+}
+
+/// Protocol for symbol occurrence abstraction, enabling testability
+protocol SymbolOccurrenceProviding {
+    var symbolMatching: SymbolMatching { get }
+}

Sources/SwiftFindRefs/IndexStore/RecordIndex.swift

@@ -0,0 +1,41 @@
+import IndexStore
+
+/// Maps record names to their source file paths
+struct RecordIndex {
+    let recordNames: [String]
+    private let recordToSource: [String: String]
+
+    /// Internal initializer for testing
+    init(recordNames: [String], recordToSource: [String: String] = [:]) {
+        self.recordNames = recordNames
+        self.recordToSource = recordToSource
+    }
+
+    func sourcePath(for recordName: String) -> String {
+        recordToSource[recordName] ?? recordName
+    }
+
+    static func build(from store: some IndexStoreProviding) -> RecordIndex {
+        var recordToSource: [String: String] = [:]
+        var allRecordNames = Set<String>()
+
+        store.forEachUnit { unitReader in
+            guard !unitReader.isSystem else { return }
+
+            unitReader.forEachDependency { dependency in
+                guard dependency.kind == .record else { return }
+                let recordName = dependency.name
+                allRecordNames.insert(recordName)
+                let filePath = dependency.filePath
+                if !filePath.isEmpty && recordToSource[recordName] == nil {
+                    recordToSource[recordName] = filePath
+                }
+            }
+        }
+
+        return RecordIndex(
+            recordNames: Array(allRecordNames),
+            recordToSource: recordToSource
+        )
+    }
+}

Sources/SwiftFindRefs/IndexStore/SymbolKind+Parsing.swift

@@ -0,0 +1,29 @@
+import IndexStore
+
+extension SymbolKind {
+    init?(parsing type: String) {
+        switch type.lowercased() {
+        case "class": self = .class
+        case "struct": self = .struct
+        case "enum": self = .enum
+        case "protocol": self = .protocol
+        case "function": self = .function
+        case "variable": self = .variable
+        case "typealias": self = .typealias
+        case "instancemethod": self = .instanceMethod
+        case "staticmethod": self = .staticMethod
+        case "classmethod": self = .classMethod
+        case "instanceproperty": self = .instanceProperty
+        case "staticproperty": self = .staticProperty
+        case "classproperty": self = .classProperty
+        case "constructor": self = .constructor
+        case "destructor": self = .destructor
+        case "field": self = .field
+        case "enumconstant": self = .enumConstant
+        case "parameter": self = .parameter
+        case "module": self = .module
+        case "extension": self = .extension
+        default: return nil
+        }
+    }
+}

Sources/SwiftFindRefs/IndexStore/SymbolMatching.swift

@@ -0,0 +1,9 @@
+import IndexStore
+
+/// Protocol for symbol matching, enabling testability
+protocol SymbolMatching {
+    var name: String { get }
+    var kind: SymbolKind { get }
+}
+
+extension Symbol: SymbolMatching {}

Sources/SwiftFindRefs/IndexStore/SymbolQuery.swift

@@ -0,0 +1,18 @@
+import IndexStore
+
+/// Encapsulates the search criteria for finding symbols
+struct SymbolQuery {
+    let name: String
+    let kind: SymbolKind?
+
+    init(name: String, kindString: String?) {
+        self.name = name
+        self.kind = kindString.flatMap { SymbolKind(parsing: $0) }
+    }
+
+    func matches(_ symbol: some SymbolMatching) -> Bool {
+        guard symbol.name == name else { return false }
+        if let kind, symbol.kind != kind { return false }
+        return true
+    }
+}