Refactor HybridValueProvider to use isolated parameters

Package.resolved

@@ -14,7 +14,7 @@
       "kind" : "remoteSourceControl",
       "location" : "https://github.com/ChimeHQ/SwiftTreeSitter",
       "state" : {
-        "revision" : "87ed52a71d4ad6b5e6a11185b42f6f74eb5b47da"
+        "revision" : "10cb68c00a9963c2884b30f168a9de377790d812"
       }
     }
   ],

Package.swift

@@ -15,7 +15,7 @@ let package = Package(
 		.library(name: "Neon", targets: ["Neon"]),
 	],
 	dependencies: [
-		.package(url: "https://github.com/ChimeHQ/SwiftTreeSitter", revision: "87ed52a71d4ad6b5e6a11185b42f6f74eb5b47da"),
+		.package(url: "https://github.com/ChimeHQ/SwiftTreeSitter", revision: "10cb68c00a9963c2884b30f168a9de377790d812"),
 		.package(url: "https://github.com/ChimeHQ/Rearrange", from: "1.8.1"),
 	],
 	targets: [

README.md

@@ -57,7 +57,7 @@ Neon's lowest-level component is called RangeState. This module contains the cor
 
 Many of these support versionable content. If you are working with a backing store structure that supports efficient versioning, like a [piece table](https://en.wikipedia.org/wiki/Piece_table), expressing this to RangeState can improve its efficiency.
 
-It might be surprising to see that `RangeValidator` and `RangeProcessor` are marked `@MainActor`. Right now, I have found no way to both support the hybrid sync/async functionality while also not being tied to a global actor. I think this is the most resonable trade-off, but I would very much like to lift this restriction. If you have ideas, I'd love to hear them.
+It might be surprising to see that many of the types in RangeState are marked `@MainActor`. Right now, I have found no way to both support the hybrid sync/async functionality while also not being tied to a global actor. I think this is the most resonable trade-off, but I would very much like to lift this restriction. However, I believe it will require [language changes](https://forums.swift.org/t/isolation-assumptions/69514/47).
 
 ### TreeSitterClient
 

Sources/Neon/TextSystemStyler.swift

@@ -44,7 +44,7 @@ public final class TextSystemStyler<Interface: TextSystemInterface> {
 	private var validationProvider: Validator.ValidationProvider {
 		.init(
 			syncValue: { [weak self] in self?.validate($0) },
-			asyncValue: { [weak self] in await self?.validate($0) ?? .stale }
+			asyncValue: { [weak self] range, _ in await self?.validate(range) ?? .stale }
 		)
 	}
 
@@ -67,7 +67,8 @@ public final class TextSystemStyler<Interface: TextSystemInterface> {
 	private func validate(_ range: Validator.ContentRange) async -> Validator.Validation {
 		guard range.version == currentVersion else { return .stale }
 
-		let application = await configuration.tokenProvider.async(range.value)
+		// https://github.com/apple/swift/pull/71143
+		let application = await configuration.tokenProvider.mainActorAsync(range.value)
 
 		applyStyles(for: application)
 

Sources/Neon/Token.swift

@@ -45,7 +45,7 @@ extension TokenProvider {
 			syncValue: { _ in
 				return TokenApplication(tokens: [])
 			},
-			asyncValue: { _ in
+			asyncValue: { _, _ in
 				return TokenApplication(tokens: [])
 			}
 		)

Sources/Neon/TreeSitterClient+Neon.swift

@@ -5,26 +5,43 @@ import SwiftTreeSitter
 import SwiftTreeSitterLayer
 import TreeSitterClient
 
+extension TokenApplication {
+	public init(namedRanges: [NamedRange], nameMap: [String : String], range: NSRange) {
+		let tokens = namedRanges.map {
+			let name = nameMap[$0.name] ?? $0.name
+
+			return Token(name: name, range: $0.range)
+		}
+
+		self.init(tokens: tokens, range: range)
+	}
+}
+
 extension TreeSitterClient {
-	public func tokenProvider(with provider: @escaping TextProvider) -> TokenProvider {
-		HybridValueProvider<NSRange, [NamedRange]>(
+	@MainActor
+	public func tokenProvider(with provider: @escaping TextProvider, nameMap: [String : String] = [:]) -> TokenProvider {
+		TokenProvider(
 			syncValue: { [highlightsProvider] range in
 				do {
-					return try highlightsProvider.sync(.init(range: range, textProvider: provider))
+					guard let namedRanges = try highlightsProvider.sync(.init(range: range, textProvider: provider)) else {
+						return nil
+					}
+
+					return TokenApplication(namedRanges: namedRanges, nameMap: nameMap, range: range)
 				} catch {
 					return []
 				}
 			},
-			asyncValue: { [highlightsProvider] range in
+			mainActorAsyncValue: { [highlightsProvider] range in
 				do {
-					return try await highlightsProvider.async(.init(range: range, textProvider: provider))
+					let namedRanges = try await highlightsProvider.mainActorAsync(.init(range: range, textProvider: provider))
+
+					return TokenApplication(namedRanges: namedRanges, nameMap: nameMap, range: range)
 				} catch {
 					return []
 				}
 			}
-		).map { namedRanges in
-			TokenApplication(tokens: namedRanges.map({ Token(name: $0.name, range: $0.range) }))
-		}
+		)
 	}
 }
 

Sources/RangeState/HybridValueProvider+RangeProcessor.swift

@@ -7,7 +7,7 @@ extension HybridValueProvider {
 		rangeProcessor: RangeProcessor,
 		inputTransformer: @escaping (Input) -> (Int, RangeFillMode),
 		syncValue: @escaping SyncValueProvider,
-		asyncValue: @escaping AsyncValueProvider
+		asyncValue: @escaping @MainActor (Input) async -> Output
 	) {
 		self.init(
 			syncValue: { input in
@@ -19,11 +19,10 @@ extension HybridValueProvider {
 
 				return nil
 			},
-			asyncValue: { input in
+			asyncValue: { input, actor in
 				let (location, fill) = inputTransformer(input)
 
-				rangeProcessor.processLocation(location, mode: fill)
-
+				await rangeProcessor.processLocation(location, mode: fill)
 				await rangeProcessor.processingCompleted()
 
 				return await asyncValue(input)
@@ -39,7 +38,7 @@ extension HybridThrowingValueProvider {
 		rangeProcessor: RangeProcessor,
 		inputTransformer: @escaping (Input) -> (Int, RangeFillMode),
 		syncValue: @escaping SyncValueProvider,
-		asyncValue: @escaping AsyncValueProvider
+		asyncValue: @escaping @MainActor (Input) async throws -> Output
 	) {
 		self.init(
 			syncValue: { input in
@@ -51,11 +50,10 @@ extension HybridThrowingValueProvider {
 
 				return nil
 			},
-			asyncValue: { input in
+			asyncValue: { input, actor in
 				let (location, fill) = inputTransformer(input)
 
-				rangeProcessor.processLocation(location, mode: fill)
-
+				await rangeProcessor.processLocation(location, mode: fill)
 				await rangeProcessor.processingCompleted()
 
 				return try await asyncValue(input)

Sources/RangeState/HybridValueProvider.swift

@@ -3,7 +3,7 @@ import Foundation
 /// A type that can perform work both synchronously and asynchronously.
 public struct HybridValueProvider<Input: Sendable, Output: Sendable> {
 	public typealias SyncValueProvider = (Input) -> Output?
-	public typealias AsyncValueProvider = (Input) async -> Output
+	public typealias AsyncValueProvider = (Input, isolated any Actor) async -> Output
 
 	public let syncValueProvider: SyncValueProvider
 	public let asyncValueProvider: AsyncValueProvider
@@ -16,19 +16,39 @@ public struct HybridValueProvider<Input: Sendable, Output: Sendable> {
 		self.asyncValueProvider = asyncValue
 	}
 
-	public func async(_ input: Input) async -> Output {
-		await asyncValueProvider(input)
+	public func async(_ input: Input, isolatedTo actor: any Actor) async -> Output {
+		await asyncValueProvider(input, actor)
 	}
 
+
 	public func sync(_ input: Input) -> Output? {
 		syncValueProvider(input)
 	}
 }
 
+extension HybridValueProvider {
+	/// Create an instance that can statically prove to the compiler that asyncValueProvider is isolated to the MainActor.
+	public init(
+		syncValue: @escaping SyncValueProvider = { _ in nil },
+		mainActorAsyncValue: @escaping @MainActor (Input) async -> Output
+	) {
+		self.syncValueProvider = syncValue
+		self.asyncValueProvider = { input, _ in
+			return await mainActorAsyncValue(input)
+		}
+	}
+
+	/// Hopefully temporary until https://github.com/apple/swift/pull/71143 is available.
+	@MainActor
+	public func mainActorAsync(_ input: Input) async -> Output {
+		await asyncValueProvider(input, MainActor.shared)
+	}
+}
+
 /// A type that can perform failable work both synchronously and asynchronously.
 public struct HybridThrowingValueProvider<Input: Sendable, Output: Sendable> {
 	public typealias SyncValueProvider = (Input) throws -> Output?
-	public typealias AsyncValueProvider = (Input) async throws -> Output
+	public typealias AsyncValueProvider = (Input, isolated any Actor) async throws -> Output
 
 	public let syncValueProvider: SyncValueProvider
 	public let asyncValueProvider: AsyncValueProvider
@@ -41,59 +61,79 @@ public struct HybridThrowingValueProvider<Input: Sendable, Output: Sendable> {
 		self.asyncValueProvider = asyncValue
 	}
 
-	public func async(_ input: Input) async throws -> Output {
-		try await asyncValueProvider(input)
+	public func async(_ input: Input, isolatedTo actor: any Actor) async throws -> Output {
+		try await asyncValueProvider(input, actor)
 	}
 
 	public func sync(_ input: Input) throws -> Output? {
 		try syncValueProvider(input)
 	}
 }
 
-extension HybridValueProvider {
-	/// Returns a new `HybridValueProvider` with a new output type.
-	public func map<T>(_ transform: @escaping (Output) -> T) -> HybridValueProvider<Input, T> where T: Sendable {
-		.init(
-			syncValue: { self.sync($0).map(transform) },
-			asyncValue: { transform(await self.async($0)) }
-		)
+extension HybridThrowingValueProvider {
+	/// Create an instance that can statically prove to the compiler that asyncValueProvider is isolated to the MainActor.
+	public init(
+		syncValue: @escaping SyncValueProvider = { _ in nil },
+		mainActorAsyncValue: @escaping @MainActor (Input) async -> Output
+	) {
+		self.syncValueProvider = syncValue
+		self.asyncValueProvider = { input, _ in
+			return await mainActorAsyncValue(input)
+		}
 	}
 
-	/// Convert to a `HybridThrowingValueProvider`.
-	public var throwing: HybridThrowingValueProvider<Input, Output> {
-		.init(
-			syncValue: self.syncValueProvider,
-			asyncValue: self.asyncValueProvider
-		)
+	/// Hopefully temporary until https://github.com/apple/swift/pull/71143 is available.
+	@MainActor
+	public func mainActorAsync(_ input: Input) async throws -> Output {
+		try await asyncValueProvider(input, MainActor.shared)
 	}
 }
 
-extension HybridThrowingValueProvider {
+// I believe these may still be implementable when https://github.com/apple/swift/pull/71143 is available.
+//extension HybridValueProvider {
+//	/// Returns a new `HybridValueProvider` with a new output type.
+//	public func map<T>(_ transform: @escaping (Output) -> T) -> HybridValueProvider<Input, T> where T: Sendable {
+//		.init(
+//			syncValue: { self.sync($0).map(transform) },
+//			asyncValue: { transform(await self.async($0, isolatedTo: $1)) }
+//		)
+//	}
+//
+//	/// Convert to a `HybridThrowingValueProvider`.
+//	public var throwing: HybridThrowingValueProvider<Input, Output> {
+//		.init(
+//			syncValue: self.syncValueProvider,
+//			asyncValue: self.asyncValueProvider
+//		)
+//	}
+//}
+
+//extension HybridThrowingValueProvider {
 	/// Returns a new `HybridThrowingValueProvider` with a new output type.
-	public func map<T>(_ transform: @escaping (Output) throws -> T) -> HybridThrowingValueProvider<Input, T> where T: Sendable {
-		.init(
-			syncValue: { try self.sync($0).map(transform) },
-			asyncValue: { try transform(try await self.async($0)) }
-		)
-	}
+//	public func map<T>(_ transform: @escaping (Output, isolated (any Actor)?) throws -> T) -> HybridThrowingValueProvider<Input, T> where T: Sendable {
+//		.init(
+//			syncValue: { input in try self.sync($0).map({ transform($0, nil) }) },
+//			asyncValue: { try transform(try await self.async($0, isolatedTo: $1), $1) }
+//		)
+//	}
 
-	/// Transforms a `HybridThrowingValueProvider` into a `HybridValueProvider`.
-	public func catching(_ block: @escaping (Input, Error) -> Output) -> HybridValueProvider<Input, Output> {
-		.init(
-			syncValue: {
-				do {
-					return try self.sync($0)
-				} catch {
-					return block($0, error)
-				}
-			},
-			asyncValue: {
-				do {
-					return try await self.async($0)
-				} catch {
-					return block($0, error)
-				}
-			}
-		)
-	}
-}
+//	/// Transforms a `HybridThrowingValueProvider` into a `HybridValueProvider`.
+//	public func catching(_ block: @escaping (Input, Error) -> Output) -> HybridValueProvider<Input, Output> {
+//		.init(
+//			syncValue: {
+//				do {
+//					return try self.sync($0)
+//				} catch {
+//					return block($0, error)
+//				}
+//			},
+//			asyncValue: {
+//				do {
+//					return try await self.async($0, isolatedTo: $1)
+//				} catch {
+//					return block($0, error)
+//				}
+//			}
+//		)
+//	}
+//}

Sources/RangeState/RangeInvalidationBuffer.swift

@@ -1,5 +1,6 @@
 import Foundation
 
+@MainActor
 public final class RangeInvalidationBuffer {
 	public typealias Handler = (RangeTarget) -> Void
 

Sources/RangeState/RangeValidator.swift

@@ -9,6 +9,7 @@ import Rearrange
 @MainActor
 public final class RangeValidator<Content: VersionedContent> {
 	public typealias Version = Content.Version
+	public typealias ValidationHandler = (NSRange) -> Void
 
 	public enum Validation: Sendable {
 		case stale
@@ -44,6 +45,7 @@ public final class RangeValidator<Content: VersionedContent> {
 	private let continuation: Sequence.Continuation
 
 	public let configuration: Configuration
+	public var validationHandler: ValidationHandler = { _ in }
 
 	public init(configuration: Configuration) {
 		self.configuration = configuration
@@ -62,7 +64,7 @@ public final class RangeValidator<Content: VersionedContent> {
 		continuation.finish()
 	}
 
-	private func beginMonitoring(_ stream: Sequence) async {
+	private nonisolated func beginMonitoring(_ stream: Sequence) async {
 		for await versionedRange in stream {
 			await self.validateRangeAsync(versionedRange)
 		}
@@ -222,7 +224,7 @@ extension RangeValidator {
 		self.pendingRequests -= 1
 		precondition(pendingRequests >= 0)
 
-		let result = await self.configuration.validationProvider.async(contentRange)
+		let result = await self.configuration.validationProvider.mainActorAsync(contentRange)
 
 		switch result {
 		case .stale:
@@ -251,5 +253,7 @@ extension RangeValidator {
 	private func handleValidatedRange(_ range: NSRange) {
 		pendingSet.remove(integersIn: range)
 		validSet.insert(range: range)
+
+		validationHandler(range)
 	}
 }

Sources/TreeSitterClient/TreeSitterClient.swift

@@ -314,7 +314,7 @@ extension TreeSitterClient {
 extension TreeSitterClient {
 	/// Execute a standard highlights.scm query.
 	public func highlights(in set: IndexSet, provider: @escaping TextProvider, mode: RangeFillMode = .required) async throws -> [NamedRange] {
-		try await highlightsProvider.async(.init(indexSet: set, textProvider: provider, mode: mode))
+		try await highlightsProvider.mainActorAsync(.init(indexSet: set, textProvider: provider, mode: mode))
 	}
 
 	/// Execute a standard highlights.scm query.
@@ -324,6 +324,6 @@ extension TreeSitterClient {
 
 	/// Execute a standard highlights.scm query.
 	public func highlights(in range: NSRange, provider: @escaping TextProvider, mode: RangeFillMode = .required) async throws -> [NamedRange] {
-		try await highlightsProvider.async(.init(range: range, textProvider: provider, mode: mode))
+		try await highlightsProvider.mainActorAsync(.init(range: range, textProvider: provider, mode: mode))
 	}
 }