Add showing KCL variables up to the point of the error (#7762)

* Add showing KCL variables up to the point of the error

* Add output of memory when sim tests fail

* Update output

* Fix so that sim tests can distinguish success from error

* Update output after adding success files

Author: jtran

rust/kcl-lib/src/errors.rs

@@ -8,6 +8,7 @@ use tower_lsp::lsp_types::{Diagnostic, DiagnosticSeverity};
 use crate::execution::{ArtifactCommand, ArtifactGraph, Operation};
 use crate::{
     ModuleId,
+    exec::KclValue,
     execution::DefaultPlanes,
     lsp::IntoDiagnostic,
     modules::{ModulePath, ModuleSource},
@@ -133,6 +134,9 @@ impl From<KclErrorWithOutputs> for KclError {
 pub struct KclErrorWithOutputs {
     pub error: KclError,
     pub non_fatal: Vec<CompilationError>,
+    /// Variables in the top-level of the root module. Note that functions will
+    /// have an invalid env ref.
+    pub variables: IndexMap<String, KclValue>,
     #[cfg(feature = "artifact-graph")]
     pub operations: Vec<Operation>,
     // TODO: Remove this field.  Doing so breaks the ts-rs output for some
@@ -151,6 +155,7 @@ impl KclErrorWithOutputs {
     pub fn new(
         error: KclError,
         non_fatal: Vec<CompilationError>,
+        variables: IndexMap<String, KclValue>,
         #[cfg(feature = "artifact-graph")] operations: Vec<Operation>,
         #[cfg(feature = "artifact-graph")] artifact_commands: Vec<ArtifactCommand>,
         #[cfg(feature = "artifact-graph")] artifact_graph: ArtifactGraph,
@@ -161,6 +166,7 @@ impl KclErrorWithOutputs {
         Self {
             error,
             non_fatal,
+            variables,
             #[cfg(feature = "artifact-graph")]
             operations,
             #[cfg(feature = "artifact-graph")]
@@ -176,6 +182,7 @@ impl KclErrorWithOutputs {
         Self {
             error,
             non_fatal: Default::default(),
+            variables: Default::default(),
             #[cfg(feature = "artifact-graph")]
             operations: Default::default(),
             #[cfg(feature = "artifact-graph")]

rust/kcl-lib/src/execution/exec_ast.rs

@@ -88,7 +88,7 @@ impl ExecutorContext {
         path: &ModulePath,
     ) -> Result<
         (Option<KclValue>, EnvironmentRef, Vec<String>, ModuleArtifactState),
-        (KclError, Option<ModuleArtifactState>),
+        (KclError, Option<EnvironmentRef>, Option<ModuleArtifactState>),
     > {
         crate::log::log(format!("enter module {path} {}", exec_state.stack()));
 
@@ -100,7 +100,7 @@ impl ExecutorContext {
         let no_prelude = self
             .handle_annotations(program.inner_attrs.iter(), crate::execution::BodyType::Root, exec_state)
             .await
-            .map_err(|err| (err, None))?;
+            .map_err(|err| (err, None, None))?;
 
         if !preserve_mem {
             exec_state.mut_stack().push_new_root_env(!no_prelude);
@@ -125,7 +125,7 @@ impl ExecutorContext {
         crate::log::log(format!("leave {path}"));
 
         result
-            .map_err(|err| (err, Some(module_artifacts.clone())))
+            .map_err(|err| (err, Some(env_ref), Some(module_artifacts.clone())))
             .map(|result| (result, env_ref, local_state.module_exports, module_artifacts))
     }
 
@@ -644,7 +644,7 @@ impl ExecutorContext {
 
         // TODO: ModuleArtifactState is getting dropped here when there's an
         // error.  Should we propagate it for non-root modules?
-        result.map_err(|(err, _)| {
+        result.map_err(|(err, _, _)| {
             if let KclError::ImportCycle { .. } = err {
                 // It was an import cycle.  Keep the original message.
                 err.override_source_ranges(vec![source_range])

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

@@ -842,7 +842,7 @@ impl ExecutorContext {
             .map_err(KclErrorWithOutputs::no_outputs)?;
 
         for modules in import_graph::import_graph(&universe, self)
-            .map_err(|err| exec_state.error_with_outputs(err, default_planes.clone()))?
+            .map_err(|err| exec_state.error_with_outputs(err, None, default_planes.clone()))?
             .into_iter()
         {
             #[cfg(not(target_arch = "wasm32"))]
@@ -982,7 +982,7 @@ impl ExecutorContext {
                         exec_state.global.module_infos[&module_id].restore_repr(repr);
                     }
                     Err(e) => {
-                        return Err(exec_state.error_with_outputs(e, default_planes));
+                        return Err(exec_state.error_with_outputs(e, None, default_planes));
                     }
                 }
             }
@@ -1020,7 +1020,7 @@ impl ExecutorContext {
             exec_state,
         )
         .await
-        .map_err(|err| exec_state.error_with_outputs(err, default_planes))?;
+        .map_err(|err| exec_state.error_with_outputs(err, None, default_planes))?;
 
         Ok((universe, root_imports))
     }
@@ -1130,7 +1130,7 @@ impl ExecutorContext {
         ));
         crate::log::log(format!("Engine stats: {:?}", self.engine.stats()));
 
-        let env_ref = result.map_err(|e| exec_state.error_with_outputs(e, default_planes))?;
+        let env_ref = result.map_err(|(err, env_ref)| exec_state.error_with_outputs(err, env_ref, default_planes))?;
 
         if !self.is_mock() {
             let mut mem = exec_state.stack().deep_clone();
@@ -1149,7 +1149,7 @@ impl ExecutorContext {
         program: NodeRef<'_, crate::parsing::ast::types::Program>,
         exec_state: &mut ExecState,
         preserve_mem: bool,
-    ) -> Result<EnvironmentRef, KclError> {
+    ) -> Result<EnvironmentRef, (KclError, Option<EnvironmentRef>)> {
         // Don't early return!  We need to build other outputs regardless of
         // whether execution failed.
 
@@ -1159,7 +1159,8 @@ impl ExecutorContext {
         let start_op = exec_state.global.root_module_artifacts.operations.len();
 
         self.eval_prelude(exec_state, SourceRange::from(program).start_as_range())
-            .await?;
+            .await
+            .map_err(|e| (e, None))?;
 
         let exec_result = self
             .exec_module_body(
@@ -1176,13 +1177,13 @@ impl ExecutorContext {
                 exec_state.global.root_module_artifacts.extend(module_artifacts);
                 env_ref
             })
-            .map_err(|(err, module_artifacts)| {
+            .map_err(|(err, env_ref, module_artifacts)| {
                 if let Some(module_artifacts) = module_artifacts {
                     // We need to extend because it may already have operations
                     // from imports.
                     exec_state.global.root_module_artifacts.extend(module_artifacts);
                 }
-                err
+                (err, env_ref)
             });
 
         #[cfg(feature = "artifact-graph")]
@@ -1208,15 +1209,21 @@ impl ExecutorContext {
         }
 
         // Ensure all the async commands completed.
-        self.engine.ensure_async_commands_completed().await?;
+        self.engine.ensure_async_commands_completed().await.map_err(|e| {
+            match &exec_result {
+                Ok(env_ref) => (e, Some(*env_ref)),
+                // Prefer the execution error.
+                Err((exec_err, env_ref)) => (exec_err.clone(), *env_ref),
+            }
+        })?;
 
         // If we errored out and early-returned, there might be commands which haven't been executed
         // and should be dropped.
         self.engine.clear_queues().await;
 
         match exec_state.build_artifact_graph(&self.engine, program).await {
             Ok(_) => exec_result,
-            Err(err) => exec_result.and(Err(err)),
+            Err(err) => exec_result.and_then(|env_ref| Err((err, Some(env_ref)))),
         }
     }
 

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

@@ -302,6 +302,7 @@ impl ExecState {
     pub(crate) fn error_with_outputs(
         &self,
         error: KclError,
+        main_ref: Option<EnvironmentRef>,
         default_planes: Option<DefaultPlanes>,
     ) -> KclErrorWithOutputs {
         let module_id_to_module_path: IndexMap<ModuleId, ModulePath> = self
@@ -314,6 +315,9 @@ impl ExecState {
         KclErrorWithOutputs::new(
             error,
             self.errors().to_vec(),
+            main_ref
+                .map(|main_ref| self.mod_local.variables(main_ref))
+                .unwrap_or_default(),
             #[cfg(feature = "artifact-graph")]
             self.global.root_module_artifacts.operations.clone(),
             #[cfg(feature = "artifact-graph")]

rust/kcl-lib/src/simulation_tests.rs

@@ -8,6 +8,7 @@ use indexmap::IndexMap;
 use crate::{
     ExecOutcome, ExecState, ExecutorContext, ModuleId,
     errors::KclError,
+    exec::KclValue,
     execution::{EnvironmentRef, ModuleArtifactState},
 };
 #[cfg(feature = "artifact-graph")]
@@ -259,24 +260,25 @@ async fn execute_test(test: &Test, render_to_png: bool, export_step: bool) {
                     panic!("Step data was empty");
                 }
             }
-            let (outcome, module_state) = exec_state.into_test_exec_outcome(env_ref, &ctx, &test.input_dir).await;
-
-            let mem_result = catch_unwind(AssertUnwindSafe(|| {
-                assert_snapshot(test, "Variables in memory after executing", || {
-                    insta::assert_json_snapshot!("program_memory", outcome.variables, {
-                         ".**.sourceRange" => Vec::new(),
-                    })
+            let ok_snap = catch_unwind(AssertUnwindSafe(|| {
+                assert_snapshot(test, "Execution success", || {
+                    insta::assert_json_snapshot!("execution_success", ())
                 })
             }));
 
+            let (outcome, module_state) = exec_state.into_test_exec_outcome(env_ref, &ctx, &test.input_dir).await;
+
+            assert_common_snapshots(test, outcome.variables);
+
             #[cfg(not(feature = "artifact-graph"))]
             drop(module_state);
             #[cfg(feature = "artifact-graph")]
             assert_artifact_snapshots(test, module_state, outcome.artifact_graph);
-            mem_result.unwrap();
+
+            ok_snap.unwrap();
         }
         Err(e) => {
-            let ok_path = test.output_dir.join("program_memory.snap");
+            let ok_path = test.output_dir.join("execution_success.snap");
             let previously_passed = std::fs::exists(&ok_path).unwrap();
             match e.error {
                 crate::errors::ExecError::Kcl(error) => {
@@ -304,6 +306,8 @@ async fn execute_test(test: &Test, render_to_png: bool, export_step: bool) {
                         })
                     }));
 
+                    assert_common_snapshots(test, error.variables);
+
                     #[cfg(feature = "artifact-graph")]
                     {
                         let module_state = e
@@ -325,6 +329,19 @@ async fn execute_test(test: &Test, render_to_png: bool, export_step: bool) {
     }
 }
 
+/// Assert snapshots that should happen both when KCL execution succeeds and
+/// when it results in an error.
+fn assert_common_snapshots(test: &Test, variables: IndexMap<String, KclValue>) {
+    let mem_result = catch_unwind(AssertUnwindSafe(|| {
+        assert_snapshot(test, "Variables in memory after executing", || {
+            insta::assert_json_snapshot!("program_memory", variables, {
+                 ".**.sourceRange" => Vec::new(),
+            })
+        })
+    }));
+    mem_result.unwrap();
+}
+
 /// Assert snapshots for artifacts that should happen both when KCL execution
 /// succeeds and when it results in an error.
 #[cfg(feature = "artifact-graph")]

rust/kcl-lib/tests/add_arrays/program_memory.snap

@@ -0,0 +1,5 @@
+---
+source: kcl-lib/src/simulation_tests.rs
+description: Variables in memory after executing add_arrays.kcl
+---
+{}

rust/kcl-lib/tests/add_lots/execution_success.snap

@@ -0,0 +1,5 @@
+---
+source: kcl-lib/src/simulation_tests.rs
+description: Execution success add_lots.kcl
+---
+null

rust/kcl-lib/tests/angled_line/execution_success.snap

@@ -0,0 +1,5 @@
+---
+source: kcl-lib/src/simulation_tests.rs
+description: Execution success angled_line.kcl
+---
+null

rust/kcl-lib/tests/any_type/execution_success.snap

@@ -0,0 +1,5 @@
+---
+source: kcl-lib/src/simulation_tests.rs
+description: Execution success any_type.kcl
+---
+null

rust/kcl-lib/tests/argument_error/program_memory.snap

@@ -0,0 +1,10 @@
+---
+source: kcl-lib/src/simulation_tests.rs
+description: Variables in memory after executing argument_error.kcl
+---
+{
+  "f": {
+    "type": "Function",
+    "value": null
+  }
+}