kcl_lib/lsp/kcl/
mod.rs

1//! Functions for the `kcl` lsp server.
2#![allow(dead_code)]
3
4use std::{
5    collections::HashMap,
6    io::Write,
7    str::FromStr,
8    sync::{Arc, Mutex},
9};
10
11use anyhow::Result;
12#[cfg(feature = "cli")]
13use clap::Parser;
14use dashmap::DashMap;
15use sha2::Digest;
16use tokio::sync::RwLock;
17use tower_lsp::{
18    jsonrpc::Result as RpcResult,
19    lsp_types::{
20        CodeAction, CodeActionKind, CodeActionOrCommand, CodeActionParams, CodeActionResponse, CompletionItem,
21        CompletionItemKind, CompletionOptions, CompletionParams, CompletionResponse, CreateFilesParams,
22        DeleteFilesParams, Diagnostic, DiagnosticOptions, DiagnosticServerCapabilities, DiagnosticSeverity,
23        DidChangeConfigurationParams, DidChangeTextDocumentParams, DidChangeWatchedFilesParams,
24        DidChangeWorkspaceFoldersParams, DidCloseTextDocumentParams, DidOpenTextDocumentParams,
25        DidSaveTextDocumentParams, DocumentDiagnosticParams, DocumentDiagnosticReport, DocumentDiagnosticReportResult,
26        DocumentFilter, DocumentFormattingParams, DocumentSymbol, DocumentSymbolParams, DocumentSymbolResponse,
27        Documentation, FoldingRange, FoldingRangeParams, FoldingRangeProviderCapability, FullDocumentDiagnosticReport,
28        Hover as LspHover, HoverContents, HoverParams, HoverProviderCapability, InitializeParams, InitializeResult,
29        InitializedParams, InlayHint, InlayHintParams, InsertTextFormat, MarkupContent, MarkupKind, MessageType, OneOf,
30        Position, RelatedFullDocumentDiagnosticReport, RenameFilesParams, RenameParams, SemanticToken,
31        SemanticTokenModifier, SemanticTokenType, SemanticTokens, SemanticTokensFullOptions, SemanticTokensLegend,
32        SemanticTokensOptions, SemanticTokensParams, SemanticTokensRegistrationOptions, SemanticTokensResult,
33        SemanticTokensServerCapabilities, ServerCapabilities, SignatureHelp, SignatureHelpOptions, SignatureHelpParams,
34        StaticRegistrationOptions, TextDocumentItem, TextDocumentRegistrationOptions, TextDocumentSyncCapability,
35        TextDocumentSyncKind, TextDocumentSyncOptions, TextEdit, WorkDoneProgressOptions, WorkspaceEdit,
36        WorkspaceFolder, WorkspaceFoldersServerCapabilities, WorkspaceServerCapabilities,
37    },
38    Client, LanguageServer,
39};
40
41use crate::{
42    docs::kcl_doc::DocData,
43    errors::Suggestion,
44    exec::KclValue,
45    execution::{cache, kcl_value::FunctionSource},
46    lsp::{
47        backend::Backend as _,
48        kcl::hover::{Hover, HoverOpts},
49        util::IntoDiagnostic,
50    },
51    parsing::{
52        ast::types::{Expr, VariableKind},
53        token::TokenStream,
54        PIPE_OPERATOR,
55    },
56    ModuleId, Program, SourceRange,
57};
58
59pub mod custom_notifications;
60mod hover;
61
62const SEMANTIC_TOKEN_TYPES: [SemanticTokenType; 10] = [
63    SemanticTokenType::NUMBER,
64    SemanticTokenType::VARIABLE,
65    SemanticTokenType::KEYWORD,
66    SemanticTokenType::TYPE,
67    SemanticTokenType::STRING,
68    SemanticTokenType::OPERATOR,
69    SemanticTokenType::COMMENT,
70    SemanticTokenType::FUNCTION,
71    SemanticTokenType::PARAMETER,
72    SemanticTokenType::PROPERTY,
73];
74
75const SEMANTIC_TOKEN_MODIFIERS: [SemanticTokenModifier; 5] = [
76    SemanticTokenModifier::DECLARATION,
77    SemanticTokenModifier::DEFINITION,
78    SemanticTokenModifier::DEFAULT_LIBRARY,
79    SemanticTokenModifier::READONLY,
80    SemanticTokenModifier::STATIC,
81];
82
83/// A subcommand for running the server.
84#[derive(Clone, Debug)]
85#[cfg_attr(feature = "cli", derive(Parser))]
86pub struct Server {
87    /// Port that the server should listen
88    #[cfg_attr(feature = "cli", clap(long, default_value = "8080"))]
89    pub socket: i32,
90
91    /// Listen over stdin and stdout instead of a tcp socket.
92    #[cfg_attr(feature = "cli", clap(short, long, default_value = "false"))]
93    pub stdio: bool,
94}
95
96/// The lsp server backend.
97#[derive(Clone)]
98pub struct Backend {
99    /// The client for the backend.
100    pub client: Client,
101    /// The file system client to use.
102    pub fs: Arc<crate::fs::FileManager>,
103    /// The workspace folders.
104    pub workspace_folders: DashMap<String, WorkspaceFolder>,
105    /// The stdlib completions for the language.
106    pub stdlib_completions: HashMap<String, CompletionItem>,
107    /// The stdlib signatures for the language.
108    pub stdlib_signatures: HashMap<String, SignatureHelp>,
109    /// For all KwArg functions in std, a map from their arg names to arg help snippets (markdown format).
110    pub stdlib_args: HashMap<String, HashMap<String, String>>,
111    /// Token maps.
112    pub(super) token_map: DashMap<String, TokenStream>,
113    /// AST maps.
114    pub ast_map: DashMap<String, crate::Program>,
115    /// Current code.
116    pub code_map: DashMap<String, Vec<u8>>,
117    /// Diagnostics.
118    pub diagnostics_map: DashMap<String, Vec<Diagnostic>>,
119    /// Symbols map.
120    pub symbols_map: DashMap<String, Vec<DocumentSymbol>>,
121    /// Semantic tokens map.
122    pub semantic_tokens_map: DashMap<String, Vec<SemanticToken>>,
123    /// The Zoo API client.
124    pub zoo_client: kittycad::Client,
125    /// If we can send telemetry for this user.
126    pub can_send_telemetry: bool,
127    /// Optional executor context to use if we want to execute the code.
128    pub executor_ctx: Arc<RwLock<Option<crate::execution::ExecutorContext>>>,
129    /// If we are currently allowed to execute the ast.
130    pub can_execute: Arc<RwLock<bool>>,
131
132    pub is_initialized: Arc<RwLock<bool>>,
133}
134
135impl Backend {
136    #[cfg(target_arch = "wasm32")]
137    pub fn new_wasm(
138        client: Client,
139        executor_ctx: Option<crate::execution::ExecutorContext>,
140        fs: crate::fs::wasm::FileSystemManager,
141        zoo_client: kittycad::Client,
142        can_send_telemetry: bool,
143    ) -> Result<Self, String> {
144        Self::with_file_manager(
145            client,
146            executor_ctx,
147            crate::fs::FileManager::new(fs),
148            zoo_client,
149            can_send_telemetry,
150        )
151    }
152
153    #[cfg(not(target_arch = "wasm32"))]
154    pub fn new(
155        client: Client,
156        executor_ctx: Option<crate::execution::ExecutorContext>,
157        zoo_client: kittycad::Client,
158        can_send_telemetry: bool,
159    ) -> Result<Self, String> {
160        Self::with_file_manager(
161            client,
162            executor_ctx,
163            crate::fs::FileManager::new(),
164            zoo_client,
165            can_send_telemetry,
166        )
167    }
168
169    fn with_file_manager(
170        client: Client,
171        executor_ctx: Option<crate::execution::ExecutorContext>,
172        fs: crate::fs::FileManager,
173        zoo_client: kittycad::Client,
174        can_send_telemetry: bool,
175    ) -> Result<Self, String> {
176        let stdlib = crate::std::StdLib::new();
177        let kcl_std = crate::docs::kcl_doc::walk_prelude();
178        let stdlib_completions = get_completions_from_stdlib(&stdlib, &kcl_std).map_err(|e| e.to_string())?;
179        let stdlib_signatures = get_signatures_from_stdlib(&stdlib, &kcl_std);
180        let stdlib_args = get_arg_maps_from_stdlib(&stdlib, &kcl_std);
181
182        Ok(Self {
183            client,
184            fs: Arc::new(fs),
185            stdlib_completions,
186            stdlib_signatures,
187            stdlib_args,
188            zoo_client,
189            can_send_telemetry,
190            can_execute: Arc::new(RwLock::new(executor_ctx.is_some())),
191            executor_ctx: Arc::new(RwLock::new(executor_ctx)),
192            workspace_folders: Default::default(),
193            token_map: Default::default(),
194            ast_map: Default::default(),
195            code_map: Default::default(),
196            diagnostics_map: Default::default(),
197            symbols_map: Default::default(),
198            semantic_tokens_map: Default::default(),
199            is_initialized: Default::default(),
200        })
201    }
202
203    fn remove_from_ast_maps(&self, filename: &str) {
204        self.ast_map.remove(filename);
205        self.symbols_map.remove(filename);
206    }
207}
208
209// Implement the shared backend trait for the language server.
210#[async_trait::async_trait]
211impl crate::lsp::backend::Backend for Backend {
212    fn client(&self) -> &Client {
213        &self.client
214    }
215
216    fn fs(&self) -> &Arc<crate::fs::FileManager> {
217        &self.fs
218    }
219
220    async fn is_initialized(&self) -> bool {
221        *self.is_initialized.read().await
222    }
223
224    async fn set_is_initialized(&self, is_initialized: bool) {
225        *self.is_initialized.write().await = is_initialized;
226    }
227
228    async fn workspace_folders(&self) -> Vec<WorkspaceFolder> {
229        // TODO: fix clone
230        self.workspace_folders.iter().map(|i| i.clone()).collect()
231    }
232
233    async fn add_workspace_folders(&self, folders: Vec<WorkspaceFolder>) {
234        for folder in folders {
235            self.workspace_folders.insert(folder.name.to_string(), folder);
236        }
237    }
238
239    async fn remove_workspace_folders(&self, folders: Vec<WorkspaceFolder>) {
240        for folder in folders {
241            self.workspace_folders.remove(&folder.name);
242        }
243    }
244
245    fn code_map(&self) -> &DashMap<String, Vec<u8>> {
246        &self.code_map
247    }
248
249    async fn insert_code_map(&self, uri: String, text: Vec<u8>) {
250        self.code_map.insert(uri, text);
251    }
252
253    async fn remove_from_code_map(&self, uri: String) -> Option<Vec<u8>> {
254        self.code_map.remove(&uri).map(|x| x.1)
255    }
256
257    async fn clear_code_state(&self) {
258        self.code_map.clear();
259        self.token_map.clear();
260        self.ast_map.clear();
261        self.diagnostics_map.clear();
262        self.symbols_map.clear();
263        self.semantic_tokens_map.clear();
264    }
265
266    fn current_diagnostics_map(&self) -> &DashMap<String, Vec<Diagnostic>> {
267        &self.diagnostics_map
268    }
269
270    async fn inner_on_change(&self, params: TextDocumentItem, force: bool) {
271        if force {
272            crate::bust_cache().await;
273        }
274
275        let filename = params.uri.to_string();
276        // We already updated the code map in the shared backend.
277
278        // Lets update the tokens.
279        let module_id = ModuleId::default();
280        let tokens = match crate::parsing::token::lex(&params.text, module_id) {
281            Ok(tokens) => tokens,
282            Err(err) => {
283                self.add_to_diagnostics(&params, &[err], true).await;
284                self.token_map.remove(&filename);
285                self.remove_from_ast_maps(&filename);
286                self.semantic_tokens_map.remove(&filename);
287                return;
288            }
289        };
290
291        // Get the previous tokens.
292        let tokens_changed = if let Some(previous_tokens) = self.token_map.get(&filename) {
293            *previous_tokens != tokens
294        } else {
295            true
296        };
297
298        let had_diagnostics = self.has_diagnostics(params.uri.as_ref()).await;
299
300        // Check if the tokens are the same.
301        if !tokens_changed && !force && !had_diagnostics {
302            // We return early here because the tokens are the same.
303            return;
304        }
305
306        if tokens_changed {
307            // Update our token map.
308            self.token_map.insert(params.uri.to_string(), tokens.clone());
309            // Update our semantic tokens.
310            self.update_semantic_tokens(&tokens, &params).await;
311        }
312
313        // Lets update the ast.
314
315        let (ast, errs) = match crate::parsing::parse_tokens(tokens.clone()).0 {
316            Ok(result) => result,
317            Err(err) => {
318                self.add_to_diagnostics(&params, &[err], true).await;
319                self.remove_from_ast_maps(&filename);
320                return;
321            }
322        };
323
324        self.add_to_diagnostics(&params, &errs, true).await;
325
326        if errs.iter().any(|e| e.severity == crate::errors::Severity::Fatal) {
327            self.remove_from_ast_maps(&filename);
328            return;
329        }
330
331        let Some(mut ast) = ast else {
332            self.remove_from_ast_maps(&filename);
333            return;
334        };
335
336        // Here we will want to store the digest and compare, but for now
337        // we're doing this in a non-load-bearing capacity so we can remove
338        // this if it backfires and only hork the LSP.
339        ast.compute_digest();
340
341        // Save it as a program.
342        let ast = crate::Program {
343            ast,
344            original_file_contents: params.text.clone(),
345        };
346
347        // Check if the ast changed.
348        let ast_changed = match self.ast_map.get(&filename) {
349            Some(old_ast) => {
350                // Check if the ast changed.
351                *old_ast.ast != *ast.ast
352            }
353            None => true,
354        };
355
356        if !ast_changed && !force && !had_diagnostics {
357            // Return early if the ast did not change and we don't need to force.
358            return;
359        }
360
361        if ast_changed {
362            self.ast_map.insert(params.uri.to_string(), ast.clone());
363            // Update the symbols map.
364            self.symbols_map.insert(
365                params.uri.to_string(),
366                ast.ast.get_lsp_symbols(&params.text).unwrap_or_default(),
367            );
368
369            // Update our semantic tokens.
370            self.update_semantic_tokens(&tokens, &params).await;
371
372            let discovered_findings = ast.lint_all().into_iter().flatten().collect::<Vec<_>>();
373            self.add_to_diagnostics(&params, &discovered_findings, false).await;
374        }
375
376        // Send the notification to the client that the ast was updated.
377        if self.can_execute().await || self.executor_ctx().await.is_none() {
378            // Only send the notification if we can execute.
379            // Otherwise it confuses the client.
380            self.client
381                .send_notification::<custom_notifications::AstUpdated>(ast.ast.clone())
382                .await;
383        }
384
385        // Execute the code if we have an executor context.
386        // This function automatically executes if we should & updates the diagnostics if we got
387        // errors.
388        if self.execute(&params, &ast).await.is_err() {
389            return;
390        }
391
392        // If we made it here we can clear the diagnostics.
393        self.clear_diagnostics_map(&params.uri, Some(DiagnosticSeverity::ERROR))
394            .await;
395    }
396}
397
398impl Backend {
399    pub async fn can_execute(&self) -> bool {
400        *self.can_execute.read().await
401    }
402
403    pub async fn executor_ctx(&self) -> tokio::sync::RwLockReadGuard<'_, Option<crate::execution::ExecutorContext>> {
404        self.executor_ctx.read().await
405    }
406
407    async fn update_semantic_tokens(&self, tokens: &TokenStream, params: &TextDocumentItem) {
408        // Update the semantic tokens map.
409        let mut semantic_tokens = vec![];
410        let mut last_position = Position::new(0, 0);
411        for token in tokens.as_slice() {
412            let Ok(token_type) = SemanticTokenType::try_from(token.token_type) else {
413                // We continue here because not all tokens can be converted this way, we will get
414                // the rest from the ast.
415                continue;
416            };
417
418            let mut token_type_index = match self.get_semantic_token_type_index(&token_type) {
419                Some(index) => index,
420                // This is actually bad this should not fail.
421                // The test for listing all semantic token types should make this never happen.
422                None => {
423                    self.client
424                        .log_message(
425                            MessageType::ERROR,
426                            format!("token type `{:?}` not accounted for", token_type),
427                        )
428                        .await;
429                    continue;
430                }
431            };
432
433            let source_range: SourceRange = token.into();
434            let position = source_range.start_to_lsp_position(&params.text);
435
436            // Calculate the token modifiers.
437            // Get the value at the current position.
438            let token_modifiers_bitset = if let Some(ast) = self.ast_map.get(params.uri.as_str()) {
439                let token_index = Arc::new(Mutex::new(token_type_index));
440                let modifier_index: Arc<Mutex<u32>> = Arc::new(Mutex::new(0));
441                crate::walk::walk(&ast.ast, |node: crate::walk::Node| {
442                    let Ok(node_range): Result<SourceRange, _> = (&node).try_into() else {
443                        return Ok(true);
444                    };
445
446                    if !node_range.contains(source_range.start()) {
447                        return Ok(true);
448                    }
449
450                    let get_modifier = |modifier: Vec<SemanticTokenModifier>| -> Result<bool> {
451                        let mut mods = modifier_index.lock().map_err(|_| anyhow::anyhow!("mutex"))?;
452                        let Some(token_modifier_index) = self.get_semantic_token_modifier_index(modifier) else {
453                            return Ok(true);
454                        };
455                        if *mods == 0 {
456                            *mods = token_modifier_index;
457                        } else {
458                            *mods |= token_modifier_index;
459                        }
460                        Ok(false)
461                    };
462
463                    match node {
464                        crate::walk::Node::TagDeclarator(_) => {
465                            return get_modifier(vec![
466                                SemanticTokenModifier::DEFINITION,
467                                SemanticTokenModifier::STATIC,
468                            ]);
469                        }
470                        crate::walk::Node::VariableDeclarator(variable) => {
471                            let sr: SourceRange = (&variable.id).into();
472                            if sr.contains(source_range.start()) {
473                                if let Expr::FunctionExpression(_) = &variable.init {
474                                    let mut ti = token_index.lock().map_err(|_| anyhow::anyhow!("mutex"))?;
475                                    *ti = match self.get_semantic_token_type_index(&SemanticTokenType::FUNCTION) {
476                                        Some(index) => index,
477                                        None => token_type_index,
478                                    };
479                                }
480
481                                return get_modifier(vec![
482                                    SemanticTokenModifier::DECLARATION,
483                                    SemanticTokenModifier::READONLY,
484                                ]);
485                            }
486                        }
487                        crate::walk::Node::Parameter(_) => {
488                            let mut ti = token_index.lock().map_err(|_| anyhow::anyhow!("mutex"))?;
489                            *ti = match self.get_semantic_token_type_index(&SemanticTokenType::PARAMETER) {
490                                Some(index) => index,
491                                None => token_type_index,
492                            };
493                            return Ok(false);
494                        }
495                        crate::walk::Node::MemberExpression(member_expression) => {
496                            let sr: SourceRange = (&member_expression.property).into();
497                            if sr.contains(source_range.start()) {
498                                let mut ti = token_index.lock().map_err(|_| anyhow::anyhow!("mutex"))?;
499                                *ti = match self.get_semantic_token_type_index(&SemanticTokenType::PROPERTY) {
500                                    Some(index) => index,
501                                    None => token_type_index,
502                                };
503                                return Ok(false);
504                            }
505                        }
506                        crate::walk::Node::ObjectProperty(object_property) => {
507                            let sr: SourceRange = (&object_property.key).into();
508                            if sr.contains(source_range.start()) {
509                                let mut ti = token_index.lock().map_err(|_| anyhow::anyhow!("mutex"))?;
510                                *ti = match self.get_semantic_token_type_index(&SemanticTokenType::PROPERTY) {
511                                    Some(index) => index,
512                                    None => token_type_index,
513                                };
514                            }
515                            return get_modifier(vec![SemanticTokenModifier::DECLARATION]);
516                        }
517                        crate::walk::Node::CallExpression(call_expr) => {
518                            let sr: SourceRange = (&call_expr.callee).into();
519                            if sr.contains(source_range.start()) {
520                                let mut ti = token_index.lock().map_err(|_| anyhow::anyhow!("mutex"))?;
521                                *ti = match self.get_semantic_token_type_index(&SemanticTokenType::FUNCTION) {
522                                    Some(index) => index,
523                                    None => token_type_index,
524                                };
525
526                                if self.stdlib_completions.contains_key(&call_expr.callee.name) {
527                                    // This is a stdlib function.
528                                    return get_modifier(vec![SemanticTokenModifier::DEFAULT_LIBRARY]);
529                                }
530
531                                return Ok(false);
532                            }
533                        }
534                        _ => {}
535                    }
536                    Ok(true)
537                })
538                .unwrap_or_default();
539
540                let t = if let Ok(guard) = token_index.lock() { *guard } else { 0 };
541                token_type_index = t;
542
543                let m = if let Ok(guard) = modifier_index.lock() {
544                    *guard
545                } else {
546                    0
547                };
548                m
549            } else {
550                0
551            };
552
553            // We need to check if we are on the last token of the line.
554            // If we are starting from the end of the last line just add 1 to the line.
555            // Check if we are on the last token of the line.
556            if let Some(line) = params.text.lines().nth(position.line as usize) {
557                if line.len() == position.character as usize {
558                    // We are on the last token of the line.
559                    // We need to add a new line.
560                    let semantic_token = SemanticToken {
561                        delta_line: position.line - last_position.line + 1,
562                        delta_start: 0,
563                        length: (token.end - token.start) as u32,
564                        token_type: token_type_index,
565                        token_modifiers_bitset,
566                    };
567
568                    semantic_tokens.push(semantic_token);
569
570                    last_position = Position::new(position.line + 1, 0);
571                    continue;
572                }
573            }
574
575            let semantic_token = SemanticToken {
576                delta_line: position.line - last_position.line,
577                delta_start: if position.line != last_position.line {
578                    position.character
579                } else {
580                    position.character - last_position.character
581                },
582                length: (token.end - token.start) as u32,
583                token_type: token_type_index,
584                token_modifiers_bitset,
585            };
586
587            semantic_tokens.push(semantic_token);
588
589            last_position = position;
590        }
591        self.semantic_tokens_map.insert(params.uri.to_string(), semantic_tokens);
592    }
593
594    async fn clear_diagnostics_map(&self, uri: &url::Url, severity: Option<DiagnosticSeverity>) {
595        let Some(mut items) = self.diagnostics_map.get_mut(uri.as_str()) else {
596            return;
597        };
598
599        // If we only want to clear a specific severity, do that.
600        if let Some(severity) = severity {
601            items.retain(|x| x.severity != Some(severity));
602        } else {
603            items.clear();
604        }
605
606        if items.is_empty() {
607            #[cfg(not(target_arch = "wasm32"))]
608            {
609                self.client.publish_diagnostics(uri.clone(), items.clone(), None).await;
610            }
611
612            // We need to drop the items here.
613            drop(items);
614
615            self.diagnostics_map.remove(uri.as_str());
616        } else {
617            // We don't need to update the map since we used get_mut.
618
619            #[cfg(not(target_arch = "wasm32"))]
620            {
621                self.client.publish_diagnostics(uri.clone(), items.clone(), None).await;
622            }
623        }
624    }
625
626    async fn add_to_diagnostics<DiagT: IntoDiagnostic + std::fmt::Debug>(
627        &self,
628        params: &TextDocumentItem,
629        diagnostics: &[DiagT],
630        clear_all_before_add: bool,
631    ) {
632        if diagnostics.is_empty() {
633            return;
634        }
635
636        if clear_all_before_add {
637            self.clear_diagnostics_map(&params.uri, None).await;
638        } else if diagnostics.iter().all(|x| x.severity() == DiagnosticSeverity::ERROR) {
639            // If the diagnostic is an error, it will be the only error we get since that halts
640            // execution.
641            // Clear the diagnostics before we add a new one.
642            self.clear_diagnostics_map(&params.uri, Some(DiagnosticSeverity::ERROR))
643                .await;
644        } else if diagnostics
645            .iter()
646            .all(|x| x.severity() == DiagnosticSeverity::INFORMATION)
647        {
648            // If the diagnostic is a lint, we will pass them all to add at once so we need to
649            // clear the old ones.
650            self.clear_diagnostics_map(&params.uri, Some(DiagnosticSeverity::INFORMATION))
651                .await;
652        }
653
654        let mut items = if let Some(items) = self.diagnostics_map.get(params.uri.as_str()) {
655            // TODO: Would be awesome to fix the clone here.
656            items.clone()
657        } else {
658            vec![]
659        };
660
661        for diagnostic in diagnostics {
662            let lsp_d = diagnostic.to_lsp_diagnostics(&params.text);
663            // Make sure we don't duplicate diagnostics.
664            for d in lsp_d {
665                if !items.iter().any(|x| x == &d) {
666                    items.push(d);
667                }
668            }
669        }
670
671        self.diagnostics_map.insert(params.uri.to_string(), items.clone());
672
673        self.client.publish_diagnostics(params.uri.clone(), items, None).await;
674    }
675
676    async fn execute(&self, params: &TextDocumentItem, ast: &Program) -> Result<()> {
677        // Check if we can execute.
678        if !self.can_execute().await {
679            return Ok(());
680        }
681
682        // Execute the code if we have an executor context.
683        let ctx = self.executor_ctx().await;
684        let Some(ref executor_ctx) = *ctx else {
685            return Ok(());
686        };
687
688        if !self.is_initialized().await {
689            // We are not initialized yet.
690            return Ok(());
691        }
692
693        match executor_ctx.run_with_caching(ast.clone()).await {
694            Err(err) => {
695                self.add_to_diagnostics(params, &[err], false).await;
696
697                // Since we already published the diagnostics we don't really care about the error
698                // string.
699                Err(anyhow::anyhow!("failed to execute code"))
700            }
701            Ok(_) => Ok(()),
702        }
703    }
704
705    pub fn get_semantic_token_type_index(&self, token_type: &SemanticTokenType) -> Option<u32> {
706        SEMANTIC_TOKEN_TYPES
707            .iter()
708            .position(|x| *x == *token_type)
709            .map(|y| y as u32)
710    }
711
712    pub fn get_semantic_token_modifier_index(&self, token_types: Vec<SemanticTokenModifier>) -> Option<u32> {
713        if token_types.is_empty() {
714            return None;
715        }
716
717        let mut modifier = None;
718        for token_type in token_types {
719            if let Some(index) = SEMANTIC_TOKEN_MODIFIERS
720                .iter()
721                .position(|x| *x == token_type)
722                .map(|y| y as u32)
723            {
724                modifier = match modifier {
725                    Some(modifier) => Some(modifier | index),
726                    None => Some(index),
727                };
728            }
729        }
730        modifier
731    }
732
733    pub async fn create_zip(&self) -> Result<Vec<u8>> {
734        // Collect all the file data we know.
735        let mut buf = vec![];
736        let mut zip = zip::ZipWriter::new(std::io::Cursor::new(&mut buf));
737        for code in self.code_map.iter() {
738            let entry = code.key();
739            let value = code.value();
740            let file_name = entry.replace("file://", "").to_string();
741
742            let options = zip::write::SimpleFileOptions::default().compression_method(zip::CompressionMethod::Stored);
743            zip.start_file(file_name, options)?;
744            zip.write_all(value)?;
745        }
746        // Apply the changes you've made.
747        // Dropping the `ZipWriter` will have the same effect, but may silently fail
748        zip.finish()?;
749
750        Ok(buf)
751    }
752
753    pub async fn send_telemetry(&self) -> Result<()> {
754        // Get information about the user.
755        let user = self
756            .zoo_client
757            .users()
758            .get_self()
759            .await
760            .map_err(|e| anyhow::anyhow!(e.to_string()))?;
761
762        // Hash the user's id.
763        // Create a SHA-256 object
764        let mut hasher = sha2::Sha256::new();
765        // Write input message
766        hasher.update(user.id);
767        // Read hash digest and consume hasher
768        let result = hasher.finalize();
769        // Get the hash as a string.
770        let user_id_hash = format!("{:x}", result);
771
772        // Get the workspace folders.
773        // The key of the workspace folder is the project name.
774        let workspace_folders = self.workspace_folders().await;
775        let project_names: Vec<&str> = workspace_folders.iter().map(|v| v.name.as_str()).collect::<Vec<_>>();
776        // Get the first name.
777        let project_name = project_names
778            .first()
779            .ok_or_else(|| anyhow::anyhow!("no project names"))?
780            .to_string();
781
782        // Send the telemetry data.
783        self.zoo_client
784            .meta()
785            .create_event(
786                vec![kittycad::types::multipart::Attachment {
787                    // Clean the URI part.
788                    name: "attachment".to_string(),
789                    filename: Some("attachment.zip".to_string()),
790                    content_type: Some("application/x-zip".to_string()),
791                    data: self.create_zip().await?,
792                }],
793                &kittycad::types::Event {
794                    // This gets generated server side so leave empty for now.
795                    attachment_uri: None,
796                    created_at: chrono::Utc::now(),
797                    event_type: kittycad::types::ModelingAppEventType::SuccessfulCompileBeforeClose,
798                    last_compiled_at: Some(chrono::Utc::now()),
799                    // We do not have project descriptions yet.
800                    project_description: None,
801                    project_name,
802                    // The UUID for the modeling app.
803                    // We can unwrap here because we know it will not panic.
804                    source_id: uuid::Uuid::from_str("70178592-dfca-47b3-bd2d-6fce2bcaee04").unwrap(),
805                    type_: kittycad::types::Type::ModelingAppEvent,
806                    user_id: user_id_hash,
807                },
808            )
809            .await
810            .map_err(|e| anyhow::anyhow!(e.to_string()))?;
811
812        Ok(())
813    }
814
815    pub async fn update_units(
816        &self,
817        params: custom_notifications::UpdateUnitsParams,
818    ) -> RpcResult<Option<custom_notifications::UpdateUnitsResponse>> {
819        {
820            let mut ctx = self.executor_ctx.write().await;
821            // Borrow the executor context mutably.
822            let Some(ref mut executor_ctx) = *ctx else {
823                self.client
824                    .log_message(MessageType::ERROR, "no executor context set to update units for")
825                    .await;
826                return Ok(None);
827            };
828
829            self.client
830                .log_message(MessageType::INFO, format!("update units: {:?}", params))
831                .await;
832
833            if executor_ctx.settings.units == params.units
834                && !self.has_diagnostics(params.text_document.uri.as_ref()).await
835            {
836                // Return early the units are the same.
837                return Ok(None);
838            }
839
840            // Set the engine units.
841            executor_ctx.update_units(params.units);
842        }
843        // Lock is dropped here since nested.
844        // This is IMPORTANT.
845
846        let new_params = TextDocumentItem {
847            uri: params.text_document.uri.clone(),
848            text: std::mem::take(&mut params.text.to_string()),
849            version: Default::default(),
850            language_id: Default::default(),
851        };
852
853        // Force re-execution.
854        self.inner_on_change(new_params, true).await;
855
856        // Check if we have diagnostics.
857        // If we do we return early, since we failed in some way.
858        if self.has_diagnostics(params.text_document.uri.as_ref()).await {
859            return Ok(None);
860        }
861
862        Ok(Some(custom_notifications::UpdateUnitsResponse {}))
863    }
864
865    pub async fn update_can_execute(
866        &self,
867        params: custom_notifications::UpdateCanExecuteParams,
868    ) -> RpcResult<custom_notifications::UpdateCanExecuteResponse> {
869        let mut can_execute = self.can_execute.write().await;
870
871        if *can_execute == params.can_execute {
872            return Ok(custom_notifications::UpdateCanExecuteResponse {});
873        }
874
875        *can_execute = params.can_execute;
876
877        Ok(custom_notifications::UpdateCanExecuteResponse {})
878    }
879}
880
881#[tower_lsp::async_trait]
882impl LanguageServer for Backend {
883    async fn initialize(&self, params: InitializeParams) -> RpcResult<InitializeResult> {
884        self.client
885            .log_message(MessageType::INFO, format!("initialize: {:?}", params))
886            .await;
887
888        Ok(InitializeResult {
889            capabilities: ServerCapabilities {
890                completion_provider: Some(CompletionOptions {
891                    resolve_provider: Some(false),
892                    trigger_characters: Some(vec![".".to_string()]),
893                    work_done_progress_options: Default::default(),
894                    all_commit_characters: None,
895                    ..Default::default()
896                }),
897                diagnostic_provider: Some(DiagnosticServerCapabilities::Options(DiagnosticOptions {
898                    ..Default::default()
899                })),
900                document_formatting_provider: Some(OneOf::Left(true)),
901                folding_range_provider: Some(FoldingRangeProviderCapability::Simple(true)),
902                hover_provider: Some(HoverProviderCapability::Simple(true)),
903                inlay_hint_provider: Some(OneOf::Left(true)),
904                rename_provider: Some(OneOf::Left(true)),
905                semantic_tokens_provider: Some(SemanticTokensServerCapabilities::SemanticTokensRegistrationOptions(
906                    SemanticTokensRegistrationOptions {
907                        text_document_registration_options: {
908                            TextDocumentRegistrationOptions {
909                                document_selector: Some(vec![DocumentFilter {
910                                    language: Some("kcl".to_string()),
911                                    scheme: Some("file".to_string()),
912                                    pattern: None,
913                                }]),
914                            }
915                        },
916                        semantic_tokens_options: SemanticTokensOptions {
917                            work_done_progress_options: WorkDoneProgressOptions::default(),
918                            legend: SemanticTokensLegend {
919                                token_types: SEMANTIC_TOKEN_TYPES.to_vec(),
920                                token_modifiers: SEMANTIC_TOKEN_MODIFIERS.to_vec(),
921                            },
922                            range: Some(false),
923                            full: Some(SemanticTokensFullOptions::Bool(true)),
924                        },
925                        static_registration_options: StaticRegistrationOptions::default(),
926                    },
927                )),
928                signature_help_provider: Some(SignatureHelpOptions {
929                    trigger_characters: None,
930                    retrigger_characters: None,
931                    ..Default::default()
932                }),
933                text_document_sync: Some(TextDocumentSyncCapability::Options(TextDocumentSyncOptions {
934                    open_close: Some(true),
935                    change: Some(TextDocumentSyncKind::FULL),
936                    ..Default::default()
937                })),
938                workspace: Some(WorkspaceServerCapabilities {
939                    workspace_folders: Some(WorkspaceFoldersServerCapabilities {
940                        supported: Some(true),
941                        change_notifications: Some(OneOf::Left(true)),
942                    }),
943                    file_operations: None,
944                }),
945                ..Default::default()
946            },
947            ..Default::default()
948        })
949    }
950
951    async fn initialized(&self, params: InitializedParams) {
952        self.do_initialized(params).await
953    }
954
955    async fn shutdown(&self) -> RpcResult<()> {
956        self.do_shutdown().await
957    }
958
959    async fn did_change_workspace_folders(&self, params: DidChangeWorkspaceFoldersParams) {
960        self.do_did_change_workspace_folders(params).await
961    }
962
963    async fn did_change_configuration(&self, params: DidChangeConfigurationParams) {
964        self.do_did_change_configuration(params).await
965    }
966
967    async fn did_change_watched_files(&self, params: DidChangeWatchedFilesParams) {
968        self.do_did_change_watched_files(params).await
969    }
970
971    async fn did_create_files(&self, params: CreateFilesParams) {
972        self.do_did_create_files(params).await
973    }
974
975    async fn did_rename_files(&self, params: RenameFilesParams) {
976        self.do_did_rename_files(params).await
977    }
978
979    async fn did_delete_files(&self, params: DeleteFilesParams) {
980        self.do_did_delete_files(params).await
981    }
982
983    async fn did_open(&self, params: DidOpenTextDocumentParams) {
984        self.do_did_open(params).await
985    }
986
987    async fn did_change(&self, params: DidChangeTextDocumentParams) {
988        self.do_did_change(params).await;
989    }
990
991    async fn did_save(&self, params: DidSaveTextDocumentParams) {
992        self.do_did_save(params).await
993    }
994
995    async fn did_close(&self, params: DidCloseTextDocumentParams) {
996        self.do_did_close(params).await;
997
998        // Inject telemetry if we can train on the user's code.
999        // Return early if we cannot.
1000        if !self.can_send_telemetry {
1001            return;
1002        }
1003
1004        // In wasm this needs to be spawn_local since fucking reqwests doesn't implement Send for wasm.
1005        #[cfg(target_arch = "wasm32")]
1006        {
1007            let be = self.clone();
1008            wasm_bindgen_futures::spawn_local(async move {
1009                if let Err(err) = be.send_telemetry().await {
1010                    be.client
1011                        .log_message(MessageType::WARNING, format!("failed to send telemetry: {}", err))
1012                        .await;
1013                }
1014            });
1015        }
1016        #[cfg(not(target_arch = "wasm32"))]
1017        if let Err(err) = self.send_telemetry().await {
1018            self.client
1019                .log_message(MessageType::WARNING, format!("failed to send telemetry: {}", err))
1020                .await;
1021        }
1022    }
1023
1024    async fn hover(&self, params: HoverParams) -> RpcResult<Option<LspHover>> {
1025        let filename = params.text_document_position_params.text_document.uri.to_string();
1026
1027        let Some(current_code) = self.code_map.get(&filename) else {
1028            return Ok(None);
1029        };
1030        let Ok(current_code) = std::str::from_utf8(&current_code) else {
1031            return Ok(None);
1032        };
1033
1034        let pos = position_to_char_index(params.text_document_position_params.position, current_code);
1035
1036        // Let's iterate over the AST and find the node that contains the cursor.
1037        let Some(ast) = self.ast_map.get(&filename) else {
1038            return Ok(None);
1039        };
1040
1041        let Some(hover) = ast
1042            .ast
1043            .get_hover_value_for_position(pos, current_code, &HoverOpts::default_for_hover())
1044        else {
1045            return Ok(None);
1046        };
1047
1048        match hover {
1049            Hover::Function { name, range } => {
1050                let (sig, docs) = if let Some(Some(result)) = with_cached_var(&name, |value| {
1051                    match value {
1052                        // User-defined or KCL std function
1053                        KclValue::Function {
1054                            value: FunctionSource::User { ast, .. },
1055                            ..
1056                        } => {
1057                            // TODO get docs from comments
1058                            Some((ast.signature(), ""))
1059                        }
1060                        _ => None,
1061                    }
1062                })
1063                .await
1064                {
1065                    result
1066                } else {
1067                    // Get the docs for this function.
1068                    let Some(completion) = self.stdlib_completions.get(&name) else {
1069                        return Ok(None);
1070                    };
1071                    let Some(docs) = &completion.documentation else {
1072                        return Ok(None);
1073                    };
1074
1075                    let docs = match docs {
1076                        Documentation::String(docs) => docs,
1077                        Documentation::MarkupContent(MarkupContent { value, .. }) => value,
1078                    };
1079
1080                    let docs = if docs.len() > 320 {
1081                        let end = docs.find("\n\n").or_else(|| docs.find("\n\r\n")).unwrap_or(320);
1082                        &docs[..end]
1083                    } else {
1084                        &**docs
1085                    };
1086
1087                    let Some(label_details) = &completion.label_details else {
1088                        return Ok(None);
1089                    };
1090
1091                    let sig = if let Some(detail) = &label_details.detail {
1092                        detail.clone()
1093                    } else {
1094                        String::new()
1095                    };
1096
1097                    (sig, docs)
1098                };
1099
1100                Ok(Some(LspHover {
1101                    contents: HoverContents::Markup(MarkupContent {
1102                        kind: MarkupKind::Markdown,
1103                        value: format!("```\n{}{}\n```\n\n{}", name, sig, docs),
1104                    }),
1105                    range: Some(range),
1106                }))
1107            }
1108            Hover::Type { name, range } => {
1109                let Some(completion) = self.stdlib_completions.get(&name) else {
1110                    return Ok(None);
1111                };
1112                let Some(docs) = &completion.documentation else {
1113                    return Ok(None);
1114                };
1115
1116                let docs = match docs {
1117                    Documentation::String(docs) => docs,
1118                    Documentation::MarkupContent(MarkupContent { value, .. }) => value,
1119                };
1120
1121                let docs = if docs.len() > 320 {
1122                    let end = docs.find("\n\n").or_else(|| docs.find("\n\r\n")).unwrap_or(320);
1123                    &docs[..end]
1124                } else {
1125                    &**docs
1126                };
1127
1128                Ok(Some(LspHover {
1129                    contents: HoverContents::Markup(MarkupContent {
1130                        kind: MarkupKind::Markdown,
1131                        value: format!("```\n{}\n```\n\n{}", name, docs),
1132                    }),
1133                    range: Some(range),
1134                }))
1135            }
1136            Hover::KwArg {
1137                name,
1138                callee_name,
1139                range,
1140            } => {
1141                // TODO handle user-defined functions too
1142
1143                let Some(arg_map) = self.stdlib_args.get(&callee_name) else {
1144                    return Ok(None);
1145                };
1146
1147                let Some(tip) = arg_map.get(&name) else {
1148                    return Ok(None);
1149                };
1150
1151                Ok(Some(LspHover {
1152                    contents: HoverContents::Markup(MarkupContent {
1153                        kind: MarkupKind::Markdown,
1154                        value: tip.clone(),
1155                    }),
1156                    range: Some(range),
1157                }))
1158            }
1159            Hover::Variable {
1160                name,
1161                ty: Some(ty),
1162                range,
1163            } => Ok(Some(LspHover {
1164                contents: HoverContents::Markup(MarkupContent {
1165                    kind: MarkupKind::Markdown,
1166                    value: format!("```\n{}: {}\n```", name, ty),
1167                }),
1168                range: Some(range),
1169            })),
1170            Hover::Variable { name, ty: None, range } => Ok(with_cached_var(&name, |value| {
1171                let mut text: String = format!("```\n{}", name);
1172                if let Some(ty) = value.principal_type() {
1173                    text.push_str(&format!(": {}", ty.human_friendly_type()));
1174                }
1175                if let Some(v) = value.value_str() {
1176                    text.push_str(&format!(" = {}", v));
1177                }
1178                text.push_str("\n```");
1179
1180                LspHover {
1181                    contents: HoverContents::Markup(MarkupContent {
1182                        kind: MarkupKind::Markdown,
1183                        value: text,
1184                    }),
1185                    range: Some(range),
1186                }
1187            })
1188            .await),
1189            Hover::Signature { .. } => Ok(None),
1190            Hover::Comment { value, range } => Ok(Some(LspHover {
1191                contents: HoverContents::Markup(MarkupContent {
1192                    kind: MarkupKind::Markdown,
1193                    value,
1194                }),
1195                range: Some(range),
1196            })),
1197        }
1198    }
1199
1200    async fn completion(&self, params: CompletionParams) -> RpcResult<Option<CompletionResponse>> {
1201        let mut completions = vec![CompletionItem {
1202            label: PIPE_OPERATOR.to_string(),
1203            label_details: None,
1204            kind: Some(CompletionItemKind::OPERATOR),
1205            detail: Some("A pipe operator.".to_string()),
1206            documentation: Some(Documentation::MarkupContent(MarkupContent {
1207                kind: MarkupKind::Markdown,
1208                value: "A pipe operator.".to_string(),
1209            })),
1210            deprecated: Some(false),
1211            preselect: None,
1212            sort_text: None,
1213            filter_text: None,
1214            insert_text: Some("|> ".to_string()),
1215            insert_text_format: Some(InsertTextFormat::PLAIN_TEXT),
1216            insert_text_mode: None,
1217            text_edit: None,
1218            additional_text_edits: None,
1219            command: None,
1220            commit_characters: None,
1221            data: None,
1222            tags: None,
1223        }];
1224
1225        // Get the current line up to cursor
1226        let Some(current_code) = self
1227            .code_map
1228            .get(params.text_document_position.text_document.uri.as_ref())
1229        else {
1230            return Ok(Some(CompletionResponse::Array(completions)));
1231        };
1232        let Ok(current_code) = std::str::from_utf8(&current_code) else {
1233            return Ok(Some(CompletionResponse::Array(completions)));
1234        };
1235
1236        // Get the current line up to cursor, with bounds checking
1237        if let Some(line) = current_code
1238            .lines()
1239            .nth(params.text_document_position.position.line as usize)
1240        {
1241            let char_pos = params.text_document_position.position.character as usize;
1242            if char_pos <= line.len() {
1243                let line_prefix = &line[..char_pos];
1244                // Get last word
1245                let last_word = line_prefix
1246                    .split(|c: char| c.is_whitespace() || c.is_ascii_punctuation())
1247                    .last()
1248                    .unwrap_or("");
1249
1250                // If the last word starts with a digit, return no completions
1251                if !last_word.is_empty() && last_word.chars().next().unwrap().is_ascii_digit() {
1252                    return Ok(None);
1253                }
1254            }
1255        }
1256
1257        completions.extend(self.stdlib_completions.values().cloned());
1258
1259        // Add more to the completions if we have more.
1260        let Some(ast) = self
1261            .ast_map
1262            .get(params.text_document_position.text_document.uri.as_ref())
1263        else {
1264            return Ok(Some(CompletionResponse::Array(completions)));
1265        };
1266
1267        let Some(current_code) = self
1268            .code_map
1269            .get(params.text_document_position.text_document.uri.as_ref())
1270        else {
1271            return Ok(Some(CompletionResponse::Array(completions)));
1272        };
1273        let Ok(current_code) = std::str::from_utf8(&current_code) else {
1274            return Ok(Some(CompletionResponse::Array(completions)));
1275        };
1276
1277        let position = position_to_char_index(params.text_document_position.position, current_code);
1278        if ast.ast.in_comment(position) {
1279            // If we are in a code comment we don't want to show completions.
1280            return Ok(None);
1281        }
1282
1283        // Get the completion items for the ast.
1284        let Ok(variables) = ast.ast.completion_items(position) else {
1285            return Ok(Some(CompletionResponse::Array(completions)));
1286        };
1287
1288        // Get our variables from our AST to include in our completions.
1289        completions.extend(variables);
1290
1291        Ok(Some(CompletionResponse::Array(completions)))
1292    }
1293
1294    async fn diagnostic(&self, params: DocumentDiagnosticParams) -> RpcResult<DocumentDiagnosticReportResult> {
1295        let filename = params.text_document.uri.to_string();
1296
1297        // Get the current diagnostics for this file.
1298        let Some(items) = self.diagnostics_map.get(&filename) else {
1299            // Send an empty report.
1300            return Ok(DocumentDiagnosticReportResult::Report(DocumentDiagnosticReport::Full(
1301                RelatedFullDocumentDiagnosticReport {
1302                    related_documents: None,
1303                    full_document_diagnostic_report: FullDocumentDiagnosticReport {
1304                        result_id: None,
1305                        items: vec![],
1306                    },
1307                },
1308            )));
1309        };
1310
1311        Ok(DocumentDiagnosticReportResult::Report(DocumentDiagnosticReport::Full(
1312            RelatedFullDocumentDiagnosticReport {
1313                related_documents: None,
1314                full_document_diagnostic_report: FullDocumentDiagnosticReport {
1315                    result_id: None,
1316                    items: items.clone(),
1317                },
1318            },
1319        )))
1320    }
1321
1322    async fn signature_help(&self, params: SignatureHelpParams) -> RpcResult<Option<SignatureHelp>> {
1323        let filename = params.text_document_position_params.text_document.uri.to_string();
1324
1325        let Some(current_code) = self.code_map.get(&filename) else {
1326            return Ok(None);
1327        };
1328        let Ok(current_code) = std::str::from_utf8(&current_code) else {
1329            return Ok(None);
1330        };
1331
1332        let pos = position_to_char_index(params.text_document_position_params.position, current_code);
1333
1334        // Let's iterate over the AST and find the node that contains the cursor.
1335        let Some(ast) = self.ast_map.get(&filename) else {
1336            return Ok(None);
1337        };
1338
1339        let Some(value) = ast.ast.get_expr_for_position(pos) else {
1340            return Ok(None);
1341        };
1342
1343        let Some(hover) =
1344            value.get_hover_value_for_position(pos, current_code, &HoverOpts::default_for_signature_help())
1345        else {
1346            return Ok(None);
1347        };
1348
1349        match hover {
1350            Hover::Function { name, range: _ } => {
1351                // Get the docs for this function.
1352                let Some(signature) = self.stdlib_signatures.get(&name) else {
1353                    return Ok(None);
1354                };
1355
1356                Ok(Some(signature.clone()))
1357            }
1358            Hover::Signature {
1359                name,
1360                parameter_index,
1361                range: _,
1362            } => {
1363                let Some(signature) = self.stdlib_signatures.get(&name) else {
1364                    return Ok(None);
1365                };
1366
1367                let mut signature = signature.clone();
1368
1369                signature.active_parameter = Some(parameter_index);
1370
1371                Ok(Some(signature))
1372            }
1373            _ => {
1374                return Ok(None);
1375            }
1376        }
1377    }
1378
1379    async fn inlay_hint(&self, _params: InlayHintParams) -> RpcResult<Option<Vec<InlayHint>>> {
1380        // TODO: do this
1381
1382        Ok(None)
1383    }
1384
1385    async fn semantic_tokens_full(&self, params: SemanticTokensParams) -> RpcResult<Option<SemanticTokensResult>> {
1386        let filename = params.text_document.uri.to_string();
1387
1388        let Some(semantic_tokens) = self.semantic_tokens_map.get(&filename) else {
1389            return Ok(None);
1390        };
1391
1392        Ok(Some(SemanticTokensResult::Tokens(SemanticTokens {
1393            result_id: None,
1394            data: semantic_tokens.clone(),
1395        })))
1396    }
1397
1398    async fn document_symbol(&self, params: DocumentSymbolParams) -> RpcResult<Option<DocumentSymbolResponse>> {
1399        let filename = params.text_document.uri.to_string();
1400
1401        let Some(symbols) = self.symbols_map.get(&filename) else {
1402            return Ok(None);
1403        };
1404
1405        Ok(Some(DocumentSymbolResponse::Nested(symbols.clone())))
1406    }
1407
1408    async fn formatting(&self, params: DocumentFormattingParams) -> RpcResult<Option<Vec<TextEdit>>> {
1409        let filename = params.text_document.uri.to_string();
1410
1411        let Some(current_code) = self.code_map.get(&filename) else {
1412            return Ok(None);
1413        };
1414        let Ok(current_code) = std::str::from_utf8(&current_code) else {
1415            return Ok(None);
1416        };
1417
1418        // Parse the ast.
1419        // I don't know if we need to do this again since it should be updated in the context.
1420        // But I figure better safe than sorry since this will write back out to the file.
1421        let module_id = ModuleId::default();
1422        let Ok(ast) = crate::parsing::parse_str(current_code, module_id).parse_errs_as_err() else {
1423            return Ok(None);
1424        };
1425        // Now recast it.
1426        let recast = ast.recast(
1427            &crate::parsing::ast::types::FormatOptions {
1428                tab_size: params.options.tab_size as usize,
1429                insert_final_newline: params.options.insert_final_newline.unwrap_or(false),
1430                use_tabs: !params.options.insert_spaces,
1431            },
1432            0,
1433        );
1434        let source_range = SourceRange::new(0, current_code.len(), module_id);
1435        let range = source_range.to_lsp_range(current_code);
1436        Ok(Some(vec![TextEdit {
1437            new_text: recast,
1438            range,
1439        }]))
1440    }
1441
1442    async fn rename(&self, params: RenameParams) -> RpcResult<Option<WorkspaceEdit>> {
1443        let filename = params.text_document_position.text_document.uri.to_string();
1444
1445        let Some(current_code) = self.code_map.get(&filename) else {
1446            return Ok(None);
1447        };
1448        let Ok(current_code) = std::str::from_utf8(&current_code) else {
1449            return Ok(None);
1450        };
1451
1452        // Parse the ast.
1453        // I don't know if we need to do this again since it should be updated in the context.
1454        // But I figure better safe than sorry since this will write back out to the file.
1455        let module_id = ModuleId::default();
1456        let Ok(mut ast) = crate::parsing::parse_str(current_code, module_id).parse_errs_as_err() else {
1457            return Ok(None);
1458        };
1459
1460        // Let's convert the position to a character index.
1461        let pos = position_to_char_index(params.text_document_position.position, current_code);
1462        // Now let's perform the rename on the ast.
1463        ast.rename_symbol(&params.new_name, pos);
1464        // Now recast it.
1465        let recast = ast.recast(&Default::default(), 0);
1466        let source_range = SourceRange::new(0, current_code.len() - 1, module_id);
1467        let range = source_range.to_lsp_range(current_code);
1468        Ok(Some(WorkspaceEdit {
1469            changes: Some(HashMap::from([(
1470                params.text_document_position.text_document.uri,
1471                vec![TextEdit {
1472                    new_text: recast,
1473                    range,
1474                }],
1475            )])),
1476            document_changes: None,
1477            change_annotations: None,
1478        }))
1479    }
1480
1481    async fn folding_range(&self, params: FoldingRangeParams) -> RpcResult<Option<Vec<FoldingRange>>> {
1482        let filename = params.text_document.uri.to_string();
1483
1484        // Get the ast.
1485        let Some(ast) = self.ast_map.get(&filename) else {
1486            return Ok(None);
1487        };
1488
1489        // Get the folding ranges.
1490        let folding_ranges = ast.ast.get_lsp_folding_ranges();
1491
1492        if folding_ranges.is_empty() {
1493            return Ok(None);
1494        }
1495
1496        Ok(Some(folding_ranges))
1497    }
1498
1499    async fn code_action(&self, params: CodeActionParams) -> RpcResult<Option<CodeActionResponse>> {
1500        let actions = params
1501            .context
1502            .diagnostics
1503            .into_iter()
1504            .filter_map(|diagnostic| {
1505                let (suggestion, range) = diagnostic.data.as_ref().and_then(|data| {
1506                    serde_json::from_value::<(Suggestion, tower_lsp::lsp_types::Range)>(data.clone()).ok()
1507                })?;
1508                let edit = TextEdit {
1509                    range,
1510                    new_text: suggestion.insert,
1511                };
1512                let changes = HashMap::from([(params.text_document.uri.clone(), vec![edit])]);
1513                Some(CodeActionOrCommand::CodeAction(CodeAction {
1514                    title: suggestion.title,
1515                    kind: Some(CodeActionKind::QUICKFIX),
1516                    diagnostics: Some(vec![diagnostic]),
1517                    edit: Some(WorkspaceEdit {
1518                        changes: Some(changes),
1519                        document_changes: None,
1520                        change_annotations: None,
1521                    }),
1522                    command: None,
1523                    is_preferred: Some(true),
1524                    disabled: None,
1525                    data: None,
1526                }))
1527            })
1528            .collect();
1529
1530        Ok(Some(actions))
1531    }
1532}
1533
1534/// Get completions from our stdlib.
1535pub fn get_completions_from_stdlib(
1536    stdlib: &crate::std::StdLib,
1537    kcl_std: &[DocData],
1538) -> Result<HashMap<String, CompletionItem>> {
1539    let mut completions = HashMap::new();
1540    let combined = stdlib.combined();
1541
1542    for internal_fn in combined.values() {
1543        completions.insert(internal_fn.name(), internal_fn.to_completion_item()?);
1544    }
1545
1546    for d in kcl_std {
1547        completions.insert(d.name().to_owned(), d.to_completion_item());
1548    }
1549
1550    let variable_kinds = VariableKind::to_completion_items();
1551    for variable_kind in variable_kinds {
1552        completions.insert(variable_kind.label.clone(), variable_kind);
1553    }
1554
1555    Ok(completions)
1556}
1557
1558/// Get signatures from our stdlib.
1559pub fn get_signatures_from_stdlib(stdlib: &crate::std::StdLib, kcl_std: &[DocData]) -> HashMap<String, SignatureHelp> {
1560    let mut signatures = HashMap::new();
1561    let combined = stdlib.combined();
1562
1563    for internal_fn in combined.values() {
1564        signatures.insert(internal_fn.name(), internal_fn.to_signature_help());
1565    }
1566
1567    for d in kcl_std {
1568        if let Some(sig) = d.to_signature_help() {
1569            signatures.insert(d.name().to_owned(), sig);
1570        }
1571    }
1572
1573    signatures
1574}
1575
1576/// Get signatures from our stdlib.
1577pub fn get_arg_maps_from_stdlib(
1578    stdlib: &crate::std::StdLib,
1579    kcl_std: &[DocData],
1580) -> HashMap<String, HashMap<String, String>> {
1581    let mut result = HashMap::new();
1582    let combined = stdlib.combined();
1583
1584    for internal_fn in combined.values() {
1585        if internal_fn.keyword_arguments() {
1586            let arg_map: HashMap<String, String> = internal_fn
1587                .args(false)
1588                .into_iter()
1589                .map(|data| {
1590                    let mut tip = "```\n".to_owned();
1591                    tip.push_str(&data.name.clone());
1592                    if !data.required {
1593                        tip.push('?');
1594                    }
1595                    if !data.type_.is_empty() {
1596                        tip.push_str(": ");
1597                        tip.push_str(&data.type_);
1598                    }
1599                    tip.push_str("\n```");
1600                    if !data.description.is_empty() {
1601                        tip.push_str("\n\n");
1602                        tip.push_str(&data.description);
1603                    }
1604                    (data.name, tip)
1605                })
1606                .collect();
1607            if !arg_map.is_empty() {
1608                result.insert(internal_fn.name(), arg_map);
1609            }
1610        }
1611    }
1612
1613    for _d in kcl_std {
1614        // TODO add KCL std fns
1615    }
1616
1617    result
1618}
1619
1620/// Convert a position to a character index from the start of the file.
1621fn position_to_char_index(position: Position, code: &str) -> usize {
1622    // Get the character position from the start of the file.
1623    let mut char_position = 0;
1624    for (index, line) in code.lines().enumerate() {
1625        if index == position.line as usize {
1626            char_position += position.character as usize;
1627            break;
1628        } else {
1629            char_position += line.len() + 1;
1630        }
1631    }
1632
1633    char_position
1634}
1635
1636async fn with_cached_var<T>(name: &str, f: impl Fn(&KclValue) -> T) -> Option<T> {
1637    let mem = cache::read_old_memory().await?;
1638    let value = mem.get(name, SourceRange::default()).ok()?;
1639
1640    Some(f(value))
1641}