better track the root source code

src/wasm-lib/kcl/src/errors.rs

@@ -150,7 +150,7 @@
             source_files: Default::default(),
         }
     }
-    pub fn into_miette_report_with_outputs(self, code: &str) -> anyhow::Result<ReportWithOutputs> {
+    pub fn into_miette_report_with_outputs(self) -> anyhow::Result<ReportWithOutputs> {
         let mut source_ranges = self.error.source_ranges();
 
         // Pop off the first source range to get the filename.
@@ -162,33 +162,23 @@
             .source_files
             .get(&first_source_range.module_id())
             .cloned()
-            .unwrap_or(ModuleSource {
-                source: code.to_string(),
-                path: self
-                    .filenames
-                    .get(&first_source_range.module_id())
-                    .ok_or_else(|| {
-                        anyhow::anyhow!(
-                            "Could not find filename for module id: {:?}",
-                            first_source_range.module_id()
-                        )
-                    })?
-                    .clone(),
-            });
+            .ok_or_else(|| {
+                anyhow::anyhow!(
+                    "Could not find source file for module id: {:?}",
+                    first_source_range.module_id()
+                )
+            })?;
         let filename = source.path.to_string();
         let kcl_source = source.source.to_string();
 
         let mut related = Vec::new();
         for source_range in source_ranges {
             let module_id = source_range.module_id();
-            let source = self.source_files.get(&module_id).cloned().unwrap_or(ModuleSource {
-                source: code.to_string(),
-                path: self
-                    .filenames
-                    .get(&module_id)
-                    .ok_or_else(|| anyhow::anyhow!("Could not find filename for module id: {:?}", module_id))?
-                    .clone(),
-            });
+            let source = self
+                .source_files
+                .get(&module_id)
+                .cloned()
+                .ok_or_else(|| anyhow::anyhow!("Could not find source file for module id: {:?}", module_id))?;
             let error = self.error.override_source_ranges(vec![source_range]);
             let report = Report {
                 error,

src/wasm-lib/kcl/src/execution/mod.rs

@@ -532,6 +532,7 @@
             }
         }
 
+        exec_state.add_root_module_contents(&program);
         let result = self.inner_run(&program.ast, &mut exec_state, true).await?;
 
         // Restore any temporary variables, then save any newly created variables back to
@@ -585,9 +586,15 @@
                             .await
                             .is_err()
                     {
-                        (true, program.ast)
+                        (true, program)
                     } else {
-                        (clear_scene, changed_program)
+                        (
+                            clear_scene,
+                            crate::Program {
+                                ast: changed_program,
+                                original_file_contents: program.original_file_contents,
+                            },
+                        )
                     }
                 }
                 CacheResult::NoAction(true) => {
@@ -608,7 +615,7 @@
 
                         return Ok(old_state.to_wasm_outcome(result_env));
                     }
-                    (true, program.ast)
+                    (true, program)
                 }
                 CacheResult::NoAction(false) => return Ok(old_state.to_wasm_outcome(result_env)),
             };
@@ -636,10 +643,11 @@
             self.send_clear_scene(&mut exec_state, Default::default())
                 .await
                 .map_err(KclErrorWithOutputs::no_outputs)?;
-            (program.ast, exec_state, false)
+            (program, exec_state, false)
         };
 
