Fix to account for cached body items when computing NodePath (#7030)

rust/kcl-lib/e2e/executor/cache.rs

@@ -1,5 +1,7 @@
 //! Cache testing framework.
 
+#[cfg(feature = "artifact-graph")]
+use kcl_lib::NodePathStep;
 use kcl_lib::{bust_cache, ExecError, ExecOutcome};
 use kcmc::{each_cmd as mcmd, ModelingCmd};
 use kittycad_modeling_cmds as kcmc;
@@ -330,6 +332,40 @@ extrude001 = extrude(profile001, length = 4)
     }
 }
 
+#[cfg(feature = "artifact-graph")]
+#[tokio::test(flavor = "multi_thread")]
+async fn kcl_test_cache_add_offset_plane_computes_node_path() {
+    let code = r#"sketch001 = startSketchOn(XY)
+profile001 = startProfile(sketch001, at = [0, 0])
+"#;
+    let code_with_more = code.to_owned()
+        + r#"plane001 = offsetPlane(XY, offset = 500)
+"#;
+
+    let result = cache_test(
+        "add_offset_plane_preserves_artifact_commands",
+        vec![
+            Variation {
+                code,
+                other_files: vec![],
+                settings: &Default::default(),
+            },
+            Variation {
+                code: code_with_more.as_str(),
+                other_files: vec![],
+                settings: &Default::default(),
+            },
+        ],
+    )
+    .await;
+
+    let second = &result.last().unwrap().2;
+
+    let v = second.artifact_graph.values().collect::<Vec<_>>();
+    let path_step = &v[2].code_ref().unwrap().node_path.steps[0];
+    assert_eq!(*path_step, NodePathStep::ProgramBodyItem { index: 2 });
+}
+
 #[tokio::test(flavor = "multi_thread")]
 async fn kcl_test_cache_empty_file_pop_cache_empty_file_planes_work() {
     // Get the current working directory.

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

@@ -750,6 +750,7 @@ 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> {
@@ -763,7 +764,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);
+        fill_in_node_paths(exec_artifact, ast, cached_body_items);
     }
 
     for artifact_command in artifact_commands {
@@ -790,6 +791,7 @@ pub(super) fn build_artifact_graph(
             &flattened_responses,
             &path_to_plane_id_map,
             ast,
+            cached_body_items,
             exec_artifacts,
         )?;
         for artifact in artifact_updates {
@@ -807,16 +809,18 @@ pub(super) fn build_artifact_graph(
 
 /// These may have been created with placeholder `CodeRef`s because we didn't
 /// have the entire AST available. Now we fill them in.
-fn fill_in_node_paths(artifact: &mut Artifact, program: &Node<Program>) {
+fn fill_in_node_paths(artifact: &mut Artifact, program: &Node<Program>, cached_body_items: usize) {
     match artifact {
         Artifact::StartSketchOnFace(face) => {
             if face.code_ref.node_path.is_empty() {
-                face.code_ref.node_path = NodePath::from_range(program, face.code_ref.range).unwrap_or_default();
+                face.code_ref.node_path =
+                    NodePath::from_range(program, cached_body_items, face.code_ref.range).unwrap_or_default();
             }
         }
         Artifact::StartSketchOnPlane(plane) => {
             if plane.code_ref.node_path.is_empty() {
-                plane.code_ref.node_path = NodePath::from_range(program, plane.code_ref.range).unwrap_or_default();
+                plane.code_ref.node_path =
+                    NodePath::from_range(program, cached_body_items, plane.code_ref.range).unwrap_or_default();
             }
         }
         _ => {}
@@ -905,6 +909,7 @@ fn artifacts_to_update(
     responses: &FnvHashMap<Uuid, OkModelingCmdResponse>,
     path_to_plane_id_map: &FnvHashMap<Uuid, Uuid>,
     ast: &Node<Program>,
+    cached_body_items: usize,
     exec_artifacts: &IndexMap<ArtifactId, Artifact>,
 ) -> Result<Vec<Artifact>, KclError> {
     let uuid = artifact_command.cmd_id;
@@ -918,7 +923,7 @@ fn artifacts_to_update(
     // correct value based on NodePath.
     let path_to_node = Vec::new();
     let range = artifact_command.range;
-    let node_path = NodePath::from_range(ast, range).unwrap_or_default();
+    let node_path = NodePath::from_range(ast, cached_body_items, range).unwrap_or_default();
     let code_ref = CodeRef {
         range,
         node_path,

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

@@ -79,6 +79,9 @@ 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.
@@ -191,6 +194,7 @@ 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,
         };
     }
 
@@ -219,6 +223,7 @@ 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,
         };
     }
 
@@ -239,6 +244,7 @@ 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 => {
@@ -255,6 +261,7 @@ 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 => {
@@ -592,7 +599,8 @@ startSketchOn(XY)
             CacheResult::ReExecute {
                 clear_scene: true,
                 reapply_settings: true,
-                program: new_program.ast
+                program: new_program.ast,
+                cached_body_items: 0,
             }
         );
     }
@@ -630,7 +638,8 @@ startSketchOn(XY)
             CacheResult::ReExecute {
                 clear_scene: true,
                 reapply_settings: true,
-                program: new_program.ast
+                program: new_program.ast,
+                cached_body_items: 0,
             }
         );
     }

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

@@ -571,7 +571,7 @@ impl ExecutorContext {
         // part of the scene).
         exec_state.mut_stack().push_new_env_for_scope();
 
-        let result = self.inner_run(&program, &mut exec_state, true).await?;
+        let result = self.inner_run(&program, 0, &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
@@ -590,12 +590,13 @@ 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, imports_info) = if let Some(OldAstState {
+        let (program, mut exec_state, preserve_mem, cached_body_items, imports_info) = if let Some(OldAstState {
             ast: old_ast,
             exec_state: mut old_state,
             settings: old_settings,
             result_env,
-        }) = cache::read_old_ast().await
+        }) =
+            cache::read_old_ast().await
         {
             let old = CacheInformation {
                 ast: &old_ast,
@@ -607,11 +608,13 @@ impl ExecutorContext {
             };
 
             // Get the program that actually changed from the old and new information.
-            let (clear_scene, program, import_check_info) = match cache::get_changed_program(old, new).await {
+            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
@@ -620,7 +623,7 @@ impl ExecutorContext {
                             .await
                             .is_err()
                     {
-                        (true, program, None)
+                        (true, program, cached_body_items, None)
                     } else {
                         (
                             clear_scene,
@@ -628,6 +631,7 @@ impl ExecutorContext {
                                 ast: changed_program,
                                 original_file_contents: program.original_file_contents,
                             },
+                            cached_body_items,
                             None,
                         )
                     }
@@ -643,7 +647,7 @@ impl ExecutorContext {
                             .await
                             .is_err()
                     {
-                        (true, program, None)
+                        (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);
@@ -676,6 +680,7 @@ impl ExecutorContext {
                                 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))
@@ -704,7 +709,7 @@ impl ExecutorContext {
                         let outcome = old_state.to_exec_outcome(result_env).await;
                         return Ok(outcome);
                     }
-                    (true, program, None)
+                    (true, program, old_ast.body.len(), None)
                 }
                 CacheResult::NoAction(false) => {
                     let outcome = old_state.to_exec_outcome(result_env).await;
@@ -736,17 +741,17 @@ impl ExecutorContext {
                     (old_state, true, None)
                 };
 
-            (program, exec_state, preserve_mem, universe_info)
+            (program, exec_state, preserve_mem, body_items, universe_info)
         } else {
             let mut exec_state = ExecState::new(self);
             self.send_clear_scene(&mut exec_state, Default::default())
                 .await
                 .map_err(KclErrorWithOutputs::no_outputs)?;
-            (program, exec_state, false, None)
+            (program, exec_state, false, 0, None)
         };
 
         let result = self
-            .run_concurrent(&program, &mut exec_state, imports_info, preserve_mem)
+            .run_concurrent(&program, cached_body_items, &mut exec_state, imports_info, preserve_mem)
             .await;
 
         if result.is_err() {
@@ -780,7 +785,7 @@ impl ExecutorContext {
         program: &crate::Program,
         exec_state: &mut ExecState,
     ) -> Result<(EnvironmentRef, Option<ModelingSessionData>), KclErrorWithOutputs> {
-        self.run_concurrent(program, exec_state, None, false).await
+        self.run_concurrent(program, 0, exec_state, None, false).await
     }
 
     /// Perform the execution of a program using a concurrent
@@ -793,6 +798,7 @@ 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,
@@ -1016,7 +1022,8 @@ impl ExecutorContext {
             }
         }
 
-        self.inner_run(program, exec_state, preserve_mem).await
+        self.inner_run(program, cached_body_items, exec_state, preserve_mem)
+            .await
     }
 
     /// Get the universe & universe map of the program.
@@ -1071,6 +1078,7 @@ 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> {
@@ -1084,7 +1092,7 @@ impl ExecutorContext {
 
         let default_planes = self.engine.get_default_planes().read().await.clone();
         let result = self
-            .execute_and_build_graph(&program.ast, exec_state, preserve_mem)
+            .execute_and_build_graph(&program.ast, cached_body_items, exec_state, preserve_mem)
             .await;
 
         crate::log::log(format!(
@@ -1131,6 +1139,7 @@ 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> {
@@ -1168,6 +1177,7 @@ impl ExecutorContext {
                 &new_commands,
                 &new_responses,
                 program,
+                cached_body_items,
                 &mut exec_state.global.artifacts,
                 initial_graph,
             );

rust/kcl-lib/src/lib.rs

@@ -99,7 +99,7 @@ pub use lsp::{
     kcl::{Backend as KclLspBackend, Server as KclLspServerSubCommand},
 };
 pub use modules::ModuleId;
-pub use parsing::ast::types::{FormatOptions, NodePath};
+pub use parsing::ast::types::{FormatOptions, NodePath, Step as NodePathStep};
 pub use settings::types::{project::ProjectConfiguration, Configuration, UnitLength};
 pub use source_range::SourceRange;
 #[cfg(not(target_arch = "wasm32"))]
@@ -250,8 +250,8 @@ impl Program {
         self.ast.lint(rule)
     }
 
-    pub fn node_path_from_range(&self, range: SourceRange) -> Option<NodePath> {
+    pub fn node_path_from_range(&self, cached_body_items: usize, range: SourceRange) -> Option<NodePath> {
-        NodePath::from_range(&self.ast, range)
+        NodePath::from_range(&self.ast, cached_body_items, range)
     }
 
     pub fn recast(&self) -> String {

rust/kcl-lib/src/parsing/ast/types/mod.rs

@@ -11,7 +11,7 @@ use std::{
 
 use anyhow::Result;
 use parse_display::{Display, FromStr};
-pub use path::NodePath;
+pub use path::{NodePath, Step};
 use schemars::JsonSchema;
 use serde::{Deserialize, Serialize};
 use tower_lsp::lsp_types::{

rust/kcl-lib/src/parsing/ast/types/path.rs

@@ -62,14 +62,21 @@ pub enum Step {
 impl NodePath {
     /// Given a program and a [`SourceRange`], return the path to the node that
     /// contains the range.
-    pub(crate) fn from_range(program: &Node<Program>, range: SourceRange) -> Option<Self> {
+    pub(crate) fn from_range(program: &Node<Program>, cached_body_items: usize, range: SourceRange) -> Option<Self> {
-        Self::from_body(&program.body, range, NodePath::default())
+        Self::from_body(&program.body, cached_body_items, range, NodePath::default())
     }
 
-    fn from_body(body: &[BodyItem], range: SourceRange, mut path: NodePath) -> Option<NodePath> {
+    fn from_body(
+        body: &[BodyItem],
+        cached_body_items: usize,
+        range: SourceRange,
+        mut path: NodePath,
+    ) -> Option<NodePath> {
         for (i, item) in body.iter().enumerate() {
             if item.contains_range(&range) {
-                path.push(Step::ProgramBodyItem { index: i });
+                path.push(Step::ProgramBodyItem {
+                    index: cached_body_items + i,
+                });
                 return Self::from_body_item(item, range, path);
             }
         }
@@ -262,7 +269,7 @@ impl NodePath {
                 }
                 if node.then_val.contains_range(&range) {
                     path.push(Step::IfExpressionThen);
-                    return Self::from_body(&node.then_val.body, range, path);
+                    return Self::from_body(&node.then_val.body, 0, range, path);
                 }
                 for else_if in &node.else_ifs {
                     if else_if.contains_range(&range) {
@@ -273,14 +280,14 @@ impl NodePath {
                         }
                         if else_if.then_val.contains_range(&range) {
                             path.push(Step::IfExpressionElseIfBody);
-                            return Self::from_body(&else_if.then_val.body, range, path);
+                            return Self::from_body(&else_if.then_val.body, 0, range, path);
                         }
                         return Some(path);
                     }
                 }
                 if node.final_else.contains_range(&range) {
                     path.push(Step::IfExpressionElse);
-                    return Self::from_body(&node.final_else.body, range, path);
+                    return Self::from_body(&node.final_else.body, 0, range, path);
                 }
             }
             Expr::LabelledExpression(node) => {
@@ -345,7 +352,7 @@ mod tests {
         // fn cube(sideLength, center) {
         //    ^^^^
         assert_eq!(
-            NodePath::from_range(&program.ast, range(38, 42)).unwrap(),
+            NodePath::from_range(&program.ast, 0, range(38, 42)).unwrap(),
             NodePath {
                 steps: vec![Step::ProgramBodyItem { index: 0 }, Step::VariableDeclarationDeclaration],
             }
@@ -353,7 +360,7 @@ mod tests {
         // fn cube(sideLength, center) {
         //                     ^^^^^^
         assert_eq!(
-            NodePath::from_range(&program.ast, range(55, 61)).unwrap(),
+            NodePath::from_range(&program.ast, 0, range(55, 61)).unwrap(),
             NodePath {
                 steps: vec![
                     Step::ProgramBodyItem { index: 0 },
@@ -366,7 +373,7 @@ mod tests {
         // |> line(endAbsolute = p1)
         //                       ^^
         assert_eq!(
-            NodePath::from_range(&program.ast, range(293, 295)).unwrap(),
+            NodePath::from_range(&program.ast, 0, range(293, 295)).unwrap(),
             NodePath {
                 steps: vec![
                     Step::ProgramBodyItem { index: 0 },
@@ -383,7 +390,7 @@ mod tests {
         // myCube = cube(sideLength = 40, center = [0, 0])
         //                                             ^
         assert_eq!(
-            NodePath::from_range(&program.ast, range(485, 486)).unwrap(),
+            NodePath::from_range(&program.ast, 0, range(485, 486)).unwrap(),
             NodePath {
                 steps: vec![
                     Step::ProgramBodyItem { index: 1 },

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

@@ -28,7 +28,7 @@ pub async fn node_path_from_range(program_ast_json: &str, range_json: &str) -> R
 
     let program: Program = serde_json::from_str(program_ast_json).map_err(|e| e.to_string())?;
     let range: SourceRange = serde_json::from_str(range_json).map_err(|e| e.to_string())?;
-    let node_path = program.node_path_from_range(range);
+    let node_path = program.node_path_from_range(0, range);
 
     JsValue::from_serde(&node_path).map_err(|e| e.to_string())
 }