Skip to main content

flams_lsp/
lib.rs

1#![cfg_attr(docsrs, feature(doc_cfg))]
2
3pub mod annotations;
4pub mod capabilities;
5pub mod documents;
6mod implementation;
7pub mod state;
8pub mod verbalizations;
9#[cfg(feature = "ws")]
10pub mod ws;
11
12//use dashmap::Entry;
13pub(crate) use std::collections::hash_map::Entry;
14pub(crate) type HMap<A, B> = rustc_hash::FxHashMap<A, B>;
15// pub(crate) type HMap<A,B> = dashmap::DashMap<A,B,rustc_hash::FxHashBuilder>
16pub(crate) type LMap<A, B> = parking_lot::RwLock<HMap<A, B>>;
17// pub(crate) type LMap<A,B> = dashmap::DashMap<A,B,rustc_hash::FxHashBuilder>
18
19use std::path::Path;
20
21pub use async_lsp;
22use async_lsp::{ClientSocket, LanguageClient, lsp_types as lsp};
23use flams_math_archives::backend::AnyBackend;
24use flams_stex::quickparse::stex::{
25    STeXParseData,
26    structs::{GetModuleError, ModuleReference, STeXModuleStore},
27};
28use flams_system::settings::Settings;
29use flams_utils::{
30    background,
31    parsing::StrParser,
32    sourcerefs::{LSPLineCol, PositionKind, StringRange},
33    unwrap,
34};
35use ftml_uris::{ArchiveId, BaseUri, DocumentUri, Language, SymbolUri};
36use state::{DocData, LSPState, UrlOrFile};
37
38use crate::verbalizations::UnlockedVerbalizationTrie;
39
40static GLOBAL_STATE: std::sync::OnceLock<LSPState> = std::sync::OnceLock::new();
41pub struct STDIOLSPServer {
42    client: ClientSocket,
43    on_port: tokio::sync::watch::Receiver<Option<u16>>,
44    workspaces: Vec<(String, lsp::Url)>,
45}
46
47impl STDIOLSPServer {
48    #[inline]
49    pub fn global_state() -> Option<&'static LSPState> {
50        GLOBAL_STATE.get()
51    }
52    fn load_all(&self) {
53        let client = self.client.clone();
54        let state = unwrap!(Self::global_state());
55        for (name, uri) in &self.workspaces {
56            tracing::info!("workspace: {name}@{uri}");
57        }
58        background(move || state.load_mathhubs(client));
59    }
60
61    fn new_router(
62        client: ClientSocket,
63        on_port: tokio::sync::watch::Receiver<Option<u16>>,
64    ) -> async_lsp::router::Router<ServerWrapper<Self>> {
65        let _ = GLOBAL_STATE.set(LSPState::default());
66        let server = ServerWrapper::new(Self {
67            client,
68            on_port,
69            workspaces: Vec::new(),
70        });
71        server.router()
72    }
73}
74
75#[allow(clippy::future_not_send)]
76/// #### Panics
77pub async fn start_lsp(on_port: tokio::sync::watch::Receiver<Option<u16>>) {
78    let (server, _client) = async_lsp::MainLoop::new_server(|client| {
79        tower::ServiceBuilder::new()
80            .layer(async_lsp::tracing::TracingLayer::default())
81            .layer(async_lsp::server::LifecycleLayer::default())
82            .layer(async_lsp::panic::CatchUnwindLayer::default())
83            .layer(async_lsp::concurrency::ConcurrencyLayer::default())
84            .layer(async_lsp::client_monitor::ClientProcessMonitorLayer::new(
85                client.clone(),
86            ))
87            .service(STDIOLSPServer::new_router(client, on_port))
88    });
89
90    #[cfg(unix)]
91    let (stdin, stdout) = (
92        async_lsp::stdio::PipeStdin::lock_tokio().expect("Failed to lock stdin"),
93        async_lsp::stdio::PipeStdout::lock_tokio().expect("Failed to lock stdout"),
94    );
95    #[cfg(not(unix))]
96    let (stdin, stdout) = (
97        tokio_util::compat::TokioAsyncReadCompatExt::compat(tokio::io::stdin()),
98        tokio_util::compat::TokioAsyncWriteCompatExt::compat_write(tokio::io::stdout()),
99    );
100
101    server
102        .run_buffered(stdin, stdout)
103        .await
104        .expect("Failed to run server");
105}
106
107impl FLAMSLSPServer for STDIOLSPServer {
108    #[inline]
109    fn client_mut(&mut self) -> &mut ClientSocket {
110        &mut self.client
111    }
112    #[inline]
113    fn client(&self) -> &ClientSocket {
114        &self.client
115    }
116    #[inline]
117    fn state(&self) -> &LSPState {
118        Self::global_state().unwrap_or_else(|| unreachable!())
119    }
120    fn initialized(&mut self) {
121        let v = *self.on_port.borrow();
122        if v.is_some() {
123            if let Err(r) = self.client.notify::<ServerURL>(ServerURL::get()) {
124                tracing::error!("failed to send notification: {}", r);
125            }
126        } else {
127            let mut port = self.on_port.clone();
128            let client = self.client.clone();
129            tokio::spawn(async move {
130                let _ = port
131                    .wait_for(|e| {
132                        e.map_or(false, |_| {
133                            if let Err(r) = client.notify::<ServerURL>(ServerURL::get()) {
134                                tracing::error!("failed to send notification: {}", r);
135                            };
136                            true
137                        })
138                    })
139                    .await;
140            });
141        }
142        tracing::info!(
143            "Using {} threads",
144            tokio::runtime::Handle::current().metrics().num_workers()
145        );
146        //#[cfg(not(debug_assertions))]
147        {
148            self.load_all();
149        }
150    }
151
152    fn initialize<I: Iterator<Item = (String, lsp::Url)> + Send + 'static>(
153        &mut self,
154        workspaces: I,
155    ) {
156        self.workspaces = workspaces.collect();
157    }
158}
159
160#[derive(serde::Serialize, serde::Deserialize)]
161struct ReloadParams {}
162struct Reload;
163impl lsp::notification::Notification for Reload {
164    type Params = ReloadParams;
165    const METHOD: &str = "flams/reload";
166}
167
168#[derive(serde::Serialize, serde::Deserialize)]
169struct InstallParams {
170    pub archives: Vec<ArchiveId>,
171    pub remote_url: String,
172}
173struct InstallArchives;
174impl lsp::notification::Notification for InstallArchives {
175    type Params = InstallParams;
176    const METHOD: &str = "flams/install";
177}
178
179#[derive(serde::Serialize, serde::Deserialize)]
180struct NewArchiveParams {
181    pub archive: ArchiveId,
182    pub urlbase: BaseUri,
183}
184struct NewArchive;
185impl lsp::notification::Notification for NewArchive {
186    type Params = NewArchiveParams;
187    const METHOD: &str = "flams/newArchive";
188}
189
190pub(crate) struct StandaloneExport;
191#[derive(serde::Serialize, serde::Deserialize)]
192pub struct StandaloneExportParams {
193    pub uri: lsp::Url,
194    pub target: std::path::PathBuf,
195}
196impl lsp::notification::Notification for StandaloneExport {
197    type Params = StandaloneExportParams;
198    const METHOD: &str = "flams/standaloneExport";
199}
200
201pub(crate) struct HtmlExport;
202impl lsp::notification::Notification for HtmlExport {
203    type Params = StandaloneExportParams;
204    const METHOD: &str = "flams/htmlExport";
205}
206
207struct UpdateMathHub;
208impl lsp::notification::Notification for UpdateMathHub {
209    type Params = Vec<()>;
210    const METHOD: &str = "flams/updateMathHub";
211}
212
213struct OpenFile;
214impl lsp::notification::Notification for OpenFile {
215    type Params = UriParams;
216    const METHOD: &str = "flams/openFile";
217}
218
219struct HTMLResult;
220impl lsp::notification::Notification for HTMLResult {
221    type Params = StringParams;
222    const METHOD: &str = "flams/htmlResult";
223}
224
225pub(crate) struct SnifyNotification;
226impl lsp::notification::Notification for SnifyNotification {
227    type Params = UriParams;
228    const METHOD: &str = "flams/snify";
229}
230
231#[derive(serde::Serialize, serde::Deserialize)]
232pub struct UriParams {
233    pub uri: lsp::Url,
234}
235pub(crate) struct HTMLRequest;
236impl lsp::request::Request for HTMLRequest {
237    type Params = UriParams;
238    type Result = Option<String>;
239    const METHOD: &'static str = "flams/htmlRequest";
240}
241
242struct BuildOne;
243impl lsp::request::Request for BuildOne {
244    type Params = UriParams;
245    type Result = ();
246    const METHOD: &str = "flams/buildOne";
247}
248struct BuildAll;
249impl lsp::request::Request for BuildAll {
250    type Params = UriParams;
251    type Result = ();
252    const METHOD: &str = "flams/buildAll";
253}
254
255pub(crate) struct QuizRequest;
256impl lsp::request::Request for QuizRequest {
257    type Params = UriParams;
258    type Result = String;
259    const METHOD: &'static str = "flams/quizRequest";
260}
261
262pub struct ServerURL;
263impl ServerURL {
264    fn get() -> StringParams {
265        let settings = Settings::get();
266        StringParams {
267            url: format!("http://{}:{}", settings.ip, settings.port()),
268        }
269    }
270}
271
272#[derive(serde::Serialize, serde::Deserialize)]
273pub struct StringParams {
274    url: String,
275}
276
277impl lsp::notification::Notification for ServerURL {
278    type Params = StringParams;
279    const METHOD: &str = "flams/serverURL";
280}
281
282pub trait ClientExt {
283    fn html_result(&self, uri: &DocumentUri);
284    fn update_mathhub(&self);
285    fn open_file(&self, path: &Path);
286}
287impl ClientExt for ClientSocket {
288    #[inline]
289    fn html_result(&self, uri: &DocumentUri) {
290        let _ = self.notify::<HTMLResult>(StringParams {
291            url: uri.to_string(),
292        });
293    }
294    #[inline]
295    fn update_mathhub(&self) {
296        if let Err(e) = self.notify::<UpdateMathHub>(Vec::new()) {
297            tracing::error!("failed to send notification: {}", e);
298        }
299    }
300
301    fn open_file(&self, path: &Path) {
302        let Ok(url) = lsp::Url::from_file_path(path) else {
303            return;
304        };
305        if let Err(e) = self.notify::<OpenFile>(UriParams { uri: url }) {
306            tracing::error!("failed to send notification: {}", e);
307        }
308    }
309}
310
311pub trait FLAMSLSPServer: 'static {
312    fn client_mut(&mut self) -> &mut ClientSocket;
313    fn client(&self) -> &ClientSocket;
314    fn state(&self) -> &LSPState;
315    #[inline]
316    fn initialized(&mut self) {}
317    #[inline]
318    fn initialize<I: Iterator<Item = (String, lsp::Url)> + Send + 'static>(
319        &mut self,
320        _workspaces: I,
321    ) {
322    }
323}
324
325pub struct ServerWrapper<T: FLAMSLSPServer> {
326    pub inner: T,
327}
328impl<T: FLAMSLSPServer> ServerWrapper<T> {
329    #[inline]
330    pub const fn new(inner: T) -> Self {
331        Self { inner }
332    }
333
334    pub fn router(self) -> async_lsp::router::Router<Self> {
335        let mut r = async_lsp::router::Router::from_language_server(self);
336        r.request::<HTMLRequest, _>(Self::html_request);
337        r.request::<QuizRequest, _>(Self::quiz_request);
338        r.request::<BuildOne, _>(Self::build_one);
339        r.request::<BuildAll, _>(Self::build_all);
340
341        r.notification::<Reload>(Self::reload);
342        r.notification::<StandaloneExport>(Self::export_standalone);
343        r.notification::<SnifyNotification>(Self::snify);
344        r.notification::<HtmlExport>(Self::export_html);
345        r.notification::<InstallArchives>(Self::install);
346        r.notification::<NewArchive>(Self::new_archive);
347        //r.request(handler)
348        r
349    }
350
351    pub fn get_progress(&self, tk: lsp::ProgressToken) -> ProgressCallbackClient {
352        match &tk {
353            lsp::ProgressToken::Number(n) if *n > 0 => {
354                TOKEN.fetch_max(*n as u32 + 1, std::sync::atomic::Ordering::Relaxed);
355            }
356            _ => (),
357        }
358        ProgressCallbackClient {
359            client: self.inner.client().clone(),
360            token: tk,
361        }
362    }
363}
364
365pub struct LSPStore<'a, 'l, const FULL: bool> {
366    pub(crate) map: &'a mut HMap<UrlOrFile, DocData>,
367    verbalizations: Option<&'a mut UnlockedVerbalizationTrie<'l>>,
368    ignore_verbalizations: &'a [&'a str],
369    cycles: Vec<DocumentUri>,
370    snify: bool,
371}
372impl<'a, 'l, const FULL: bool> LSPStore<'a, 'l, FULL> {
373    #[inline]
374    pub const fn new(
375        map: &'a mut HMap<UrlOrFile, DocData>,
376        verbalizations: Option<&'a mut UnlockedVerbalizationTrie<'l>>,
377        ignore_verbalizations: &'a [&'a str],
378        snify: bool,
379    ) -> Self {
380        Self {
381            map,
382            cycles: Vec::new(),
383            ignore_verbalizations,
384            verbalizations,
385            snify,
386        }
387    }
388
389    pub fn load(&mut self, p: &Path, uri: &DocumentUri) -> Option<STeXParseData> {
390        let text = std::fs::read_to_string(p).ok()?;
391        let r = flams_stex::quickparse::stex::quickparse(uri, &text, p, &AnyBackend::Global, self)
392            .lock();
393        Some(r)
394    }
395
396    fn load_as_false(&mut self, p: &Path, uri: &DocumentUri) -> Option<STeXParseData> {
397        if !FULL {
398            self.load(p, uri)
399        } else {
400            let mut nstore = if let Some(verb) = &mut self.verbalizations {
401                LSPStore::<'_, '_, false>::new(self.map, Some(verb), &[], false)
402            } else {
403                LSPStore::<'_, '_, false>::new(self.map, None, &[], false)
404            };
405            nstore.cycles = std::mem::take(&mut self.cycles);
406            let r = nstore.load(p, uri);
407            self.cycles = nstore.cycles;
408            r
409        }
410    }
411}
412
413impl<const FULL: bool> STeXModuleStore for &mut LSPStore<'_, '_, FULL> {
414    const FULL: bool = FULL;
415    fn get_module(
416        &mut self,
417        module: &ModuleReference,
418        _in_path: Option<&std::sync::Arc<Path>>,
419    ) -> Result<STeXParseData, GetModuleError> {
420        let Some(p) = module.full_path.as_ref() else {
421            return Err(GetModuleError::NotFound(module.uri.clone()));
422        };
423        let uri = &module.in_doc;
424        if let Some(i) = self.cycles.iter().position(|u| u == uri) {
425            let mut ret = self.cycles[i..].to_vec();
426            ret.push(uri.clone());
427            return Err(GetModuleError::Cycle(ret));
428        }
429
430        macro_rules! do_return {
431            ($e:expr) => {{
432                /*if TRACK_DEPS {
433                  if let Some(in_path) = in_path {
434                    if module.uri != *flams_ontology::metatheory::URI {
435                      $e.lock().dependents.insert_clone(in_path);
436                    }
437                  }
438                }*/
439                return Ok($e);
440            }};
441        }
442
443        let lsp_uri = UrlOrFile::File(p.clone());
444        //let lsp_uri = lsp::Url::from_file_path(p).map_err(|_| GetModuleError::NotFound(module.uri.clone()))?;
445        match self.map.get(&lsp_uri).as_deref() {
446            Some(DocData::Data(d, _)) => do_return!(d.clone()),
447            Some(DocData::Doc(d)) if d.is_up_to_date() => do_return!(d.annotations.clone()),
448            _ => (),
449        }
450
451        self.cycles.push(uri.clone());
452        let r = self
453            .load_as_false(p, uri)
454            .inspect(|ret| match self.map.entry(lsp_uri) {
455                Entry::Vacant(e) => {
456                    e.insert(DocData::Data(ret.clone(), FULL));
457                }
458                Entry::Occupied(mut e) => {
459                    e.get_mut().merge(DocData::Data(ret.clone(), FULL));
460                }
461            });
462        /*if TRACK_DEPS {
463          if let Some(r) = &r {
464            if let Some(in_path) = in_path {
465              if module.uri != *flams_ontology::metatheory::URI {
466                r.lock().dependencies.insert_clone(in_path);
467              }
468            }
469          }
470        }*/
471        self.cycles.pop();
472        r.ok_or_else(|| GetModuleError::NotFound(module.uri.clone()))
473    }
474    fn add_verbalization(
475        &mut self,
476        s: &str,
477        symbol: &ftml_uris::prelude::SymbolUri,
478        language: Language,
479    ) {
480        if let Some(verb) = &mut self.verbalizations {
481            verb.insert(language, s, symbol);
482        }
483    }
484    fn add_text<Pos: flams_utils::sourcerefs::StringPosition>(
485        &self,
486        r: StringRange<Pos>,
487        text: &str,
488        language: Language,
489        needs_usemodule: &dyn Fn(&SymbolUri) -> bool,
490    ) -> Option<flams_stex::quickparse::stex::structs::STeXToken<Pos>> {
491        //Vec<(StringRange<Pos>, SmallVec<SymbolUri, 1>)> {
492        if self.snify
493            && Pos::KIND == PositionKind::LineColUtf16
494            && Self::FULL
495            && let Some(verb) = &self.verbalizations
496        {
497            let v = verb.find_all_in(
498                language,
499                text,
500                r.start,
501                self.ignore_verbalizations,
502                needs_usemodule,
503            );
504            if v.is_empty() {
505                None
506            } else {
507                Some(flams_stex::quickparse::stex::structs::STeXToken::Vec(
508                    v.into_iter()
509                        .map(|(r, s)| {
510                            flams_stex::quickparse::stex::structs::STeXToken::SnifySuggestion {
511                                range: r,
512                                symbols: s,
513                            }
514                        })
515                        .collect(),
516                ))
517            }
518        } else {
519            None
520        }
521    }
522}
523
524pub trait IsLSPRange {
525    fn into_range(self) -> lsp::Range;
526    fn from_range(range: lsp::Range) -> Self;
527}
528
529impl IsLSPRange for StringRange<LSPLineCol> {
530    fn into_range(self) -> lsp::Range {
531        lsp::Range {
532            start: lsp::Position {
533                line: self.start.line,
534                character: self.start.col,
535            },
536            end: lsp::Position {
537                line: self.end.line,
538                character: self.end.col,
539            },
540        }
541    }
542    fn from_range(range: lsp::Range) -> Self {
543        Self {
544            start: LSPLineCol::new(range.start.line, range.start.character),
545            end: LSPLineCol::new(range.end.line, range.end.character),
546        }
547    }
548}
549
550pub struct ProgressCallbackServer {
551    client: ClientSocket,
552    token: u32,
553    handle: tokio::task::JoinHandle<()>,
554    progress: Option<parking_lot::Mutex<(u32, u32)>>,
555}
556
557lazy_static::lazy_static! {
558  static ref TOKEN:triomphe::Arc<std::sync::atomic::AtomicU32> = triomphe::Arc::new(std::sync::atomic::AtomicU32::new(42));
559}
560
561impl ProgressCallbackServer {
562    #[inline]
563    pub const fn client_mut(&mut self) -> &mut ClientSocket {
564        &mut self.client
565    }
566
567    #[inline]
568    pub fn client(&self) -> ClientSocket {
569        self.client.clone()
570    }
571
572    pub fn with<R>(
573        client: ClientSocket,
574        title: String,
575        total: Option<u32>,
576        f: impl FnOnce(Self) -> R,
577    ) -> R {
578        let slf = Self::new(client, title, total);
579        f(slf)
580    }
581
582    #[must_use]
583    #[allow(clippy::let_underscore_future)]
584    pub fn new(mut client: ClientSocket, title: String, total: Option<u32>) -> Self {
585        let token = TOKEN.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
586        let f = client.work_done_progress_create(lsp::WorkDoneProgressCreateParams {
587            token: lsp::NumberOrString::Number(token as _),
588        });
589        let mut c = client.clone();
590        let handle = tokio::spawn(async move {
591            if let Err(e) = f.await {
592                tracing::error!("Error: {}", e);
593            } else {
594                let _ = c.progress(async_lsp::lsp_types::ProgressParams {
595                    token: async_lsp::lsp_types::ProgressToken::Number(token as _),
596                    value: async_lsp::lsp_types::ProgressParamsValue::WorkDone(
597                        async_lsp::lsp_types::WorkDoneProgress::Begin(
598                            async_lsp::lsp_types::WorkDoneProgressBegin {
599                                message: None,
600                                title,
601                                percentage: total.map(|_| 0),
602                                cancellable: None,
603                            },
604                        ),
605                    ),
606                });
607            }
608        });
609        //tracing::info!("New progress with id {token}");
610        Self {
611            client,
612            token,
613            handle,
614            progress: total.map(|i| parking_lot::Mutex::new((0, i))),
615        }
616    }
617
618    pub fn update(&self, message: String, add_step: Option<u32>) {
619        let (message, percentage) = if let Some(i) = add_step {
620            if let Some(lock) = self.progress.as_ref() {
621                let mut lock = lock.lock();
622                lock.0 += i;
623                (
624                    format!("{}/{}:{message}", lock.0, lock.1),
625                    Some(100 * lock.0 / lock.1),
626                )
627            } else {
628                (message, None)
629            }
630        } else if let Some(lock) = self.progress.as_ref() {
631            let lock = lock.lock();
632            (
633                format!("{}/{}:{message}", lock.0, lock.1),
634                Some(100 * lock.0 / lock.1),
635            )
636        } else {
637            (message, None)
638        };
639        let token = async_lsp::lsp_types::ProgressToken::Number(self.token as _);
640        //tracing::info!("updating progress {}",self.token);
641        while !self.handle.is_finished() {
642            std::thread::sleep(std::time::Duration::from_millis(10));
643        }
644        let _ = self
645            .client
646            .clone()
647            .progress(async_lsp::lsp_types::ProgressParams {
648                token,
649                value: async_lsp::lsp_types::ProgressParamsValue::WorkDone(
650                    async_lsp::lsp_types::WorkDoneProgress::Report(
651                        async_lsp::lsp_types::WorkDoneProgressReport {
652                            message: Some(message),
653                            percentage,
654                            cancellable: None,
655                        },
656                    ),
657                ),
658            });
659    }
660}
661
662impl Drop for ProgressCallbackServer {
663    fn drop(&mut self) {
664        let token = async_lsp::lsp_types::ProgressToken::Number(self.token as _);
665        let _ = self.client.progress(async_lsp::lsp_types::ProgressParams {
666            token,
667            value: async_lsp::lsp_types::ProgressParamsValue::WorkDone(
668                async_lsp::lsp_types::WorkDoneProgress::End(
669                    async_lsp::lsp_types::WorkDoneProgressEnd {
670                        message: Some("done".to_string()),
671                    },
672                ),
673            ),
674        });
675    }
676}
677
678pub struct ProgressCallbackClient {
679    client: ClientSocket,
680    token: async_lsp::lsp_types::ProgressToken,
681}
682
683impl ProgressCallbackClient {
684    pub fn finish(mut self) {
685        let _ = self.client.progress(async_lsp::lsp_types::ProgressParams {
686            token: self.token,
687            value: async_lsp::lsp_types::ProgressParamsValue::WorkDone(
688                async_lsp::lsp_types::WorkDoneProgress::End(
689                    async_lsp::lsp_types::WorkDoneProgressEnd {
690                        message: Some("done".to_string()),
691                    },
692                ),
693            ),
694        });
695    }
696
697    #[allow(clippy::let_underscore_future)]
698    pub fn finish_delay(self) {
699        let _ = tokio::spawn(async move {
700            tokio::time::sleep(tokio::time::Duration::from_secs_f32(0.5)).await;
701            self.finish();
702        });
703    }
704}