Support custom preparationBatchSize defined via SourceKit's options

Documentation/Configuration File.md

@@ -62,3 +62,7 @@ The structure of the file is currently not guaranteed to be stable. Options may
 - `sourcekitdRequestTimeout: number`: The maximum duration that a sourcekitd request should be allowed to execute before being declared as timed out. In general, editors should cancel requests that they are no longer interested in, but in case editors don't cancel requests, this ensures that a long-running non-cancelled request is not blocking sourcekitd and thus most semantic functionality. In particular, VS Code does not cancel the semantic tokens request, which can cause a long-running AST build that blocks sourcekitd.
 - `semanticServiceRestartTimeout: number`: If a request to sourcekitd or clangd exceeds this timeout, we assume that the semantic service provider is hanging for some reason and won't recover. To restore semantic functionality, we terminate and restart it.
 - `buildServerWorkspaceRequestsTimeout: number`: Duration how long to wait for responses to `workspace/buildTargets` or `buildTarget/sources` request by the build server before defaulting to an empty response.
+- `preparationBatchingStrategy: object`: Defines the batch size for target preparation. If nil, defaults to preparing 1 target at a time.
+  - This is a tagged union discriminated by the `strategy` field. Each case has the following structure:
+  - `strategy: "fixedTargetBatchSize"`: Prepare a fixed number of targets in a single batch. `batchSize`: The number of targets to prepare in each batch.
+    - `batchSize: integer`: The number of targets to prepare in each batch.

SourceKitLSPDevUtils/Sources/ConfigSchemaGen/JSONSchema.swift

@@ -38,6 +38,8 @@ struct JSONSchema: Encodable {
     case additionalProperties
     case markdownDescription
     case markdownEnumDescriptions
+    case oneOf
+    case const
   }
   var _schema: String?
   var id: String?
