Skip to content

Commit

Permalink
better track the root source code (#5558)
Browse files Browse the repository at this point in the history 
* better track the root source code

Signed-off-by: Jess Frazelle <[email protected]>

* fix mut self

Signed-off-by: Jess Frazelle <[email protected]>

---------

Signed-off-by: Jess Frazelle <[email protected]>
  • Loading branch information
@jessfraz
jessfraz authored Feb 27, 2025
1 parent 850c5c6 commit 5ce22e2
Show file tree
Hide file tree
Showing 10 changed files with 127 additions and 93 deletions.
34 changes: 12 additions & 22 deletions src/wasm-lib/kcl/src/errors.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -150,7 +150,7 @@ impl KclErrorWithOutputs {
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.
Expand All @@ -162,33 +162,23 @@ impl KclErrorWithOutputs {
.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,
Expand Down
21 changes: 15 additions & 6 deletions src/wasm-lib/kcl/src/execution/mod.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -532,6 +532,7 @@ impl ExecutorContext {
}
}

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
Expand Down Expand Up @@ -585,9 +586,15 @@ impl ExecutorContext {
.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) => {
Expand All @@ -608,7 +615,7 @@ impl ExecutorContext {

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)),
};
Expand Down Expand Up @@ -636,10 +643,11 @@ impl ExecutorContext {
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;
Expand All @@ -650,7 +658,7 @@ impl ExecutorContext {

// 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,
Expand Down Expand Up @@ -691,6 +699,7 @@ impl ExecutorContext {
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
}

Expand Down
23 changes: 20 additions & 3 deletions src/wasm-lib/kcl/src/execution/state.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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());
}

Expand Down Expand Up @@ -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
}
}
Expand Down
27 changes: 18 additions & 9 deletions src/wasm-lib/kcl/src/lib.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -134,6 +134,11 @@ use crate::log::{log, logln};
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"))]
Expand All @@ -147,15 +152,24 @@ impl Program {
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 {
Expand All @@ -168,9 +182,10 @@ impl Program {
}

/// 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(),
})
}

Expand All @@ -192,12 +207,6 @@ impl Program {
}
}

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;
Expand Down
30 changes: 18 additions & 12 deletions src/wasm-lib/kcl/src/lsp/kcl/mod.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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,
},
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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,
};
Expand All @@ -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.
Expand All @@ -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;
}

Expand Down Expand Up @@ -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);
};
Expand Down Expand Up @@ -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);
};

Expand Down Expand Up @@ -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)));
};

Expand Down Expand Up @@ -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);
};

Expand Down Expand Up @@ -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);
Expand Down
Loading

0 comments on commit 5ce22e2

Please sign in to comment.