lsp_core/
backend.rs

1use std::{collections::HashMap, sync::Arc};
2
3use bevy_ecs::{
4    bundle::Bundle,
5    component::Component,
6    entity::Entity,
7    schedule::ScheduleLabel,
8    world::{CommandQueue, World},
9};
10use completion::CompletionRequest;
11use futures::lock::Mutex;
12use goto_type::GotoTypeRequest;
13use lsp_types::{request::SemanticTokensRefresh, *};
14use references::ReferencesRequest;
15use request::{GotoTypeDefinitionParams, GotoTypeDefinitionResponse};
16use ropey::Rope;
17use tower_lsp::{jsonrpc::Result, LanguageServer};
18use tracing::info;
19
20use crate::{feature::goto_definition::GotoDefinitionRequest, prelude::*, Startup};
21
22#[derive(Debug)]
23pub struct Backend {
24    entities: Arc<Mutex<HashMap<String, Entity>>>,
25    sender: CommandSender,
26    #[allow(unused)]
27    client: tower_lsp::Client,
28    semantic_tokens: Vec<SemanticTokenType>,
29}
30
31impl Backend {
32    pub fn new(
33        sender: CommandSender,
34        client: tower_lsp::Client,
35        tokens: Vec<SemanticTokenType>,
36    ) -> Self {
37        Self {
38            entities: Default::default(),
39            sender,
40            client,
41            semantic_tokens: tokens,
42        }
43    }
44
45    async fn run<T: Send + Sync + 'static>(
46        &self,
47        f: impl FnOnce(&mut World) -> T + Send + Sync + 'static,
48    ) -> Option<T> {
49        let (tx, rx) = futures::channel::oneshot::channel();
50        let mut commands = CommandQueue::default();
51        commands.push(move |world: &mut World| {
52            let o = f(world);
53            if let Err(_) = tx.send(o) {
54                tracing::error!("Failed to run schedule for {}", stringify!(T));
55            };
56        });
57
58        if let Err(e) = self.sender.0.unbounded_send(commands) {
59            tracing::error!("Failed to send commands {}", e);
60            return None;
61        }
62
63        rx.await.ok()
64    }
65
66    async fn run_schedule<T: Component>(
67        &self,
68        entity: Entity,
69        schedule: impl ScheduleLabel + Clone,
70        param: impl Bundle,
71    ) -> Option<T> {
72        let (tx, rx) = futures::channel::oneshot::channel();
73
74        let mut commands = CommandQueue::default();
75        commands.push(move |world: &mut World| {
76            world.entity_mut(entity).insert(param);
77            world.run_schedule(schedule.clone());
78            if let Err(_) = tx.send(world.entity_mut(entity).take::<T>()) {
79                tracing::error!("Failed to run schedule {:?}", schedule);
80            };
81        });
82
83        if let Err(e) = self.sender.0.unbounded_send(commands) {
84            tracing::error!("Failed to send commands {}", e);
85            return None;
86        }
87
88        rx.await.unwrap_or_default()
89    }
90}
91
92#[tower_lsp::async_trait]
93impl LanguageServer for Backend {
94    #[tracing::instrument(skip(self, init))]
95    async fn initialize(&self, init: InitializeParams) -> Result<InitializeResult> {
96        info!("Initialize");
97        let workspaces = init.workspace_folders.clone().unwrap_or_default();
98        let config: Config =
99            serde_json::from_value(init.initialization_options.clone().unwrap_or_default())
100                .unwrap_or_default();
101
102        let server_config = ServerConfig { config, workspaces };
103        info!("Initialize {:?}", server_config);
104        let document_selectors: Vec<_> = [
105            ("sparql", server_config.config.sparql.unwrap_or(true)),
106            ("turtle", server_config.config.turtle.unwrap_or(true)),
107            ("jsonld", server_config.config.jsonld.unwrap_or(true)),
108        ]
109        .into_iter()
110        .filter(|(_, x)| *x)
111        .map(|(x, _)| DocumentFilter {
112            language: Some(String::from(x)),
113            scheme: None,
114            pattern: None,
115        })
116        .collect();
117
118        self.run(|world| {
119            world.insert_resource(server_config);
120            world.run_schedule(Startup);
121        })
122        .await;
123
124        // let triggers = L::TRIGGERS.iter().copied().map(String::from).collect();
125        Ok(InitializeResult {
126            server_info: None,
127            capabilities: ServerCapabilities {
128                inlay_hint_provider: Some(OneOf::Left(true)),
129                text_document_sync: Some(TextDocumentSyncCapability::Kind(
130                    TextDocumentSyncKind::FULL,
131                )),
132                code_action_provider: None,
133                completion_provider: Some(CompletionOptions {
134                    resolve_provider: Some(false),
135                    trigger_characters: Some(vec![String::from(":")]),
136                    work_done_progress_options: Default::default(),
137                    all_commit_characters: None,
138                    completion_item: None,
139                }),
140                // implementation_provider: Some(ImplementationProviderCapability::Simple(true)),
141                type_definition_provider: Some(TypeDefinitionProviderCapability::Simple(true)),
142                references_provider: Some(OneOf::Left(true)),
143                hover_provider: Some(HoverProviderCapability::Simple(true)),
144                definition_provider: Some(OneOf::Left(true)),
145                document_formatting_provider: Some(OneOf::Left(true)),
146                semantic_tokens_provider: Some(
147                    SemanticTokensServerCapabilities::SemanticTokensRegistrationOptions(
148                        SemanticTokensRegistrationOptions {
149                            text_document_registration_options: {
150                                TextDocumentRegistrationOptions {
151                                    document_selector: Some(document_selectors),
152                                }
153                            },
154                            semantic_tokens_options: SemanticTokensOptions {
155                                work_done_progress_options: WorkDoneProgressOptions::default(),
156                                legend: SemanticTokensLegend {
157                                    token_types: self.semantic_tokens.clone(),
158                                    token_modifiers: vec![],
159                                },
160                                range: Some(false),
161                                full: Some(SemanticTokensFullOptions::Bool(true)),
162                            },
163                            static_registration_options: StaticRegistrationOptions::default(),
164                        },
165                    ),
166                ),
167                rename_provider: Some(OneOf::Right(RenameOptions {
168                    prepare_provider: Some(true),
169                    work_done_progress_options: Default::default(),
170                })),
171                ..ServerCapabilities::default()
172            },
173        })
174    }
175
176    async fn did_change_workspace_folders(&self, params: DidChangeWorkspaceFoldersParams) -> () {
177        self.run(move |world| {
178            let mut config = world.resource_mut::<ServerConfig>();
179            let WorkspaceFoldersChangeEvent { added, removed } = params.event;
180
181            for r in removed {
182                if let Some(idx) = config.workspaces.iter().position(|x| x == &r) {
183                    config.workspaces.remove(idx);
184                }
185            }
186
187            // This is nice and all, but we don't bubble this event up in the world
188            config.workspaces.extend(added);
189        })
190        .await;
191        ()
192    }
193
194    #[tracing::instrument(skip(self, params), fields(uri = %params.text_document.uri.as_str()))]
195    async fn semantic_tokens_full(
196        &self,
197        params: SemanticTokensParams,
198    ) -> Result<Option<SemanticTokensResult>> {
199        info!("semantic tokens full");
200        let uri = params.text_document.uri.as_str();
201        let entity = {
202            let e = {
203                let map = self.entities.lock().await;
204                if let Some(entity) = map.get(uri) {
205                    Some(entity.clone())
206                } else {
207                    info!("Didn't find entity {} retrying", uri);
208                    None
209                }
210            };
211
212            if let Some(e) = e {
213                e
214            } else {
215                let map = self.entities.lock().await;
216                if let Some(entity) = map.get(uri) {
217                    entity.clone()
218                } else {
219                    info!("Didn't find entty {} stopping", uri);
220                    return Ok(None);
221                }
222            }
223        };
224
225        if let Some(res) = self
226            .run_schedule::<HighlightRequest>(entity, SemanticLabel, HighlightRequest(vec![]))
227            .await
228        {
229            Ok(Some(SemanticTokensResult::Tokens(
230                lsp_types::SemanticTokens {
231                    result_id: None,
232                    data: res.0,
233                },
234            )))
235        } else {
236            info!("resulitng in no tokens");
237            Ok(None)
238        }
239    }
240
241    #[tracing::instrument(skip(self))]
242    async fn shutdown(&self) -> Result<()> {
243        info!("Shutting down!");
244
245        Ok(())
246    }
247
248    #[tracing::instrument(skip(self, params), fields(uri = %params.text_document_position.text_document.uri.as_str()))]
249    async fn references(&self, params: ReferenceParams) -> Result<Option<Vec<Location>>> {
250        let entity = {
251            let map = self.entities.lock().await;
252            if let Some(entity) = map.get(params.text_document_position.text_document.uri.as_str())
253            {
254                entity.clone()
255            } else {
256                return Ok(None);
257            }
258        };
259
260        let mut pos = params.text_document_position.position;
261        pos.character = if pos.character > 0 {
262            pos.character - 1
263        } else {
264            pos.character
265        };
266
267        let arr = self
268            .run_schedule::<ReferencesRequest>(
269                entity,
270                ReferencesLabel,
271                (PositionComponent(pos), ReferencesRequest(Vec::new())),
272            )
273            .await
274            .map(|x| x.0);
275
276        Ok(arr)
277    }
278
279    #[tracing::instrument(skip(self, params), fields(uri = %params.text_document.uri.as_str()))]
280    async fn prepare_rename(
281        &self,
282        params: TextDocumentPositionParams,
283    ) -> Result<Option<PrepareRenameResponse>> {
284        let entity = {
285            let map = self.entities.lock().await;
286            if let Some(entity) = map.get(params.text_document.uri.as_str()) {
287                entity.clone()
288            } else {
289                return Ok(None);
290            }
291        };
292
293        let mut pos = params.position;
294        pos.character = if pos.character > 0 {
295            pos.character - 1
296        } else {
297            pos.character
298        };
299
300        let resp = self
301            .run_schedule::<PrepareRenameRequest>(
302                entity,
303                PrepareRenameLabel,
304                PositionComponent(pos),
305            )
306            .await
307            .map(|x| PrepareRenameResponse::RangeWithPlaceholder {
308                range: x.range,
309                placeholder: x.placeholder,
310            });
311
312        Ok(resp)
313    }
314
315    #[tracing::instrument(skip(self, params), fields(uri = %params.text_document_position.text_document.uri.as_str()))]
316    async fn rename(&self, params: RenameParams) -> Result<Option<WorkspaceEdit>> {
317        let entity = {
318            let map = self.entities.lock().await;
319            if let Some(entity) = map.get(params.text_document_position.text_document.uri.as_str())
320            {
321                entity.clone()
322            } else {
323                return Ok(None);
324            }
325        };
326
327        let mut pos = params.text_document_position.position;
328        pos.character = if pos.character > 0 {
329            pos.character - 1
330        } else {
331            pos.character
332        };
333
334        let mut change_map: HashMap<lsp_types::Url, Vec<TextEdit>> = HashMap::new();
335        if let Some(changes) = self
336            .run_schedule::<RenameEdits>(
337                entity,
338                RenameLabel,
339                (
340                    PositionComponent(pos),
341                    RenameEdits(Vec::new(), params.new_name),
342                ),
343            )
344            .await
345        {
346            for (url, change) in changes.0 {
347                let entry = change_map.entry(url);
348                entry.or_default().push(change);
349            }
350        }
351        Ok(Some(WorkspaceEdit::new(change_map)))
352    }
353
354    async fn hover(&self, params: HoverParams) -> Result<Option<lsp_types::Hover>> {
355        let request: HoverRequest = HoverRequest::default();
356
357        let entity = {
358            let map = self.entities.lock().await;
359            if let Some(entity) = map.get(
360                params
361                    .text_document_position_params
362                    .text_document
363                    .uri
364                    .as_str(),
365            ) {
366                entity.clone()
367            } else {
368                return Ok(None);
369            }
370        };
371
372        let mut pos = params.text_document_position_params.position;
373        pos.character = if pos.character > 0 {
374            pos.character - 1
375        } else {
376            pos.character
377        };
378
379        if let Some(hover) = self
380            .run_schedule::<HoverRequest>(entity, HoverLabel, (request, PositionComponent(pos)))
381            .await
382        {
383            if hover.0.len() > 0 {
384                return Ok(Some(lsp_types::Hover {
385                    contents: lsp_types::HoverContents::Array(
386                        hover.0.into_iter().map(MarkedString::String).collect(),
387                    ),
388                    range: hover.1,
389                }));
390            }
391        }
392
393        Ok(None)
394    }
395
396    async fn inlay_hint(&self, params: InlayHintParams) -> Result<Option<Vec<InlayHint>>> {
397        info!("Inlay hints called");
398        let uri = params.text_document.uri.as_str();
399        let entity = {
400            let map = self.entities.lock().await;
401            if let Some(entity) = map.get(uri) {
402                entity.clone()
403            } else {
404                info!("Didn't find entity {}", uri);
405                return Ok(None);
406            }
407        };
408
409        let request = self
410            .run_schedule::<InlayRequest>(entity, InlayLabel, InlayRequest(None))
411            .await;
412
413        Ok(request.and_then(|x| x.0))
414    }
415
416    #[tracing::instrument(skip(self))]
417    async fn formatting(&self, params: DocumentFormattingParams) -> Result<Option<Vec<TextEdit>>> {
418        let uri = params.text_document.uri.as_str();
419        let entity = {
420            let map = self.entities.lock().await;
421            if let Some(entity) = map.get(uri) {
422                entity.clone()
423            } else {
424                info!("Didn't find entity {}", uri);
425                return Ok(None);
426            }
427        };
428
429        let request = self
430            .run_schedule::<FormatRequest>(entity, FormatLabel, FormatRequest(None))
431            .await;
432        Ok(request.and_then(|x| x.0))
433    }
434
435    #[tracing::instrument(skip(self, params), fields(uri = %params.text_document.uri.as_str()))]
436    async fn did_open(&self, params: DidOpenTextDocumentParams) {
437        let item = params.text_document;
438        let url = item.uri.as_str().to_string();
439
440        tracing::info!("Did open");
441
442        let lang_id = Some(item.language_id.clone());
443        let spawn = spawn_or_insert(
444            item.uri.clone(),
445            (
446                Source(item.text.clone()),
447                Label(item.uri.clone()),
448                RopeC(Rope::from_str(&item.text)),
449                Wrapped(item),
450                DocumentLinks(Vec::new()),
451                Open,
452                Types(HashMap::new()),
453            ),
454            lang_id,
455            (),
456        );
457
458        let entity = self
459            .run(|world| {
460                let id = spawn(world);
461                world.run_schedule(ParseLabel);
462                world.flush();
463                info!("Running diagnostics");
464                world.run_schedule(DiagnosticsLabel);
465                info!("Done diagnostics");
466                id
467            })
468            .await;
469
470        if let Some(entity) = entity {
471            self.entities.lock().await.insert(url, entity);
472        }
473
474        let _ = self.client.send_request::<SemanticTokensRefresh>(()).await;
475        info!("Semantic tokens refresh");
476    }
477
478    #[tracing::instrument(skip(self, params), fields(uri = %params.text_document.uri.as_str()))]
479    async fn did_change(&self, params: DidChangeTextDocumentParams) {
480        let entity = {
481            let map = self.entities.lock().await;
482            if let Some(entity) = map.get(params.text_document.uri.as_str()) {
483                entity.clone()
484            } else {
485                info!("Didn't find entity {}", params.text_document.uri.as_str());
486                return;
487            }
488        };
489
490        let change = {
491            if let Some(c) = params.content_changes.into_iter().next() {
492                c
493            } else {
494                return;
495            }
496        };
497
498        self.run(move |world| {
499            let rope_c = RopeC(Rope::from_str(&change.text));
500            world
501                .entity_mut(entity)
502                .insert((Source(change.text), rope_c));
503            world.run_schedule(ParseLabel);
504            world.flush();
505            info!("Running diagnostics");
506            world.run_schedule(DiagnosticsLabel);
507            info!("Running diagnostics done");
508        })
509        .await;
510    }
511
512    #[tracing::instrument(skip(self, params), fields(uri = %params.text_document.uri.as_str()))]
513    async fn did_save(&self, params: DidSaveTextDocumentParams) {
514        let _ = params;
515
516        info!("Did save");
517        self.run(move |world| {
518            world.run_schedule(SaveLabel);
519
520            info!("Ran OnSave Schedule");
521        })
522        .await;
523    }
524
525    #[tracing::instrument(skip(self, params), fields(uri = %params.text_document_position_params.text_document.uri.as_str()))]
526    async fn goto_definition(
527        &self,
528        params: GotoDefinitionParams,
529    ) -> Result<Option<GotoDefinitionResponse>> {
530        let entity = {
531            let map = self.entities.lock().await;
532            if let Some(entity) = map.get(
533                params
534                    .text_document_position_params
535                    .text_document
536                    .uri
537                    .as_str(),
538            ) {
539                entity.clone()
540            } else {
541                return Ok(None);
542            }
543        };
544
545        let mut pos = params.text_document_position_params.position;
546        pos.character = if pos.character > 0 {
547            pos.character - 1
548        } else {
549            pos.character
550        };
551
552        let arr = self
553            .run_schedule::<GotoDefinitionRequest>(
554                entity,
555                GotoDefinitionLabel,
556                (PositionComponent(pos), GotoDefinitionRequest(Vec::new())),
557            )
558            .await
559            .map(|x| GotoDefinitionResponse::Array(x.0));
560
561        Ok(arr)
562    }
563
564    #[tracing::instrument(skip(self, params), fields(uri = %params.text_document_position_params.text_document.uri.as_str()))]
565    async fn goto_type_definition(
566        &self,
567        params: GotoTypeDefinitionParams,
568    ) -> Result<Option<GotoTypeDefinitionResponse>> {
569        let entity = {
570            let map = self.entities.lock().await;
571            if let Some(entity) = map.get(
572                params
573                    .text_document_position_params
574                    .text_document
575                    .uri
576                    .as_str(),
577            ) {
578                entity.clone()
579            } else {
580                return Ok(None);
581            }
582        };
583
584        let mut pos = params.text_document_position_params.position;
585        pos.character = if pos.character > 0 {
586            pos.character - 1
587        } else {
588            pos.character
589        };
590
591        let arr = self
592            .run_schedule::<GotoTypeRequest>(
593                entity,
594                GotoTypeLabel,
595                (PositionComponent(pos), GotoTypeRequest(Vec::new())),
596            )
597            .await
598            .map(|x| GotoTypeDefinitionResponse::Array(x.0));
599
600        Ok(arr)
601    }
602
603    #[tracing::instrument(skip(self, params), fields(uri = %params.text_document_position.text_document.uri.as_str()))]
604    async fn completion(&self, params: CompletionParams) -> Result<Option<CompletionResponse>> {
605        let entity = {
606            let map = self.entities.lock().await;
607            if let Some(entity) = map.get(params.text_document_position.text_document.uri.as_str())
608            {
609                entity.clone()
610            } else {
611                return Ok(None);
612            }
613        };
614
615        // Problem: whne the cursor is at the end of en ident, that ident is not in range of the
616        // cursor
617        let mut pos = params.text_document_position.position;
618        pos.character = if pos.character > 0 {
619            pos.character - 1
620        } else {
621            pos.character
622        };
623
624        let completions: Option<Vec<lsp_types::CompletionItem>> = self
625            .run_schedule::<CompletionRequest>(
626                entity,
627                CompletionLabel,
628                (CompletionRequest(vec![]), PositionComponent(pos)),
629            )
630            .await
631            .map(|x| x.0.into_iter().map(|x| x.into()).collect());
632
633        Ok(completions.map(|c| CompletionResponse::Array(c)))
634    }
635}