Implement a syntactic workspace-wide test index

Sources/SKCore/BuildServerBuildSystem.swift

@@ -316,6 +316,17 @@ extension BuildServerBuildSystem: BuildSystem {
 
     return .unhandled
   }
+
+  public func testFiles() async -> [DocumentURI] {
+    // BuildServerBuildSystem does not support syntactic test discovery
+    // (https://github.com/apple/sourcekit-lsp/issues/1173).
+    return []
+  }
+
+  public func addTestFilesDidChangeCallback(_ callback: @escaping () async -> Void) {
+    // BuildServerBuildSystem does not support syntactic test discovery
+    // (https://github.com/apple/sourcekit-lsp/issues/1173).
+  }
 }
 
 private func loadBuildServerConfig(path: AbsolutePath, fileSystem: FileSystem) throws -> BuildServerConfig {

Sources/SKCore/BuildSystem.swift

@@ -87,6 +87,18 @@ public protocol BuildSystem: AnyObject, Sendable {
   func filesDidChange(_ events: [FileEvent]) async
 
   func fileHandlingCapability(for uri: DocumentURI) async -> FileHandlingCapability
+
+  /// Returns the list of files that might contain test cases.
+  ///
+  /// The returned file list is an over-approximation. It might contain tests from non-test targets or files that don't
+  /// actually contain any tests. Keeping this list as minimal as possible helps reduce the amount of work that the
+  /// syntactic test indexer needs to perform.
+  func testFiles() async -> [DocumentURI]
+
+  /// Adds a callback that should be called when the value returned by `testFiles()` changes.
+  ///
+  /// The callback might also be called without an actual change to `testFiles`.
+  func addTestFilesDidChangeCallback(_ callback: @Sendable @escaping () async -> Void) async
 }
 
 public let buildTargetsNotSupported =

Sources/SKCore/BuildSystemManager.swift

@@ -176,6 +176,10 @@ extension BuildSystemManager {
       fallbackBuildSystem != nil ? .fallback : .unhandled
     )
   }
+
+  public func testFiles() async -> [DocumentURI] {
+    return await buildSystem?.testFiles() ?? []
+  }
 }
 
 extension BuildSystemManager: BuildSystemDelegate {

Sources/SKCore/CompilationDatabaseBuildSystem.swift

@@ -39,6 +39,9 @@ public actor CompilationDatabaseBuildSystem {
   /// Delegate to handle any build system events.
   public weak var delegate: BuildSystemDelegate? = nil
 
+  /// Callbacks that should be called if the list of possible test files has changed.
+  public var testFilesDidChangeCallbacks: [() async -> Void] = []
+
   public func setDelegate(_ delegate: BuildSystemDelegate?) async {
     self.delegate = delegate
   }
@@ -167,6 +170,9 @@ extension CompilationDatabaseBuildSystem: BuildSystem {
     if let delegate = self.delegate {
       await delegate.fileBuildSettingsChanged(self.watchedFiles)
     }
+    for testFilesDidChangeCallback in testFilesDidChangeCallbacks {
+      await testFilesDidChangeCallback()
+    }
   }
 
   public func filesDidChange(_ events: [FileEvent]) async {
@@ -185,4 +191,12 @@ extension CompilationDatabaseBuildSystem: BuildSystem {
       return .unhandled
     }
   }
+
+  public func testFiles() async -> [DocumentURI] {
+    return compdb?.allCommands.map { DocumentURI($0.url) } ?? []
+  }
+
+  public func addTestFilesDidChangeCallback(_ callback: @escaping () async -> Void) async {
+    testFilesDidChangeCallbacks.append(callback)
+  }
 }

Sources/SKSupport/AsyncQueue.swift

@@ -16,6 +16,8 @@ import Foundation
 /// array.
 private protocol AnyTask: Sendable {
   func waitForCompletion() async
+
+  func cancel()
 }
 
 extension Task: AnyTask {
@@ -89,6 +91,16 @@ public final class AsyncQueue<TaskMetadata: DependencyTracker>: Sendable {
 
   public init() {}
 
+  public func cancelTasks(where filter: (TaskMetadata) -> Bool) {
+    pendingTasks.withLock { pendingTasks in
+      for task in pendingTasks {
+        if filter(task.metadata) {
+          task.task.cancel()
+        }
+      }
+    }
+  }
+
   /// Schedule a new closure to be executed on the queue.
   ///
   /// If this is a serial queue, all previously added tasks are guaranteed to

Sources/SKSwiftPMWorkspace/SwiftPMBuildSystem.swift

@@ -81,6 +81,9 @@ public actor SwiftPMBuildSystem {
     self.delegate = delegate
   }
 
+  /// Callbacks that should be called if the list of possible test files has changed.
+  public var testFilesDidChangeCallbacks: [() async -> Void] = []
+
   let workspacePath: TSCAbsolutePath
   /// The directory containing `Package.swift`.
   public var projectRoot: TSCAbsolutePath
@@ -267,6 +270,9 @@ extension SwiftPMBuildSystem {
     }
     await delegate.fileBuildSettingsChanged(self.watchedFiles)
     await delegate.fileHandlingCapabilityChanged()
+    for testFilesDidChangeCallback in testFilesDidChangeCallbacks {
+      await testFilesDidChangeCallback()
+    }
   }
 }
 
@@ -380,6 +386,15 @@ extension SwiftPMBuildSystem: SKCore.BuildSystem {
       return .unhandled
     }
   }