-        let result = self.inner_run(&program, &mut exec_state, preserve_mem).await;
+        exec_state.add_root_module_contents(&program);
+        let result = self.inner_run(&program.ast, &mut exec_state, preserve_mem).await;
 
         if result.is_err() {
             cache::bust_cache().await;
@@ -650,7 +658,7 @@
 
         // Save this as the last successful execution to the cache.
         cache::write_old_ast(OldAstState {
-            ast: program,
+            ast: program.ast,
             exec_state: exec_state.clone(),
             settings: self.settings.clone(),
             result_env: result.0,
@@ -691,6 +699,7 @@
         self.send_clear_scene(exec_state, Default::default())
             .await
             .map_err(KclErrorWithOutputs::no_outputs)?;
+        exec_state.add_root_module_contents(program);
         self.inner_run(&program.ast, exec_state, false).await
     }
 

src/wasm-lib/kcl/src/execution/state.rs

@@ -182,8 +182,27 @@ impl ExecState {
         self.global.path_to_source_id.insert(path.clone(), id);
     }
 
+    pub(crate) fn add_root_module_contents(&mut self, program: &crate::Program) {
+        let root_id = ModuleId::default();
+        // Get the path for the root module.
+        let path = self
+            .global
+            .path_to_source_id
+            .iter()
+            .find(|(_, v)| **v == root_id)
+            .unwrap()
+            .0
+            .clone();
+        self.add_id_to_source(
+            root_id,
+            ModuleSource {
+                path,
+                source: program.original_file_contents.to_string(),
+            },
+        );
+    }
+
     pub(super) fn add_id_to_source(&mut self, id: ModuleId, source: ModuleSource) {
-        debug_assert!(!self.global.id_to_source.contains_key(&id));
         self.global.id_to_source.insert(id, source.clone());
     }
 
@@ -251,8 +270,6 @@ impl GlobalState {
         global
             .path_to_source_id
             .insert(ModulePath::Local { value: root_path }, root_id);
-        // Ideally we'd have a way to set the root module's source here, but
-        // we don't have a way to get the source from the executor settings.
         global
     }
 }

src/wasm-lib/kcl/src/lib.rs

@@ -134,6 +134,11 @@
 pub struct Program {
     #[serde(flatten)]
     pub ast: parsing::ast::types::Node<parsing::ast::types::Program>,
+    // The ui doesn't need to know about this.
+    // It's purely used for saving the contents of the original file, so we can use it for errors.
+    // Because in the case of the root file, we don't want to read the file from disk again.
+    #[serde(skip)]
+    pub original_file_contents: String,
 }
 
 #[cfg(any(test, feature = "lsp-test-util"))]
@@ -147,15 +152,24 @@
         let tokens = parsing::token::lex(input, module_id)?;
         let (ast, errs) = parsing::parse_tokens(tokens).0?;
 
-        Ok((ast.map(|ast| Program { ast }), errs))
+        Ok((
+            ast.map(|ast| Program {
+                ast,
+                original_file_contents: input.to_string(),
+            }),
+            errs,
+        ))
     }
 
     pub fn parse_no_errs(input: &str) -> Result<Program, KclError> {
         let module_id = ModuleId::default();
         let tokens = parsing::token::lex(input, module_id)?;
         let ast = parsing::parse_tokens(tokens).parse_errs_as_err()?;
 
-        Ok(Program { ast })
+        Ok(Program {
+            ast,
+            original_file_contents: input.to_string(),
+        })
     }
 
     pub fn compute_digest(&mut self) -> parsing::ast::digest::Digest {
@@ -168,9 +182,10 @@
     }
 
     /// Change the meta settings for the kcl file.
-    pub fn change_meta_settings(&mut self, settings: crate::MetaSettings) -> Result<Self, KclError> {
+    pub fn change_meta_settings(&self, settings: crate::MetaSettings) -> Result<Self, KclError> {
         Ok(Self {
             ast: self.ast.change_meta_settings(settings)?,
+            original_file_contents: self.original_file_contents.clone(),
         })
     }
 
@@ -192,12 +207,6 @@
     }
 }
 
