@@ -32,9 +32,9 @@ final class SemanticTokenHighlightProvider<
3232 typealias EditCallback = @MainActor ( Result < IndexSet , any Error > ) -> Void
3333 typealias HighlightCallback = @MainActor ( Result < [ HighlightRange ] , any Error > ) -> Void
3434
35- private let tokenMap : SemanticTokenMap
36- private let documentURI : String
37- private weak var languageServer : LanguageServer < DocumentType > ?
35+ private var tokenMap : SemanticTokenMap ?
36+ private var documentURI : String ?
37+ weak var languageServer : LanguageServer < DocumentType > ?
3838 private weak var textView : TextView ?
3939
4040 private var lastEditCallback : EditCallback ?
@@ -45,13 +45,23 @@ final class SemanticTokenHighlightProvider<
4545 textView? . documentRange ?? . zero
4646 }
4747
48- init ( tokenMap: SemanticTokenMap , languageServer: LanguageServer < DocumentType > , documentURI: String ) {
48+ init (
49+ tokenMap: SemanticTokenMap ? = nil ,
50+ languageServer: LanguageServer < DocumentType > ? = nil ,
51+ documentURI: String ? = nil
52+ ) {
4953 self . tokenMap = tokenMap
5054 self . languageServer = languageServer
5155 self . documentURI = documentURI
5256 self . storage = Storage ( )
5357 }
5458
59+ func setUp( server: LanguageServer < DocumentType > , document: DocumentType ) {
60+ languageServer = server
61+ documentURI = document. languageServerURI
62+ tokenMap = server. highlightMap
63+ }
64+
5565 // MARK: - Language Server Content Lifecycle
5666
5767 /// Called when the language server finishes sending a document update.
@@ -95,7 +105,8 @@ final class SemanticTokenHighlightProvider<
95105 textView: TextView ,
96106 lastResultId: String
97107 ) async throws {
98- guard let response = try await languageServer. requestSemanticTokens (
108+ guard let documentURI,
109+ let response = try await languageServer. requestSemanticTokens (
99110 for: documentURI,
100111 previousResultId: lastResultId
101112 ) else {
@@ -112,7 +123,7 @@ final class SemanticTokenHighlightProvider<
112123 /// Requests and applies tokens for an entire document. This does not require a previous response id, and should be
113124 /// used in place of `requestDeltaTokens` when that's the case.
114125 private func requestTokens( languageServer: LanguageServer < DocumentType > , textView: TextView ) async throws {
115- guard let response = try await languageServer. requestSemanticTokens ( for: documentURI) else {
126+ guard let documentURI , let response = try await languageServer. requestSemanticTokens ( for: documentURI) else {
116127 return
117128 }
118129 await applyEntireResponse ( response, callback: lastEditCallback)
@@ -159,7 +170,7 @@ final class SemanticTokenHighlightProvider<
159170 return
160171 }
161172
162- guard let lspRange = textView. lspRangeFrom ( nsRange: range) else {
173+ guard let lspRange = textView. lspRangeFrom ( nsRange: range) , let tokenMap else {
163174 completion ( . failure( HighlightError . lspRangeFailure) )
164175 return
165176 }
0 commit comments