+
+  public func testFiles() -> [DocumentURI] {
+    // We should only include source files from test targets (https://github.com/apple/sourcekit-lsp/issues/1174).
+    return fileToTarget.map { DocumentURI($0.key.asURL) }
+  }
+
+  public func addTestFilesDidChangeCallback(_ callback: @Sendable @escaping () async -> Void) async {
+    testFilesDidChangeCallbacks.append(callback)
+  }
 }
 
 extension SwiftPMBuildSystem {

Sources/SKTestSupport/IndexedSingleSwiftFileTestProject.swift

@@ -68,6 +68,25 @@ public struct IndexedSingleSwiftFileTestProject {
 
     if let sdk = TibsBuilder.defaultSDKPath {
       compilerArguments += ["-sdk", sdk]
+
+      // The following are needed so we can import XCTest
+      let sdkUrl = URL(fileURLWithPath: sdk)
+      let usrLibDir =
+        sdkUrl
+        .deletingLastPathComponent()
+        .deletingLastPathComponent()
+        .appendingPathComponent("usr")
+        .appendingPathComponent("lib")
+      let frameworksDir =
+        sdkUrl
+        .deletingLastPathComponent()
+        .deletingLastPathComponent()
+        .appendingPathComponent("Library")
+        .appendingPathComponent("Frameworks")
+      compilerArguments += [
+        "-I", usrLibDir.path,
+        "-F", frameworksDir.path,
+      ]
     }
 
     let compilationDatabase = JSONCompilationDatabase(

Sources/SKTestSupport/MultiFileTestProject.swift

@@ -142,4 +142,13 @@ public class MultiFileTestProject {
     }
     return DocumentPositions(markedText: fileData.markedText)[marker]
   }
+
+  public func range(from fromMarker: String, to toMarker: String, in fileName: String) throws -> Range<Position> {
+    return try position(of: fromMarker, in: fileName)..<position(of: toMarker, in: fileName)
+  }
+
+  public func location(from fromMarker: String, to toMarker: String, in fileName: String) throws -> Location {
+    let range = try self.range(from: fromMarker, to: toMarker, in: fileName)
+    return Location(uri: try self.uri(for: fileName), range: range)
+  }
 }

Sources/SKTestSupport/SwiftPMTestProject.swift

@@ -40,6 +40,7 @@ public class SwiftPMTestProject: MultiFileTestProject {
     manifest: String = SwiftPMTestProject.defaultPackageManifest,
     workspaces: (URL) -> [WorkspaceFolder] = { [WorkspaceFolder(uri: DocumentURI($0))] },
     build: Bool = false,
+    allowBuildFailure: Bool = false,
     usePullDiagnostics: Bool = true,
     testName: String = #function
   ) async throws {
@@ -66,7 +67,11 @@ public class SwiftPMTestProject: MultiFileTestProject {
     )
 
     if build {
-      try await Self.build(at: self.scratchDirectory)
+      if allowBuildFailure {
+        try? await Self.build(at: self.scratchDirectory)
+      } else {
+        try await Self.build(at: self.scratchDirectory)
+      }
     }
     // Wait for the indexstore-db to finish indexing
     _ = try await testClient.send(PollIndexRequest())

Sources/SKTestSupport/TestSourceKitLSPClient.swift

@@ -362,7 +362,7 @@ public struct DocumentPositions {
     }
   }
 
-  init(markedText: String) {
+  public init(markedText: String) {
     let (markers, textWithoutMarker) = extractMarkers(markedText)
     self.init(markers: markers, textWithoutMarkers: textWithoutMarker)
   }

Sources/SourceKitLSP/CMakeLists.txt

@@ -2,6 +2,7 @@
 add_library(SourceKitLSP STATIC
   CapabilityRegistry.swift
   DocumentManager.swift
+  DocumentSnapshot+FromFileContents.swift
   IndexOutOfDateChecker.swift
   IndexStoreDB+MainFilesProvider.swift
   LanguageService.swift
@@ -12,6 +13,7 @@ add_library(SourceKitLSP STATIC
   SourceKitLSPCommandMetadata.swift
   SourceKitLSPServer.swift
   SourceKitLSPServer+Options.swift
+  SymbolLocation+DocumentURI.swift
   TestDiscovery.swift
   WorkDoneProgressManager.swift
   Workspace.swift
@@ -44,9 +46,10 @@ target_sources(SourceKitLSP PRIVATE
   Swift/SwiftLanguageService.swift
   Swift/SwiftTestingScanner.swift
   Swift/SymbolInfo.swift
+  Swift/SyntacticTestIndex.swift
   Swift/SyntaxHighlightingToken.swift
-  Swift/SyntaxHighlightingTokens.swift
   Swift/SyntaxHighlightingTokenParser.swift
+  Swift/SyntaxHighlightingTokens.swift
   Swift/SyntaxTreeManager.swift
   Swift/VariableTypeInfo.swift
 )

Sources/SourceKitLSP/DocumentSnapshot+FromFileContents.swift

@@ -0,0 +1,35 @@
+//===----------------------------------------------------------------------===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2014 - 2024 Apple Inc. and the Swift project authors
+// Licensed under Apache License v2.0 with Runtime Library Exception
+//
+// See https://swift.org/LICENSE.txt for license information
+// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
+//
+//===----------------------------------------------------------------------===//
+
+import LanguageServerProtocol
+import SKSupport
+
+public extension DocumentSnapshot {
+  /// Creates a `DocumentSnapshot` with the file contents from disk.
+  ///
+  /// Throws an error if the file could not be read.
+  /// Returns `nil` if the `uri` is not a file URL.
+  init?(withContentsFromDisk uri: DocumentURI, language: Language) throws {
+    guard let url = uri.fileURL else {
+      return nil
+    }
+    try self.init(withContentsFromDisk: url, language: language)
+  }
+
+  /// Creates a `DocumentSnapshot` with the file contents from disk.
+  ///
+  /// Throws an error if the file could not be read.
+  init(withContentsFromDisk url: URL, language: Language) throws {
+    let contents = try String(contentsOf: url)
+    self.init(uri: DocumentURI(url), language: language, version: 0, lineTable: LineTable(contents))
+  }
+}

Sources/SourceKitLSP/IndexOutOfDateChecker.swift

@@ -13,13 +13,19 @@
 import Foundation
 import IndexStoreDB
 import LSPLogging
+import LanguageServerProtocol
 
 /// Helper class to check if symbols from the index are up-to-date or if the source file has been modified after it was
-/// indexed.
+/// indexed. Modifications include both changes to the file on disk as well as modifications to the file that have not
+/// been saved to disk (ie. changes that only live in `DocumentManager`).
 ///
 /// The checker caches mod dates of source files. It should thus not be long lived. Its intended lifespan is the
 /// evaluation of a single request.
 struct IndexOutOfDateChecker {
+  /// The `DocumentManager` that holds the in-memory file contents. We consider the index out-of-date for all files that
+  /// have in-memory changes.
+  private let documentManager: DocumentManager
+
   /// The last modification time of a file. Can also represent the fact that the file does not exist.
   private enum ModificationTime {
     case fileDoesNotExist
@@ -37,12 +43,19 @@ struct IndexOutOfDateChecker {
     }
   }
 
-  /// File paths to modification times that have already been computed.
-  private var modTimeCache: [String: ModificationTime] = [:]
+  /// File URLs to modification times that have already been computed.
+  private var modTimeCache: [URL: ModificationTime] = [:]
+
+  /// Caches whether a file URL has modifications in `documentManager` that haven't been saved to disk yet.
+  private var hasFileInMemoryModificationsCache: [URL: Bool] = [:]
 
-  private func modificationDateUncached(of path: String) throws -> ModificationTime {
+  init(documentManager: DocumentManager) {
+    self.documentManager = documentManager
+  }
+
+  private func modificationDateUncached(of url: URL) throws -> ModificationTime {
     do {
-      let attributes = try FileManager.default.attributesOfItem(atPath: path)
+      let attributes = try FileManager.default.attributesOfItem(atPath: url.path)
       guard let modificationDate = attributes[FileAttributeKey.modificationDate] as? Date else {
         throw Error.fileAttributesDontHaveModificationDate
       }
@@ -52,20 +65,46 @@ struct IndexOutOfDateChecker {
     }
   }
 
-  private mutating func modificationDate(of path: String) throws -> ModificationTime {
-    if let cached = modTimeCache[path] {
+  private mutating func modificationDate(of url: URL) throws -> ModificationTime {
+    if let cached = modTimeCache[url] {
       return cached
     }
-    let modTime = try modificationDateUncached(of: path)
-    modTimeCache[path] = modTime
+    let modTime = try modificationDateUncached(of: url)
+    modTimeCache[url] = modTime
     return modTime
   }
 
+  private func hasFileInMemoryModificationsUncached(at url: URL) -> Bool {
+    guard let document = try? documentManager.latestSnapshot(DocumentURI(url)) else {
+      return false
+    }
+
+    guard let onDiskFileContents = try? String(contentsOf: url, encoding: .utf8) else {
+      // If we can't read the file on disk, it's an in-memory document
+      return true
+    }
+    return onDiskFileContents != document.lineTable.content
+  }
+
+  /// Returns `true` if the file has modified in-memory state, ie. if the version stored in the `DocumentManager` is
+  /// different than the version on disk.
+  public mutating func fileHasInMemoryModifications(_ url: URL) -> Bool {
+    if let cached = hasFileInMemoryModificationsCache[url] {
+      return cached
+    }
+    let hasInMemoryModifications = hasFileInMemoryModificationsUncached(at: url)
+    hasFileInMemoryModificationsCache[url] = hasInMemoryModifications
+    return hasInMemoryModifications
+  }
+
   /// Returns `true` if the source file for the given symbol location exists and has not been modified after it has been
   /// indexed.
   mutating func isUpToDate(_ symbolLocation: SymbolLocation) -> Bool {
+    if fileHasInMemoryModifications(URL(fileURLWithPath: symbolLocation.path)) {
+      return false
+    }
     do {
-      let sourceFileModificationDate = try modificationDate(of: symbolLocation.path)
+      let sourceFileModificationDate = try modificationDate(of: URL(fileURLWithPath: symbolLocation.path))
       switch sourceFileModificationDate {
       case .fileDoesNotExist:
         return false
@@ -81,8 +120,11 @@ struct IndexOutOfDateChecker {
   /// Return `true` if a unit file has been indexed for the given file path after its last modification date.
   ///
   /// This means that at least a single build configuration of this file has been indexed since its last modification.
-  mutating func indexHasUpToDateUnit(for filePath: String, index: IndexStoreDB) -> Bool {
-    guard let lastUnitDate = index.dateOfLatestUnitFor(filePath: filePath) else {
+  mutating func indexHasUpToDateUnit(for filePath: URL, index: IndexStoreDB) -> Bool {
+    if fileHasInMemoryModifications(filePath) {
+      return false
+    }
+    guard let lastUnitDate = index.dateOfLatestUnitFor(filePath: filePath.path) else {
       return false
     }
     do {