@@ -60,6 +62,9 @@ struct JSONSchema: Encodable {
   /// https://github.com/microsoft/vscode-wiki/blob/main/Setting-Descriptions.md
   var markdownEnumDescriptions: [String]?
 
+  var oneOf: [JSONSchema]?
+  var const: String?
+
   func encode(to encoder: any Encoder) throws {
     // Manually implement encoding to use `encodeIfPresent` for HeapBox-ed fields
     var container = encoder.container(keyedBy: CodingKeys.self)
@@ -82,6 +87,10 @@ struct JSONSchema: Encodable {
     if let markdownEnumDescriptions {
       try container.encode(markdownEnumDescriptions, forKey: .markdownEnumDescriptions)
     }
+    if let oneOf, !oneOf.isEmpty {
+      try container.encode(oneOf, forKey: .oneOf)
+    }
+    try container.encodeIfPresent(const, forKey: .const)
   }
 }
 
@@ -126,13 +135,53 @@ struct JSONSchemaBuilder {
       schema.properties = properties
       schema.required = required
     case .enum(let enumInfo):
-      schema.type = "string"
-      schema.enum = enumInfo.cases.map(\.name)
-      // Set `markdownEnumDescriptions` for better rendering in VSCode rich hover
-      // Unlike `description`, `enumDescriptions` field is not a part of JSON Schema spec,
-      // so we only set `markdownEnumDescriptions` here.
-      if enumInfo.cases.contains(where: { $0.description != nil }) {
-        schema.markdownEnumDescriptions = enumInfo.cases.map { $0.description ?? "" }
+      let hasAssociatedTypes = enumInfo.cases.contains { !($0.associatedProperties?.isEmpty ?? true) }
+
+      if hasAssociatedTypes {
+        let discriminatorFieldName = enumInfo.discriminatorFieldName ?? "type"
+        var oneOfSchemas: [JSONSchema] = []
+
+        for caseInfo in enumInfo.cases {
+          var caseSchema = JSONSchema()
+          caseSchema.type = "object"
+          caseSchema.description = caseInfo.description
+          caseSchema.markdownDescription = caseInfo.description
+
+          var caseProperties: [String: JSONSchema] = [:]
+          var caseRequired: [String] = [discriminatorFieldName]
+
+          var discriminatorSchema = JSONSchema()
+          discriminatorSchema.const = caseInfo.name
+          caseProperties[discriminatorFieldName] = discriminatorSchema
+
+          if let associatedProperties = caseInfo.associatedProperties {
+            for property in associatedProperties {
+              let propertyType = property.type
+              var propertySchema = try buildJSONSchema(from: propertyType)
+              propertySchema.description = property.description
+              propertySchema.markdownDescription = property.description
+              caseProperties[property.name] = propertySchema
+              if !propertyType.isOptional {
+                caseRequired.append(property.name)
+              }
+            }
+          }
+
+          caseSchema.properties = caseProperties
+          caseSchema.required = caseRequired
+          oneOfSchemas.append(caseSchema)
+        }
+
+        schema.oneOf = oneOfSchemas
+      } else {
+        schema.type = "string"
+        schema.enum = enumInfo.cases.map(\.name)
+        // Set `markdownEnumDescriptions` for better rendering in VSCode rich hover
+        // Unlike `description`, `enumDescriptions` field is not a part of JSON Schema spec,
+        // so we only set `markdownEnumDescriptions` here.
+        if enumInfo.cases.contains(where: { $0.description != nil }) {
+          schema.markdownEnumDescriptions = enumInfo.cases.map { $0.description ?? "" }
+        }
       }
     }
     return schema

SourceKitLSPDevUtils/Sources/ConfigSchemaGen/OptionDocument.swift

@@ -68,14 +68,39 @@ struct OptionDocumentBuilder {
           try appendProperty(property, indentLevel: indentLevel + 1)
         }
       case .enum(let schema):
-        for caseInfo in schema.cases {
-          // Add detailed description for each case if available
-          guard let description = caseInfo.description else {
-            continue
+        let hasAssociatedTypes = schema.cases.contains {
+          $0.associatedProperties != nil && !$0.associatedProperties!.isEmpty
+        }
+
+        if hasAssociatedTypes {
+          let discriminatorFieldName = schema.discriminatorFieldName ?? "type"
+          doc +=
+            "\(indent)  - This is a tagged union discriminated by the `\(discriminatorFieldName)` field. Each case has the following structure:\n"
+
+          for caseInfo in schema.cases {
+            doc += """
+              \(indent)  - `\(discriminatorFieldName): "\(caseInfo.name)"`
+              """
+            if let description = caseInfo.description {
+              doc += ": " + description.split(separator: "\n").joined(separator: "\n\(indent)    ")
+            }
+            doc += "\n"
+
+            if let associatedProperties = caseInfo.associatedProperties {
+              for assocProp in associatedProperties {
+                try appendProperty(assocProp, indentLevel: indentLevel + 2)
+              }
+            }
+          }
+        } else {
+          for caseInfo in schema.cases {
+            guard let description = caseInfo.description else {
+              continue
+            }
+            doc += "\(indent)  - `\(caseInfo.name)`"
+            doc += ": " + description.split(separator: "\n").joined(separator: "\n\(indent)    ")
+            doc += "\n"
           }
-          doc += "\(indent)  - `\(caseInfo.name)`"
-          doc += ": " + description.split(separator: "\n").joined(separator: "\n\(indent)    ")
-          doc += "\n"
         }
       default: break
       }
@@ -102,8 +127,15 @@ struct OptionDocumentBuilder {
     case .struct(let structInfo):
       return structInfo.name
     case .enum(let enumInfo):
-      let cases = enumInfo.cases.map { "\"\($0.name)\"" }.joined(separator: "|")
-      return shouldWrap ? "(\(cases))" : cases
+      let hasAssociatedTypes = enumInfo.cases.contains {
+        $0.associatedProperties != nil && !$0.associatedProperties!.isEmpty
+      }
+      if hasAssociatedTypes {
+        return "object"
+      } else {
+        let cases = enumInfo.cases.map { "\"\($0.name)\"" }.joined(separator: "|")
+        return shouldWrap ? "(\(cases))" : cases
+      }
     }
   }
 }

SourceKitLSPDevUtils/Sources/ConfigSchemaGen/OptionSchema.swift

@@ -31,11 +31,13 @@ struct OptionTypeSchama {
   struct Case {
     var name: String
     var description: String?
+    var associatedProperties: [Property]?
   }
 
   struct Enum {
     var name: String
     var cases: [Case]
+    var discriminatorFieldName: String?
   }
 
   enum Kind {
@@ -146,14 +148,13 @@ struct OptionSchemaContext {
   }
 
   private func buildEnumCases(_ node: EnumDeclSyntax) throws -> OptionTypeSchama.Enum {
+    let discriminatorFieldName = Self.extractDiscriminatorFieldName(node.leadingTrivia)
+
     let cases = try node.memberBlock.members.flatMap { member -> [OptionTypeSchama.Case] in
       guard let caseDecl = member.decl.as(EnumCaseDeclSyntax.self) else {
         return []
       }
       return try caseDecl.elements.compactMap {
-        guard $0.parameterClause == nil else {
-          throw ConfigSchemaGenError("Associated values in enum cases are not supported: \(caseDecl)")
-        }
         let name: String
         if let rawValue = $0.rawValue?.value {
           if let stringLiteral = rawValue.as(StringLiteralExprSyntax.self),
@@ -172,11 +173,45 @@ struct OptionSchemaContext {
         if description?.contains("- Note: Internal option") ?? false {
           return nil
         }
-        return OptionTypeSchama.Case(name: name, description: description)
+
+        var associatedProperties: [OptionTypeSchama.Property]? = nil
+        if let parameterClause = $0.parameterClause {
+          let caseDescription = description
+          associatedProperties = try parameterClause.parameters.map { param in
+            let propertyName: String
+            if let firstName = param.firstName, firstName.tokenKind != .wildcard {
+              propertyName = firstName.text
+            } else if let secondName = param.secondName {
+              propertyName = secondName.text
+            } else {
+              propertyName = name
+            }
+
+            let propertyType = try resolveType(param.type)
+            let propertyDescription =
+              Self.extractParameterDescription(
+                from: caseDescription,
+                parameterName: propertyName
+              ) ?? Self.extractDocComment(param.leadingTrivia)
+
+            return OptionTypeSchama.Property(
+              name: propertyName,
+              type: propertyType,
+              description: propertyDescription,
+              defaultValue: nil
+            )
+          }
+        }
+
+        return OptionTypeSchama.Case(
+          name: name,
+          description: description,
+          associatedProperties: associatedProperties
+        )
       }
     }
     let typeName = node.name.text
-    return .init(name: typeName, cases: cases)
+    return .init(name: typeName, cases: cases, discriminatorFieldName: discriminatorFieldName)
   }
 
   private func buildStructProperties(_ node: StructDeclSyntax) throws -> OptionTypeSchama.Struct {
@@ -234,4 +269,52 @@ struct OptionSchemaContext {
     }
     return docLines.joined(separator: " ")
   }
+
+  private static func extractDiscriminatorFieldName(_ trivia: Trivia) -> String? {
+    let docLines = trivia.flatMap { piece -> [Substring] in
+      switch piece {
+      case .docBlockComment(let text):
+        assert(text.hasPrefix("/**") && text.hasSuffix("*/"), "Unexpected doc block comment format: \(text)")
+        return text.dropFirst(3).dropLast(2).split { $0.isNewline }
+      case .docLineComment(let text):
+        assert(text.hasPrefix("///"), "Unexpected doc line comment format: \(text)")
+        let text = text.dropFirst(3)
+        return [text]
+      default:
+        return []
+      }
+    }
+
+    for line in docLines {
+      var trimmed = line
+      while trimmed.first?.isWhitespace == true {
+        trimmed = trimmed.dropFirst()
+      }
+      if trimmed.hasPrefix("- discriminator:") {
+        let fieldName = trimmed.dropFirst("- discriminator:".count).trimmingCharacters(in: .whitespaces)
+        return fieldName.isEmpty ? nil : fieldName
+      }
+    }
+    return nil
+  }
+
+  private static func extractParameterDescription(from docComment: String?, parameterName: String) -> String? {
+    guard let docComment = docComment else {
+      return nil
+    }
+
+    let pattern = "`\(parameterName)`:"
+    guard let range = docComment.range(of: pattern) else {
+      return nil
+    }
+
+    let afterPattern = docComment[range.upperBound...]
+    let lines = afterPattern.split(separator: "\n", maxSplits: 1, omittingEmptySubsequences: false)
+    guard let firstLine = lines.first else {
+      return nil
+    }
+
+    let description = firstLine.trimmingCharacters(in: .whitespaces)
+    return description.isEmpty ? nil : description
+  }
 }

Sources/SKOptions/PreparationBatchingStrategy.swift

@@ -0,0 +1,54 @@
+//===----------------------------------------------------------------------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+/// Defines the batch size for target preparation.
+///
+/// If nil, SourceKit-LSP will default to preparing 1 target at a time.
+///
+/// - discriminator: strategy
+public enum PreparationBatchingStrategy: Sendable, Equatable {
+  /// Prepare a fixed number of targets in a single batch.
+  ///
+  /// `batchSize`: The number of targets to prepare in each batch.
+  case fixedTargetBatchSize(batchSize: Int)
+}
+
+extension PreparationBatchingStrategy: Codable {
+  private enum CodingKeys: String, CodingKey {
+    case strategy
+    case batchSize
+  }
+
+  private enum StrategyValue: String, Codable {
+    case fixedTargetBatchSize
+  }
+
+  public init(from decoder: Decoder) throws {
+    let container = try decoder.container(keyedBy: CodingKeys.self)
+    let strategy = try container.decode(StrategyValue.self, forKey: .strategy)
+
+    switch strategy {
+    case .fixedTargetBatchSize:
+      let batchSize = try container.decode(Int.self, forKey: .batchSize)
+      self = .fixedTargetBatchSize(batchSize: batchSize)
+    }
+  }
+
+  public func encode(to encoder: Encoder) throws {
+    var container = encoder.container(keyedBy: CodingKeys.self)
+    switch self {
+    case .fixedTargetBatchSize(let batchSize):
+      try container.encode(StrategyValue.fixedTargetBatchSize, forKey: .strategy)
+      try container.encode(batchSize, forKey: .batchSize)
+    }
+  }
+}

Sources/SKOptions/SourceKitLSPOptions.swift

@@ -449,6 +449,10 @@ public struct SourceKitLSPOptions: Sendable, Codable, Equatable {
     return .seconds(15)
   }
 
+  /// Defines the batch size for target preparation.
+  /// If nil, defaults to preparing 1 target at a time.
+  public var preparationBatchingStrategy: PreparationBatchingStrategy?
+
   public init(
     swiftPM: SwiftPMOptions? = .init(),
     fallbackBuildSystem: FallbackBuildSystemOptions? = .init(),
@@ -462,6 +466,7 @@ public struct SourceKitLSPOptions: Sendable, Codable, Equatable {
     generatedFilesPath: String? = nil,
     backgroundIndexing: Bool? = nil,
     backgroundPreparationMode: BackgroundPreparationMode? = nil,
+    preparationBatchingStrategy: PreparationBatchingStrategy? = nil,
     cancelTextDocumentRequestsOnEditAndClose: Bool? = nil,
     experimentalFeatures: Set<ExperimentalFeature>? = nil,
     swiftPublishDiagnosticsDebounceDuration: Double? = nil,
@@ -482,6 +487,7 @@ public struct SourceKitLSPOptions: Sendable, Codable, Equatable {
     self.defaultWorkspaceType = defaultWorkspaceType
     self.backgroundIndexing = backgroundIndexing
     self.backgroundPreparationMode = backgroundPreparationMode
+    self.preparationBatchingStrategy = preparationBatchingStrategy
     self.cancelTextDocumentRequestsOnEditAndClose = cancelTextDocumentRequestsOnEditAndClose
     self.experimentalFeatures = experimentalFeatures
     self.swiftPublishDiagnosticsDebounceDuration = swiftPublishDiagnosticsDebounceDuration
@@ -545,6 +551,7 @@ public struct SourceKitLSPOptions: Sendable, Codable, Equatable {
       generatedFilesPath: override?.generatedFilesPath ?? base.generatedFilesPath,
       backgroundIndexing: override?.backgroundIndexing ?? base.backgroundIndexing,
       backgroundPreparationMode: override?.backgroundPreparationMode ?? base.backgroundPreparationMode,
+      preparationBatchingStrategy: override?.preparationBatchingStrategy ?? base.preparationBatchingStrategy,
       cancelTextDocumentRequestsOnEditAndClose: override?.cancelTextDocumentRequestsOnEditAndClose
         ?? base.cancelTextDocumentRequestsOnEditAndClose,
       experimentalFeatures: override?.experimentalFeatures ?? base.experimentalFeatures,