Execution refactoring (#7376)

* Move import graph to execution

Signed-off-by: Nick Cameron <nrc@ncameron.org>

* Refactor artifact handling

Signed-off-by: Nick Cameron <nrc@ncameron.org>

* Refactor caching to separate global state from per-module state

Signed-off-by: Nick Cameron <nrc@ncameron.org>

---------

Signed-off-by: Nick Cameron <nrc@ncameron.org>

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

@@ -676,6 +676,7 @@ impl EdgeCut {
 #[serde(rename_all = "camelCase")]
 pub struct ArtifactGraph {
     map: IndexMap<ArtifactId, Artifact>,
+    item_count: usize,
 }
 
 impl ArtifactGraph {
@@ -711,10 +712,10 @@ pub(super) fn build_artifact_graph(
     artifact_commands: &[ArtifactCommand],
     responses: &IndexMap<Uuid, WebSocketResponse>,
     ast: &Node<Program>,
-    cached_body_items: usize,
     exec_artifacts: &mut IndexMap<ArtifactId, Artifact>,
     initial_graph: ArtifactGraph,
 ) -> Result<ArtifactGraph, KclError> {
+    let item_count = initial_graph.item_count;
     let mut map = initial_graph.into_map();
 
     let mut path_to_plane_id_map = FnvHashMap::default();
@@ -725,7 +726,7 @@ pub(super) fn build_artifact_graph(
     for exec_artifact in exec_artifacts.values_mut() {
         // Note: We only have access to the new AST. So if these artifacts
         // somehow came from cached AST, this won't fill in anything.
-        fill_in_node_paths(exec_artifact, ast, cached_body_items);
+        fill_in_node_paths(exec_artifact, ast, item_count);
     }
 
     for artifact_command in artifact_commands {
@@ -752,7 +753,7 @@ pub(super) fn build_artifact_graph(
             &flattened_responses,
             &path_to_plane_id_map,
             ast,
-            cached_body_items,
+            item_count,
             exec_artifacts,
         )?;
         for artifact in artifact_updates {
@@ -765,7 +766,10 @@ pub(super) fn build_artifact_graph(
         merge_artifact_into_map(&mut map, exec_artifact.clone());
     }
 
-    Ok(ArtifactGraph { map })
+    Ok(ArtifactGraph {
+        map,
+        item_count: item_count + ast.body.len(),
+    })
 }
 
 /// These may have been created with placeholder `CodeRef`s because we didn't

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

@@ -6,25 +6,31 @@ use itertools::{EitherOrBoth, Itertools};
 use tokio::sync::RwLock;
 
 use crate::{
-    execution::{annotations, memory::Stack, state::ModuleInfoMap, EnvironmentRef, ExecState, ExecutorSettings},
+    execution::{
+        annotations,
+        memory::Stack,
+        state::{self as exec_state, ModuleInfoMap},
+        EnvironmentRef, ExecutorSettings,
+    },
     parsing::ast::types::{Annotation, Node, Program},
     walk::Node as WalkNode,
+    ExecOutcome, ExecutorContext,
 };
 
 lazy_static::lazy_static! {
     /// A static mutable lock for updating the last successful execution state for the cache.
-    static ref OLD_AST: Arc<RwLock<Option<OldAstState>>> = Default::default();
+    static ref OLD_AST: Arc<RwLock<Option<GlobalState>>> = Default::default();
     // The last successful run's memory. Not cleared after an unssuccessful run.
     static ref PREV_MEMORY: Arc<RwLock<Option<(Stack, ModuleInfoMap)>>> = Default::default();
 }
 
 /// Read the old ast memory from the lock.
-pub(crate) async fn read_old_ast() -> Option<OldAstState> {
+pub(super) async fn read_old_ast() -> Option<GlobalState> {
     let old_ast = OLD_AST.read().await;
     old_ast.clone()
 }
 
-pub(super) async fn write_old_ast(old_state: OldAstState) {
+pub(super) async fn write_old_ast(old_state: GlobalState) {
     let mut old_ast = OLD_AST.write().await;
     *old_ast = Some(old_state);
 }
@@ -34,7 +40,7 @@ pub(crate) async fn read_old_memory() -> Option<(Stack, ModuleInfoMap)> {
     old_mem.clone()
 }
 
-pub(super) async fn write_old_memory(mem: (Stack, ModuleInfoMap)) {
+pub(crate) async fn write_old_memory(mem: (Stack, ModuleInfoMap)) {
     let mut old_mem = PREV_MEMORY.write().await;
     *old_mem = Some(mem);
 }
@@ -56,16 +62,73 @@ pub struct CacheInformation<'a> {
     pub settings: &'a ExecutorSettings,
 }
 
-/// The old ast and program memory.
+/// The cached state of the whole program.
 #[derive(Debug, Clone)]
-pub struct OldAstState {
+pub(super) struct GlobalState {
-    /// The ast.
+    pub(super) main: ModuleState,
-    pub ast: Node<Program>,
     /// The exec state.
-    pub exec_state: ExecState,
+    pub(super) exec_state: exec_state::GlobalState,
     /// The last settings used for execution.
-    pub settings: crate::execution::ExecutorSettings,
+    pub(super) settings: ExecutorSettings,
-    pub result_env: EnvironmentRef,
+}
+
+impl GlobalState {
+    pub fn new(
+        state: exec_state::ExecState,
+        settings: ExecutorSettings,
+        ast: Node<Program>,
+        result_env: EnvironmentRef,
+    ) -> Self {
+        Self {
+            main: ModuleState {
+                ast,
+                exec_state: state.mod_local,
+                result_env,
+            },
+            exec_state: state.global,
+            settings,
+        }
+    }
+
+    pub fn with_settings(mut self, settings: ExecutorSettings) -> GlobalState {
+        self.settings = settings;
+        self
+    }
+
+    pub fn reconstitute_exec_state(&self) -> exec_state::ExecState {
+        exec_state::ExecState {
+            global: self.exec_state.clone(),
+            mod_local: self.main.exec_state.clone(),
+        }
+    }
+
+    pub async fn into_exec_outcome(self, ctx: &ExecutorContext) -> ExecOutcome {
+        // Fields are opt-in so that we don't accidentally leak private internal
+        // state when we add more to ExecState.
+        ExecOutcome {
+            variables: self.main.exec_state.variables(self.main.result_env),
+            filenames: self.exec_state.filenames(),
+            #[cfg(feature = "artifact-graph")]
+            operations: self.exec_state.artifacts.operations,
+            #[cfg(feature = "artifact-graph")]
+            artifact_commands: self.exec_state.artifacts.commands,
+            #[cfg(feature = "artifact-graph")]
+            artifact_graph: self.exec_state.artifacts.graph,
+            errors: self.exec_state.errors,
+            default_planes: ctx.engine.get_default_planes().read().await.clone(),
+        }
+    }
+}
+
+/// Per-module cached state
+#[derive(Debug, Clone)]
+pub(super) struct ModuleState {
+    /// The AST of the module.
+    pub(super) ast: Node<Program>,
+    /// The ExecState of the module.
+    pub(super) exec_state: exec_state::ModuleState,
+    /// The memory env for the module.
+    pub(super) result_env: EnvironmentRef,
 }
 
 /// The result of a cache check.
@@ -79,9 +142,6 @@ pub(super) enum CacheResult {
         reapply_settings: bool,
         /// The program that needs to be executed.
         program: Node<Program>,
-        /// The number of body items that were cached and omitted from the
-        /// program that needs to be executed. Used to compute [`crate::NodePath`].
-        cached_body_items: usize,
     },
     /// Check only the imports, and not the main program.
     /// Before sending this we already checked the main program and it is the same.
@@ -146,7 +206,6 @@ pub(super) async fn get_changed_program(old: CacheInformation<'_>, new: CacheInf
         // We know they have the same imports because the ast is the same.
         // If we have no imports, we can skip this.
         if !old.ast.has_import_statements() {
-            println!("No imports, no need to check.");
             return CacheResult::NoAction(reapply_settings);
         }
 
@@ -194,7 +253,6 @@ pub(super) async fn get_changed_program(old: CacheInformation<'_>, new: CacheInf
             clear_scene: true,
             reapply_settings: true,
             program: new.ast.clone(),
-            cached_body_items: 0,
         };
     }
 
@@ -223,7 +281,6 @@ fn generate_changed_program(old_ast: Node<Program>, mut new_ast: Node<Program>,
             clear_scene: true,
             reapply_settings,
             program: new_ast,
-            cached_body_items: 0,
         };
     }
 
@@ -244,7 +301,6 @@ fn generate_changed_program(old_ast: Node<Program>, mut new_ast: Node<Program>,
                 clear_scene: true,
                 reapply_settings,
                 program: new_ast,
-                cached_body_items: 0,
             }
         }
         std::cmp::Ordering::Greater => {
@@ -261,7 +317,6 @@ fn generate_changed_program(old_ast: Node<Program>, mut new_ast: Node<Program>,
                 clear_scene: false,
                 reapply_settings,
                 program: new_ast,
-                cached_body_items: old_ast.body.len(),
             }
         }
         std::cmp::Ordering::Equal => {
@@ -600,7 +655,6 @@ startSketchOn(XY)
                 clear_scene: true,
                 reapply_settings: true,
                 program: new_program.ast,
-                cached_body_items: 0,
             }
         );
     }
@@ -639,7 +693,6 @@ startSketchOn(XY)
                 clear_scene: true,
                 reapply_settings: true,
                 program: new_program.ast,
-                cached_body_items: 0,
             }
         );
     }

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

@@ -31,7 +31,7 @@ pub(crate) type Universe = HashMap<String, DependencyInfo>;
 /// run concurrently. Each "stage" is blocking in this model, which will
 /// change in the future. Don't use this function widely, yet.
 #[allow(clippy::iter_over_hash_type)]
-pub fn import_graph(progs: &Universe, ctx: &ExecutorContext) -> Result<Vec<Vec<String>>, KclError> {
+pub(crate) fn import_graph(progs: &Universe, ctx: &ExecutorContext) -> Result<Vec<Vec<String>>, KclError> {
     let mut graph = Graph::new();
 
     for (name, (_, _, path, repr)) in progs.iter() {
@@ -120,7 +120,7 @@ fn topsort(all_modules: &[&str], graph: Graph) -> Result<Vec<Vec<String>>, KclEr
 
 type ImportDependencies = Vec<(String, AstNode<ImportStatement>, ModulePath)>;
 
-pub(crate) fn import_dependencies(
+fn import_dependencies(
     path: &ModulePath,
     repr: &ModuleRepr,
     ctx: &ExecutorContext,

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

@@ -5,9 +5,8 @@ use std::sync::Arc;
 use anyhow::Result;
 #[cfg(feature = "artifact-graph")]
 pub use artifact::{Artifact, ArtifactCommand, ArtifactGraph, CodeRef, StartSketchOnFace, StartSketchOnPlane};
-use cache::OldAstState;
+use cache::GlobalState;
 pub use cache::{bust_cache, clear_mem_cache};
-#[cfg(feature = "artifact-graph")]
 pub use cad_op::{Group, Operation};
 pub use geometry::*;
 pub use id_generator::IdGenerator;
@@ -27,13 +26,12 @@ use serde::{Deserialize, Serialize};
 pub use state::{ExecState, MetaSettings};
 use uuid::Uuid;
 
-#[cfg(feature = "artifact-graph")]
-use crate::execution::artifact::build_artifact_graph;
 use crate::{
     engine::EngineManager,
     errors::{KclError, KclErrorDetails},
     execution::{
         cache::{CacheInformation, CacheResult},
+        import_graph::{Universe, UniverseMap},
         typed_path::TypedPath,
         types::{UnitAngle, UnitLen},
     },
@@ -41,7 +39,6 @@ use crate::{
     modules::{ModuleId, ModulePath, ModuleRepr},
     parsing::ast::types::{Expr, ImportPath, NodeRef},
     source_range::SourceRange,
-    walk::{Universe, UniverseMap},
     CompilationError, ExecError, KclErrorWithOutputs,
 };
 
@@ -55,6 +52,7 @@ pub mod fn_call;
 mod geometry;
 mod id_generator;
 mod import;
+mod import_graph;
 pub(crate) mod kcl_value;
 mod memory;
 mod state;
@@ -568,7 +566,7 @@ impl ExecutorContext {
         // part of the scene).
         exec_state.mut_stack().push_new_env_for_scope();
 
-        let result = self.inner_run(&program, 0, &mut exec_state, true).await?;
+        let result = self.inner_run(&program, &mut exec_state, true).await?;
 
         // Restore any temporary variables, then save any newly created variables back to
         // memory in case another run wants to use them. Note this is just saved to the preserved
@@ -576,7 +574,7 @@ impl ExecutorContext {
 
         let mut mem = exec_state.stack().clone();
         let module_infos = exec_state.global.module_infos.clone();
-        let outcome = exec_state.to_mock_exec_outcome(result.0).await;
+        let outcome = exec_state.to_mock_exec_outcome(result.0, self).await;
 
         mem.squash_env(result.0);
         cache::write_old_memory((mem, module_infos)).await;
@@ -587,169 +585,176 @@ impl ExecutorContext {
     pub async fn run_with_caching(&self, program: crate::Program) -> Result<ExecOutcome, KclErrorWithOutputs> {
         assert!(!self.is_mock());
 
-        let (program, mut exec_state, preserve_mem, cached_body_items, imports_info) = if let Some(OldAstState {
+        let (program, exec_state, result) = match cache::read_old_ast().await {
-            ast: old_ast,
+            Some(mut cached_state) => {
-            exec_state: mut old_state,
+                let old = CacheInformation {
-            settings: old_settings,
+                    ast: &cached_state.main.ast,
-            result_env,
+                    settings: &cached_state.settings,
-        }) =
+                };
-            cache::read_old_ast().await
+                let new = CacheInformation {
-        {
+                    ast: &program.ast,
-            let old = CacheInformation {
+                    settings: &self.settings,
-                ast: &old_ast,
-                settings: &old_settings,
-            };
-            let new = CacheInformation {
-                ast: &program.ast,
-                settings: &self.settings,
-            };
-
-            // Get the program that actually changed from the old and new information.
-            let (clear_scene, program, body_items, import_check_info) = match cache::get_changed_program(old, new).await
-            {
-                CacheResult::ReExecute {
-                    clear_scene,
-                    reapply_settings,
-                    program: changed_program,
-                    cached_body_items,
-                } => {
-                    if reapply_settings
-                        && self
-                            .engine
-                            .reapply_settings(&self.settings, Default::default(), old_state.id_generator())
-                            .await
-                            .is_err()
-                    {
-                        (true, program, cached_body_items, None)
-                    } else {
-                        (
-                            clear_scene,
-                            crate::Program {
-                                ast: changed_program,
-                                original_file_contents: program.original_file_contents,
-                            },
-                            cached_body_items,
-                            None,
-                        )
-                    }
-                }
-                CacheResult::CheckImportsOnly {
-                    reapply_settings,
-                    ast: changed_program,
-                } => {
-                    if reapply_settings
-                        && self
-                            .engine
-                            .reapply_settings(&self.settings, Default::default(), old_state.id_generator())
-                            .await
-                            .is_err()
-                    {
-                        (true, program, old_ast.body.len(), None)
-                    } else {
-                        // We need to check our imports to see if they changed.
-                        let mut new_exec_state = ExecState::new(self);
-                        let (new_universe, new_universe_map) = self.get_universe(&program, &mut new_exec_state).await?;
-                        let mut clear_scene = false;
-
-                        let mut keys = new_universe.keys().clone().collect::<Vec<_>>();
-                        keys.sort();
-                        for key in keys {
-                            let (_, id, _, _) = &new_universe[key];
-                            if let (Some(source0), Some(source1)) =
-                                (old_state.get_source(*id), new_exec_state.get_source(*id))
-                            {
-                                if source0.source != source1.source {
-                                    clear_scene = true;
-                                    break;
-                                }
-                            }
-                        }
-
-                        if !clear_scene {
-                            // Return early we don't need to clear the scene.
-                            let outcome = old_state.to_exec_outcome(result_env).await;
-                            return Ok(outcome);
-                        }
-
-                        (
-                            clear_scene,
-                            crate::Program {
-                                ast: changed_program,
-                                original_file_contents: program.original_file_contents,
-                            },
-                            old_ast.body.len(),
-                            // We only care about this if we are clearing the scene.
-                            if clear_scene {
-                                Some((new_universe, new_universe_map, new_exec_state))
-                            } else {
-                                None
-                            },
-                        )
-                    }
-                }
-                CacheResult::NoAction(true) => {
-                    if self
-                        .engine
-                        .reapply_settings(&self.settings, Default::default(), old_state.id_generator())
-                        .await
-                        .is_ok()
-                    {
-                        // We need to update the old ast state with the new settings!!
-                        cache::write_old_ast(OldAstState {
-                            ast: old_ast,
-                            exec_state: old_state.clone(),
-                            settings: self.settings.clone(),
-                            result_env,
-                        })
-                        .await;
-
-                        let outcome = old_state.to_exec_outcome(result_env).await;
-                        return Ok(outcome);
-                    }
-                    (true, program, old_ast.body.len(), None)
-                }
-                CacheResult::NoAction(false) => {
-                    let outcome = old_state.to_exec_outcome(result_env).await;
-                    return Ok(outcome);
-                }
-            };
-
-            let (exec_state, preserve_mem, universe_info) =
-                if let Some((new_universe, new_universe_map, mut new_exec_state)) = import_check_info {
-                    // Clear the scene if the imports changed.
-                    self.send_clear_scene(&mut new_exec_state, Default::default())
-                        .await
-                        .map_err(KclErrorWithOutputs::no_outputs)?;
-
-                    (new_exec_state, false, Some((new_universe, new_universe_map)))
-                } else if clear_scene {
-                    // Pop the execution state, since we are starting fresh.
-                    let mut exec_state = old_state;
-                    exec_state.reset(self);
-
-                    self.send_clear_scene(&mut exec_state, Default::default())
-                        .await
-                        .map_err(KclErrorWithOutputs::no_outputs)?;
-
-                    (exec_state, false, None)
-                } else {
-                    old_state.mut_stack().restore_env(result_env);
-
-                    (old_state, true, None)
                 };
 
-            (program, exec_state, preserve_mem, body_items, universe_info)
+                // Get the program that actually changed from the old and new information.
-        } else {
+                let (clear_scene, program, import_check_info) = match cache::get_changed_program(old, new).await {
-            let mut exec_state = ExecState::new(self);
+                    CacheResult::ReExecute {
-            self.send_clear_scene(&mut exec_state, Default::default())
+                        clear_scene,
-                .await
+                        reapply_settings,
-                .map_err(KclErrorWithOutputs::no_outputs)?;
+                        program: changed_program,
-            (program, exec_state, false, 0, None)
+                    } => {
-        };
+                        if reapply_settings
+                            && self
+                                .engine
+                                .reapply_settings(
+                                    &self.settings,
+                                    Default::default(),
+                                    &mut cached_state.main.exec_state.id_generator,
+                                )
+                                .await
+                                .is_err()
+                        {
+                            (true, program, None)
+                        } else {
+                            (
+                                clear_scene,
+                                crate::Program {
+                                    ast: changed_program,
+                                    original_file_contents: program.original_file_contents,
+                                },
+                                None,
+                            )
+                        }
+                    }
+                    CacheResult::CheckImportsOnly {
+                        reapply_settings,
+                        ast: changed_program,
+                    } => {
+                        if reapply_settings
+                            && self
+                                .engine
+                                .reapply_settings(
+                                    &self.settings,
+                                    Default::default(),
+                                    &mut cached_state.main.exec_state.id_generator,
+                                )
+                                .await
+                                .is_err()
+                        {
+                            (true, program, None)
+                        } else {
+                            // We need to check our imports to see if they changed.
+                            let mut new_exec_state = ExecState::new(self);
+                            let (new_universe, new_universe_map) =
+                                self.get_universe(&program, &mut new_exec_state).await?;
 
-        let result = self
+                            let clear_scene = new_universe.keys().any(|key| {
-            .run_concurrent(&program, cached_body_items, &mut exec_state, imports_info, preserve_mem)
+                                let id = new_universe[key].1;
-            .await;
+                                match (
+                                    cached_state.exec_state.get_source(id),
+                                    new_exec_state.global.get_source(id),
+                                ) {
+                                    (Some(s0), Some(s1)) => s0.source != s1.source,
+                                    _ => false,
+                                }
+                            });
+
+                            if !clear_scene {
+                                // Return early we don't need to clear the scene.
+                                return Ok(cached_state.into_exec_outcome(self).await);
+                            }
+
+                            (
+                                true,
+                                crate::Program {
+                                    ast: changed_program,
+                                    original_file_contents: program.original_file_contents,
+                                },
+                                Some((new_universe, new_universe_map, new_exec_state)),
+                            )
+                        }
+                    }
+                    CacheResult::NoAction(true) => {
+                        if self
+                            .engine
+                            .reapply_settings(
+                                &self.settings,
+                                Default::default(),
+                                &mut cached_state.main.exec_state.id_generator,
+                            )
+                            .await
+                            .is_ok()
+                        {
+                            // We need to update the old ast state with the new settings!!
+                            cache::write_old_ast(GlobalState::with_settings(
+                                cached_state.clone(),
+                                self.settings.clone(),
+                            ))
+                            .await;
+
+                            return Ok(cached_state.into_exec_outcome(self).await);
+                        }
+                        (true, program, None)
+                    }
+                    CacheResult::NoAction(false) => {
+                        return Ok(cached_state.into_exec_outcome(self).await);
+                    }
+                };
+
+                let (exec_state, result) = match import_check_info {
+                    Some((new_universe, new_universe_map, mut new_exec_state)) => {
+                        // Clear the scene if the imports changed.
+                        self.send_clear_scene(&mut new_exec_state, Default::default())
+                            .await
+                            .map_err(KclErrorWithOutputs::no_outputs)?;
+
+                        let result = self
+                            .run_concurrent(
+                                &program,
+                                &mut new_exec_state,
+                                Some((new_universe, new_universe_map)),
+                                false,
+                            )
+                            .await;
+
+                        (new_exec_state, result)
+                    }
+                    None if clear_scene => {
+                        // Pop the execution state, since we are starting fresh.
+                        let mut exec_state = cached_state.reconstitute_exec_state();
+                        exec_state.reset(self);
+
+                        self.send_clear_scene(&mut exec_state, Default::default())
+                            .await
+                            .map_err(KclErrorWithOutputs::no_outputs)?;
+
+                        let result = self.run_concurrent(&program, &mut exec_state, None, false).await;
+
+                        (exec_state, result)
+                    }
+                    None => {
+                        let mut exec_state = cached_state.reconstitute_exec_state();
+                        exec_state.mut_stack().restore_env(cached_state.main.result_env);
+
+                        let result = self.run_concurrent(&program, &mut exec_state, None, true).await;
+
+                        (exec_state, result)
+                    }
+                };
+
+                (program, exec_state, result)
+            }
+            None => {
+                let mut exec_state = ExecState::new(self);
+                self.send_clear_scene(&mut exec_state, Default::default())
+                    .await
+                    .map_err(KclErrorWithOutputs::no_outputs)?;
+
+                let result = self.run_concurrent(&program, &mut exec_state, None, false).await;
+
+                (program, exec_state, result)
+            }
+        };
 
         if result.is_err() {
             cache::bust_cache().await;
@@ -759,15 +764,15 @@ impl ExecutorContext {
         let result = result?;
 
         // Save this as the last successful execution to the cache.
-        cache::write_old_ast(OldAstState {
+        cache::write_old_ast(GlobalState::new(
-            ast: program.ast,
+            exec_state.clone(),
-            exec_state: exec_state.clone(),
+            self.settings.clone(),
-            settings: self.settings.clone(),
+            program.ast,
-            result_env: result.0,
+            result.0,
-        })
+        ))
         .await;
 
-        let outcome = exec_state.to_exec_outcome(result.0).await;
+        let outcome = exec_state.to_exec_outcome(result.0, self).await;
         Ok(outcome)
     }
 
@@ -782,11 +787,11 @@ impl ExecutorContext {
         program: &crate::Program,
         exec_state: &mut ExecState,
     ) -> Result<(EnvironmentRef, Option<ModelingSessionData>), KclErrorWithOutputs> {
-        self.run_concurrent(program, 0, exec_state, None, false).await
+        self.run_concurrent(program, exec_state, None, false).await
     }
 
     /// Perform the execution of a program using a concurrent
-    /// execution model. This has the same signature as [Self::run].
+    /// execution model.
     ///
     /// You can optionally pass in some initialization memory for partial
     /// execution.
@@ -795,13 +800,12 @@ impl ExecutorContext {
     pub async fn run_concurrent(
         &self,
         program: &crate::Program,
-        cached_body_items: usize,
         exec_state: &mut ExecState,
         universe_info: Option<(Universe, UniverseMap)>,
         preserve_mem: bool,
     ) -> Result<(EnvironmentRef, Option<ModelingSessionData>), KclErrorWithOutputs> {
         // Reuse our cached universe if we have one.
-        #[allow(unused_variables)]
+
         let (universe, universe_map) = if let Some((universe, universe_map)) = universe_info {
             (universe, universe_map)
         } else {
@@ -815,29 +819,8 @@ impl ExecutorContext {
             .await
             .map_err(KclErrorWithOutputs::no_outputs)?;
 
-        for modules in crate::walk::import_graph(&universe, self)
+        for modules in import_graph::import_graph(&universe, self)
-            .map_err(|err| {
+            .map_err(|err| exec_state.error_with_outputs(err, default_planes.clone()))?
-                let module_id_to_module_path: IndexMap<ModuleId, ModulePath> = exec_state
-                    .global
-                    .path_to_source_id
-                    .iter()
-                    .map(|(k, v)| ((*v), k.clone()))
-                    .collect();
-
-                KclErrorWithOutputs::new(
-                    err,
-                    exec_state.errors().to_vec(),
-                    #[cfg(feature = "artifact-graph")]
-                    exec_state.global.operations.clone(),
-                    #[cfg(feature = "artifact-graph")]
-                    exec_state.global.artifact_commands.clone(),
-                    #[cfg(feature = "artifact-graph")]
-                    exec_state.global.artifact_graph.clone(),
-                    module_id_to_module_path,
-                    exec_state.global.id_to_source.clone(),
-                    default_planes.clone(),
-                )
-            })?
             .into_iter()
         {
             #[cfg(not(target_arch = "wasm32"))]
@@ -859,7 +842,6 @@ impl ExecutorContext {
                 let module_path = module_path.clone();
                 let source_range = SourceRange::from(import_stmt);
 
-                #[cfg(feature = "artifact-graph")]
                 match &module_path {
                     ModulePath::Main => {
                         // This should never happen.
@@ -868,7 +850,7 @@ impl ExecutorContext {
                         // We only want to display the top-level module imports in
                         // the Feature Tree, not transitive imports.
                         if universe_map.contains_key(value) {
-                            exec_state.global.operations.push(Operation::GroupBegin {
+                            exec_state.push_op(Operation::GroupBegin {
                                 group: Group::ModuleInstance {
                                     name: value.file_name().unwrap_or_default(),
                                     module_id,
@@ -878,7 +860,7 @@ impl ExecutorContext {
                             // Due to concurrent execution, we cannot easily
                             // group operations by module. So we leave the
                             // group empty and close it immediately.
-                            exec_state.global.operations.push(Operation::GroupEnd);
+                            exec_state.push_op(Operation::GroupEnd);
                         }
                     }
                     ModulePath::Std { .. } => {
@@ -923,7 +905,6 @@ impl ExecutorContext {
                 #[cfg(target_arch = "wasm32")]
                 {
                     wasm_bindgen_futures::spawn_local(async move {
-                        //set.spawn(async move {
                         let mut exec_state = exec_state;
                         let exec_ctxt = exec_ctxt;
 
@@ -993,33 +974,13 @@ impl ExecutorContext {
                         exec_state.global.module_infos[&module_id].restore_repr(repr);
                     }
                     Err(e) => {
-                        let module_id_to_module_path: IndexMap<ModuleId, ModulePath> = exec_state
+                        return Err(exec_state.error_with_outputs(e, default_planes));
-                            .global
-                            .path_to_source_id
-                            .iter()
-                            .map(|(k, v)| ((*v), k.clone()))
-                            .collect();
-
-                        return Err(KclErrorWithOutputs::new(
-                            e,
-                            exec_state.errors().to_vec(),
-                            #[cfg(feature = "artifact-graph")]
-                            exec_state.global.operations.clone(),
-                            #[cfg(feature = "artifact-graph")]
-                            exec_state.global.artifact_commands.clone(),
-                            #[cfg(feature = "artifact-graph")]
-                            exec_state.global.artifact_graph.clone(),
-                            module_id_to_module_path,
-                            exec_state.global.id_to_source.clone(),
-                            default_planes,
-                        ));
                     }
                 }
             }
         }
 
-        self.inner_run(program, cached_body_items, exec_state, preserve_mem)
+        self.inner_run(program, exec_state, preserve_mem).await
-            .await
     }
 
     /// Get the universe & universe map of the program.
@@ -1035,7 +996,7 @@ impl ExecutorContext {
 
         let default_planes = self.engine.get_default_planes().read().await.clone();
 
-        let root_imports = crate::walk::import_universe(
+        let root_imports = import_graph::import_universe(
             self,
             &ModulePath::Main,
             &ModuleRepr::Kcl(program.ast.clone(), None),
@@ -1043,29 +1004,7 @@ impl ExecutorContext {
             exec_state,
         )
         .await
-        .map_err(|err| {
+        .map_err(|err| exec_state.error_with_outputs(err, default_planes))?;
-            println!("Error: {err:?}");
-            let module_id_to_module_path: IndexMap<ModuleId, ModulePath> = exec_state
-                .global
-                .path_to_source_id
-                .iter()
-                .map(|(k, v)| ((*v), k.clone()))
-                .collect();
-
-            KclErrorWithOutputs::new(
-                err,
-                exec_state.errors().to_vec(),
-                #[cfg(feature = "artifact-graph")]
-                exec_state.global.operations.clone(),
-                #[cfg(feature = "artifact-graph")]
-                exec_state.global.artifact_commands.clone(),
-                #[cfg(feature = "artifact-graph")]
-                exec_state.global.artifact_graph.clone(),
-                module_id_to_module_path,
-                exec_state.global.id_to_source.clone(),
-                default_planes,
-            )
-        })?;
 
         Ok((universe, root_imports))
     }
@@ -1075,7 +1014,6 @@ impl ExecutorContext {
     async fn inner_run(
         &self,
         program: &crate::Program,
-        cached_body_items: usize,
         exec_state: &mut ExecState,
         preserve_mem: bool,
     ) -> Result<(EnvironmentRef, Option<ModelingSessionData>), KclErrorWithOutputs> {
@@ -1089,7 +1027,7 @@ impl ExecutorContext {
 
         let default_planes = self.engine.get_default_planes().read().await.clone();
         let result = self
-            .execute_and_build_graph(&program.ast, cached_body_items, exec_state, preserve_mem)
+            .execute_and_build_graph(&program.ast, exec_state, preserve_mem)
             .await;
 
         crate::log::log(format!(
@@ -1098,28 +1036,7 @@ impl ExecutorContext {
         ));
         crate::log::log(format!("Engine stats: {:?}", self.engine.stats()));
 
-        let env_ref = result.map_err(|e| {
+        let env_ref = result.map_err(|e| exec_state.error_with_outputs(e, default_planes))?;
-            let module_id_to_module_path: IndexMap<ModuleId, ModulePath> = exec_state
-                .global
-                .path_to_source_id
-                .iter()
-                .map(|(k, v)| ((*v), k.clone()))
-                .collect();
-
-            KclErrorWithOutputs::new(
-                e,
-                exec_state.errors().to_vec(),
-                #[cfg(feature = "artifact-graph")]
-                exec_state.global.operations.clone(),
-                #[cfg(feature = "artifact-graph")]
-                exec_state.global.artifact_commands.clone(),
-                #[cfg(feature = "artifact-graph")]
-                exec_state.global.artifact_graph.clone(),
-                module_id_to_module_path,
-                exec_state.global.id_to_source.clone(),
-                default_planes.clone(),
-            )
-        })?;
 
         if !self.is_mock() {
             let mut mem = exec_state.stack().deep_clone();
@@ -1136,7 +1053,6 @@ impl ExecutorContext {
     async fn execute_and_build_graph(
         &self,
         program: NodeRef<'_, crate::parsing::ast::types::Program>,
-        #[cfg_attr(not(feature = "artifact-graph"), expect(unused))] cached_body_items: usize,
         exec_state: &mut ExecState,
         preserve_mem: bool,
     ) -> Result<EnvironmentRef, KclError> {
@@ -1163,40 +1079,9 @@ impl ExecutorContext {
         // and should be dropped.
         self.engine.clear_queues().await;
 
-        #[cfg(feature = "artifact-graph")]
+        match exec_state.build_artifact_graph(&self.engine, program).await {
-        {
+            Ok(_) => exec_result.map(|(_, env_ref, _)| env_ref),
-            let new_commands = self.engine.take_artifact_commands().await;
+            Err(err) => exec_result.and(Err(err)),
-            let new_responses = self.engine.take_responses().await;
-            let initial_graph = exec_state.global.artifact_graph.clone();
-
-            // Build the artifact graph.
-            let graph_result = build_artifact_graph(
-                &new_commands,
-                &new_responses,
-                program,
-                cached_body_items,
-                &mut exec_state.global.artifacts,
-                initial_graph,
-            );
-            // Move the artifact commands and responses into ExecState to
-            // simplify cache management and error creation.
-            exec_state.global.artifact_commands.extend(new_commands);
-            exec_state.global.artifact_responses.extend(new_responses);
-
-            match graph_result {
-                Ok(artifact_graph) => {
-                    exec_state.global.artifact_graph = artifact_graph;
-                    exec_result.map(|(_, env_ref, _)| env_ref)
-                }
-                Err(err) => {
-                    // Prefer the exec error.
-                    exec_result.and(Err(err))
-                }
-            }
-        }
-        #[cfg(not(feature = "artifact-graph"))]
-        {
-            exec_result.map(|(_, env_ref, _)| env_ref)
         }
     }
 
@@ -2198,7 +2083,7 @@ w = f() + f()
         }
 
         // Get the id_generator from the first execution.
-        let id_generator = cache::read_old_ast().await.unwrap().exec_state.mod_local.id_generator;
+        let id_generator = cache::read_old_ast().await.unwrap().main.exec_state.id_generator;
 
         let code = r#"sketch001 = startSketchOn(XZ)
 |> startProfile(at = [62.74, 206.13])
@@ -2219,7 +2104,7 @@ w = f() + f()
         // Execute the program.
         ctx.run_with_caching(program).await.unwrap();
 
-        let new_id_generator = cache::read_old_ast().await.unwrap().exec_state.mod_local.id_generator;
+        let new_id_generator = cache::read_old_ast().await.unwrap().main.exec_state.id_generator;
 
         assert_eq!(id_generator, new_id_generator);
     }

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

@@ -12,19 +12,19 @@ use uuid::Uuid;
 use crate::execution::{Artifact, ArtifactCommand, ArtifactGraph, ArtifactId};
 use crate::{
     errors::{KclError, KclErrorDetails, Severity},
+    exec::DefaultPlanes,
     execution::{
         annotations,
         cad_op::Operation,
         id_generator::IdGenerator,
         memory::{ProgramMemory, Stack},
-        types,
+        types::{self, NumericType},
-        types::NumericType,
         EnvironmentRef, ExecOutcome, ExecutorSettings, KclValue, UnitAngle, UnitLen,
     },
     modules::{ModuleId, ModuleInfo, ModuleLoader, ModulePath, ModuleRepr, ModuleSource},
-    parsing::ast::types::Annotation,
+    parsing::ast::types::{Annotation, NodeRef},
     source_range::SourceRange,
-    CompilationError,
+    CompilationError, EngineManager, ExecutorContext, KclErrorWithOutputs,
 };
 
 /// State for executing a program.
@@ -32,7 +32,6 @@ use crate::{
 pub struct ExecState {
     pub(super) global: GlobalState,
     pub(super) mod_local: ModuleState,
-    pub(super) exec_context: Option<super::ExecutorContext>,
 }
 
 pub type ModuleInfoMap = IndexMap<ModuleId, ModuleInfo>;
@@ -45,33 +44,39 @@ pub(super) struct GlobalState {
     pub id_to_source: IndexMap<ModuleId, ModuleSource>,
     /// Map from module ID to module info.
     pub module_infos: ModuleInfoMap,
-    /// Output map of UUIDs to artifacts.
-    #[cfg(feature = "artifact-graph")]
-    pub artifacts: IndexMap<ArtifactId, Artifact>,
-    /// Output commands to allow building the artifact graph by the caller.
-    /// These are accumulated in the [`ExecutorContext`] but moved here for
-    /// convenience of the execution cache.
-    #[cfg(feature = "artifact-graph")]
-    pub artifact_commands: Vec<ArtifactCommand>,
-    /// Responses from the engine for `artifact_commands`.  We need to cache
-    /// this so that we can build the artifact graph.  These are accumulated in
-    /// the [`ExecutorContext`] but moved here for convenience of the execution
-    /// cache.
-    #[cfg(feature = "artifact-graph")]
-    pub artifact_responses: IndexMap<Uuid, WebSocketResponse>,
-    /// Output artifact graph.
-    #[cfg(feature = "artifact-graph")]
-    pub artifact_graph: ArtifactGraph,
-    /// Operations that have been performed in execution order, for display in
-    /// the Feature Tree.
-    #[cfg(feature = "artifact-graph")]
-    pub operations: Vec<Operation>,
     /// Module loader.
     pub mod_loader: ModuleLoader,
     /// Errors and warnings.
     pub errors: Vec<CompilationError>,
+    #[allow(dead_code)]
+    pub artifacts: ArtifactState,
 }
 
+#[cfg(feature = "artifact-graph")]
+#[derive(Debug, Clone, Default)]
+pub(super) struct ArtifactState {
+    /// Output map of UUIDs to artifacts.
+    pub artifacts: IndexMap<ArtifactId, Artifact>,
+    /// Output commands to allow building the artifact graph by the caller.
+    /// These are accumulated in the [`ExecutorContext`] but moved here for
+    /// convenience of the execution cache.
+    pub commands: Vec<ArtifactCommand>,
+    /// Responses from the engine for `artifact_commands`.  We need to cache
+    /// this so that we can build the artifact graph.  These are accumulated in
+    /// the [`ExecutorContext`] but moved here for convenience of the execution
+    /// cache.
+    pub responses: IndexMap<Uuid, WebSocketResponse>,
+    /// Output artifact graph.
+    pub graph: ArtifactGraph,
+    /// Operations that have been performed in execution order, for display in
+    /// the Feature Tree.
+    pub operations: Vec<Operation>,
+}
+
+#[cfg(not(feature = "artifact-graph"))]
+#[derive(Debug, Clone, Default)]
+pub(super) struct ArtifactState {}
+
 #[derive(Debug, Clone)]
 pub(super) struct ModuleState {
     /// The id generator for this module.
@@ -98,7 +103,6 @@ impl ExecState {
         ExecState {
             global: GlobalState::new(&exec_context.settings),
             mod_local: ModuleState::new(ModulePath::Main, ProgramMemory::new(), Default::default()),
-            exec_context: Some(exec_context.clone()),
         }
     }
 
@@ -108,7 +112,6 @@ impl ExecState {
         *self = ExecState {
             global,
             mod_local: ModuleState::new(self.mod_local.path.clone(), ProgramMemory::new(), Default::default()),
-            exec_context: Some(exec_context.clone()),
         };
     }
 
@@ -130,45 +133,26 @@ impl ExecState {
     /// Convert to execution outcome when running in WebAssembly.  We want to
     /// reduce the amount of data that crosses the WASM boundary as much as
     /// possible.
-    pub async fn to_exec_outcome(self, main_ref: EnvironmentRef) -> ExecOutcome {
+    pub async fn to_exec_outcome(self, main_ref: EnvironmentRef, ctx: &ExecutorContext) -> ExecOutcome {
         // Fields are opt-in so that we don't accidentally leak private internal
         // state when we add more to ExecState.
         ExecOutcome {
-            variables: self
+            variables: self.mod_local.variables(main_ref),
-                .stack()
+            filenames: self.global.filenames(),
-                .find_all_in_env(main_ref)
-                .map(|(k, v)| (k.clone(), v.clone()))
-                .collect(),
             #[cfg(feature = "artifact-graph")]
-            operations: self.global.operations,
+            operations: self.global.artifacts.operations,
             #[cfg(feature = "artifact-graph")]
-            artifact_commands: self.global.artifact_commands,
+            artifact_commands: self.global.artifacts.commands,
             #[cfg(feature = "artifact-graph")]
-            artifact_graph: self.global.artifact_graph,
+            artifact_graph: self.global.artifacts.graph,
             errors: self.global.errors,
-            filenames: self
+            default_planes: ctx.engine.get_default_planes().read().await.clone(),
-                .global
-                .path_to_source_id
-                .iter()
-                .map(|(k, v)| ((*v), k.clone()))
-                .collect(),
-            default_planes: if let Some(ctx) = &self.exec_context {
-                ctx.engine.get_default_planes().read().await.clone()
-            } else {
-                None
-            },
         }
     }
 
-    pub async fn to_mock_exec_outcome(self, main_ref: EnvironmentRef) -> ExecOutcome {
+    pub async fn to_mock_exec_outcome(self, main_ref: EnvironmentRef, ctx: &ExecutorContext) -> ExecOutcome {
-        // Fields are opt-in so that we don't accidentally leak private internal
-        // state when we add more to ExecState.
         ExecOutcome {
-            variables: self
+            variables: self.mod_local.variables(main_ref),
-                .stack()
-                .find_all_in_env(main_ref)
-                .map(|(k, v)| (k.clone(), v.clone()))
-                .collect(),
             #[cfg(feature = "artifact-graph")]
             operations: Default::default(),
             #[cfg(feature = "artifact-graph")]
@@ -177,11 +161,7 @@ impl ExecState {
             artifact_graph: Default::default(),
             errors: self.global.errors,
             filenames: Default::default(),
-            default_planes: if let Some(ctx) = &self.exec_context {
+            default_planes: ctx.engine.get_default_planes().read().await.clone(),
-                ctx.engine.get_default_planes().read().await.clone()
-            } else {
-                None
-            },
         }
     }
 
@@ -204,12 +184,12 @@ impl ExecState {
     #[cfg(feature = "artifact-graph")]
     pub(crate) fn add_artifact(&mut self, artifact: Artifact) {
         let id = artifact.id();
-        self.global.artifacts.insert(id, artifact);
+        self.global.artifacts.artifacts.insert(id, artifact);
     }
 
     pub(crate) fn push_op(&mut self, op: Operation) {
         #[cfg(feature = "artifact-graph")]
-        self.global.operations.push(op);
+        self.global.artifacts.operations.push(op);
         #[cfg(not(feature = "artifact-graph"))]
         drop(op);
     }
@@ -251,10 +231,6 @@ impl ExecState {
         self.global.id_to_source.insert(id, source.clone());
     }
 
-    pub(super) fn get_source(&self, id: ModuleId) -> Option<&ModuleSource> {
-        self.global.id_to_source.get(&id)
-    }
-
     pub(super) fn add_module(&mut self, id: ModuleId, path: ModulePath, repr: ModuleRepr) {
         debug_assert!(self.global.path_to_source_id.contains_key(&path));
         let module_info = ModuleInfo { id, repr, path };
@@ -300,6 +276,71 @@ impl ExecState {
     pub(crate) fn pipe_value(&self) -> Option<&KclValue> {
         self.mod_local.pipe_value.as_ref()
     }
+
+    pub(crate) fn error_with_outputs(
+        &self,
+        error: KclError,
+        default_planes: Option<DefaultPlanes>,
+    ) -> KclErrorWithOutputs {
+        let module_id_to_module_path: IndexMap<ModuleId, ModulePath> = self
+            .global
+            .path_to_source_id
+            .iter()
+            .map(|(k, v)| ((*v), k.clone()))
+            .collect();
+
+        KclErrorWithOutputs::new(
+            error,
+            self.errors().to_vec(),
+            #[cfg(feature = "artifact-graph")]
+            self.global.artifacts.operations.clone(),
+            #[cfg(feature = "artifact-graph")]
+            self.global.artifacts.commands.clone(),
+            #[cfg(feature = "artifact-graph")]
+            self.global.artifacts.graph.clone(),
+            module_id_to_module_path,
+            self.global.id_to_source.clone(),
+            default_planes,
+        )
+    }
+
+    #[cfg(feature = "artifact-graph")]
+    pub(crate) async fn build_artifact_graph(
+        &mut self,
+        engine: &Arc<Box<dyn EngineManager>>,
+        program: NodeRef<'_, crate::parsing::ast::types::Program>,
+    ) -> Result<(), KclError> {
+        let new_commands = engine.take_artifact_commands().await;
+        let new_responses = engine.take_responses().await;
+        let initial_graph = self.global.artifacts.graph.clone();
+
+        // Build the artifact graph.
+        let graph_result = crate::execution::artifact::build_artifact_graph(
+            &new_commands,
+            &new_responses,
+            program,
+            &mut self.global.artifacts.artifacts,
+            initial_graph,
+        );
+        // Move the artifact commands and responses into ExecState to
+        // simplify cache management and error creation.
+        self.global.artifacts.commands.extend(new_commands);
+        self.global.artifacts.responses.extend(new_responses);
+
+        let artifact_graph = graph_result?;
+        self.global.artifacts.graph = artifact_graph;
+
+        Ok(())
+    }
+
+    #[cfg(not(feature = "artifact-graph"))]
+    pub(crate) async fn build_artifact_graph(
+        &mut self,
+        _engine: &Arc<Box<dyn EngineManager>>,
+        _program: NodeRef<'_, crate::parsing::ast::types::Program>,
+    ) -> Result<(), KclError> {
+        Ok(())
+    }
 }
 
 impl GlobalState {
@@ -307,16 +348,7 @@ impl GlobalState {
         let mut global = GlobalState {
             path_to_source_id: Default::default(),
             module_infos: Default::default(),
-            #[cfg(feature = "artifact-graph")]
             artifacts: Default::default(),
-            #[cfg(feature = "artifact-graph")]
-            artifact_commands: Default::default(),
-            #[cfg(feature = "artifact-graph")]
-            artifact_responses: Default::default(),
-            #[cfg(feature = "artifact-graph")]
-            artifact_graph: Default::default(),
-            #[cfg(feature = "artifact-graph")]
-            operations: Default::default(),
             mod_loader: Default::default(),
             errors: Default::default(),
             id_to_source: Default::default(),
@@ -339,6 +371,14 @@ impl GlobalState {
             .insert(ModulePath::Local { value: root_path }, root_id);
         global
     }
+
+    pub(super) fn filenames(&self) -> IndexMap<ModuleId, ModulePath> {
+        self.path_to_source_id.iter().map(|(k, v)| ((*v), k.clone())).collect()
+    }
+
+    pub(super) fn get_source(&self, id: ModuleId) -> Option<&ModuleSource> {
+        self.id_to_source.get(&id)
+    }
 }
 
 impl ModuleState {
@@ -358,6 +398,13 @@ impl ModuleState {
             },
         }
     }
+
+    pub(super) fn variables(&self, main_ref: EnvironmentRef) -> IndexMap<String, KclValue> {
+        self.stack
+            .find_all_in_env(main_ref)
+            .map(|(k, v)| (k.clone(), v.clone()))
+            .collect()
+    }
 }
 
 #[derive(Debug, Default, Clone, Deserialize, Serialize, PartialEq, Eq, ts_rs::TS, JsonSchema)]

rust/kcl-lib/src/simulation_tests.rs

@@ -163,7 +163,7 @@ async fn execute_test(test: &Test, render_to_png: bool, export_step: bool) {
     // Run the program.
     let exec_res = crate::test_server::execute_and_snapshot_ast(ast, Some(test.entry_point.clone()), export_step).await;
     match exec_res {
-        Ok((exec_state, env_ref, png, step)) => {
+        Ok((exec_state, ctx, env_ref, png, step)) => {
             let fail_path = test.output_dir.join("execution_error.snap");
             if std::fs::exists(&fail_path).unwrap() {
                 panic!("This test case is expected to fail, but it passed. If this is intended, and the test should actually be passing now, please delete kcl-lib/{}", fail_path.to_string_lossy())
@@ -181,7 +181,7 @@ async fn execute_test(test: &Test, render_to_png: bool, export_step: bool) {
                     panic!("Step data was empty");
                 }
             }
-            let outcome = exec_state.to_exec_outcome(env_ref).await;
+            let outcome = exec_state.to_exec_outcome(env_ref, &ctx).await;
 
             let mem_result = catch_unwind(AssertUnwindSafe(|| {
                 assert_snapshot(test, "Variables in memory after executing", || {

rust/kcl-lib/src/test_server.rs

@@ -36,7 +36,16 @@ pub async fn execute_and_snapshot_ast(
     ast: Program,
     current_file: Option<PathBuf>,
     with_export_step: bool,
-) -> Result<(ExecState, EnvironmentRef, image::DynamicImage, Option<Vec<u8>>), ExecErrorWithState> {
+) -> Result<
+    (
+        ExecState,
+        ExecutorContext,
+        EnvironmentRef,
+        image::DynamicImage,
+        Option<Vec<u8>>,
+    ),
+    ExecErrorWithState,
+> {
     let ctx = new_context(true, current_file).await?;
     let (exec_state, env, img) = match do_execute_and_snapshot(&ctx, ast).await {
         Ok((exec_state, env_ref, img)) => (exec_state, env_ref, img),
@@ -64,7 +73,7 @@ pub async fn execute_and_snapshot_ast(
         step = files.into_iter().next().map(|f| f.contents);
     }
     ctx.close().await;
-    Ok((exec_state, env, img, step))
+    Ok((exec_state, ctx, env, img, step))
 }
 
 pub async fn execute_and_snapshot_no_auth(

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

@@ -3,10 +3,7 @@
 mod ast_node;
 mod ast_visitor;
 mod ast_walk;
-mod import_graph;
 
-pub use ast_node::Node;
+pub(crate) use ast_node::Node;
-pub use ast_visitor::Visitable;
+pub(crate) use ast_visitor::Visitable;
-pub use ast_walk::walk;
+pub(crate) use ast_walk::walk;
-pub use import_graph::import_graph;
-pub(crate) use import_graph::{import_universe, Universe, UniverseMap};