-impl From<parsing::ast::types::Node<parsing::ast::types::Program>> for Program {
-    fn from(ast: parsing::ast::types::Node<parsing::ast::types::Program>) -> Program {
-        Self { ast }
-    }
-}
-
 #[inline]
 fn try_f64_to_usize(f: f64) -> Option<usize> {
     let i = f as usize;

src/wasm-lib/kcl/src/lsp/kcl/mod.rs

@@ -46,7 +46,7 @@ use crate::{
     errors::Suggestion,
     lsp::{backend::Backend as _, util::IntoDiagnostic},
     parsing::{
-        ast::types::{Expr, Node, VariableKind},
+        ast::types::{Expr, VariableKind},
         token::TokenStream,
         PIPE_OPERATOR,
     },
@@ -102,7 +102,7 @@ pub struct Backend {
     /// Token maps.
     pub(super) token_map: DashMap<String, TokenStream>,
     /// AST maps.
-    pub ast_map: DashMap<String, Node<crate::parsing::ast::types::Program>>,
+    pub ast_map: DashMap<String, crate::Program>,
     /// Current code.
     pub code_map: DashMap<String, Vec<u8>>,
     /// Diagnostics.
@@ -327,11 +327,17 @@ impl crate::lsp::backend::Backend for Backend {
         // this if it backfires and only hork the LSP.
         ast.compute_digest();
 
+        // Save it as a program.
+        let ast = crate::Program {
+            ast,
+            original_file_contents: params.text.clone(),
+        };
+
         // Check if the ast changed.
         let ast_changed = match self.ast_map.get(&filename) {
             Some(old_ast) => {
                 // Check if the ast changed.
-                *old_ast != ast
+                *old_ast.ast != *ast.ast
             }
             None => true,
         };
@@ -346,7 +352,7 @@ impl crate::lsp::backend::Backend for Backend {
             // Update the symbols map.
             self.symbols_map.insert(
                 params.uri.to_string(),
-                ast.get_lsp_symbols(&params.text).unwrap_or_default(),
+                ast.ast.get_lsp_symbols(&params.text).unwrap_or_default(),
             );
 
             // Update our semantic tokens.
@@ -361,14 +367,14 @@ impl crate::lsp::backend::Backend for Backend {
             // Only send the notification if we can execute.
             // Otherwise it confuses the client.
             self.client
-                .send_notification::<custom_notifications::AstUpdated>(ast.clone())
+                .send_notification::<custom_notifications::AstUpdated>(ast.ast.clone())
                 .await;
         }
 
         // Execute the code if we have an executor context.
         // This function automatically executes if we should & updates the diagnostics if we got
         // errors.
-        if self.execute(&params, &ast.into()).await.is_err() {
+        if self.execute(&params, &ast).await.is_err() {
             return;
         }
 
@@ -421,7 +427,7 @@ impl Backend {
             let token_modifiers_bitset = if let Some(ast) = self.ast_map.get(params.uri.as_str()) {
                 let token_index = Arc::new(Mutex::new(token_type_index));
                 let modifier_index: Arc<Mutex<u32>> = Arc::new(Mutex::new(0));
-                crate::walk::walk(&ast, |node: crate::walk::Node| {
+                crate::walk::walk(&ast.ast, |node: crate::walk::Node| {
                     let Ok(node_range): Result<SourceRange, _> = (&node).try_into() else {
                         return Ok(true);
                     };
@@ -1021,7 +1027,7 @@ impl LanguageServer for Backend {
             return Ok(None);
         };
 
-        let Some(hover) = ast.get_hover_value_for_position(pos, current_code) else {
+        let Some(hover) = ast.ast.get_hover_value_for_position(pos, current_code) else {
             return Ok(None);
         };
 
@@ -1150,13 +1156,13 @@ impl LanguageServer for Backend {
         };
 
         let position = position_to_char_index(params.text_document_position.position, current_code);
-        if ast.get_non_code_meta_for_position(position).is_some() {
+        if ast.ast.get_non_code_meta_for_position(position).is_some() {
             // If we are in a code comment we don't want to show completions.
             return Ok(None);
         }
 
         // Get the completion items for the ast.
-        let Ok(variables) = ast.completion_items() else {
+        let Ok(variables) = ast.ast.completion_items() else {
             return Ok(Some(CompletionResponse::Array(completions)));
         };
 
@@ -1211,7 +1217,7 @@ impl LanguageServer for Backend {
             return Ok(None);
         };
 
-        let Some(value) = ast.get_expr_for_position(pos) else {
+        let Some(value) = ast.ast.get_expr_for_position(pos) else {
             return Ok(None);
         };
 
@@ -1360,7 +1366,7 @@ impl LanguageServer for Backend {
         };
 
         // Get the folding ranges.
-        let folding_ranges = ast.get_lsp_folding_ranges();
+        let folding_ranges = ast.ast.get_lsp_folding_ranges();
 
         if folding_ranges.is_empty() {
             return Ok(None);