Fix a units bug in multi-module projects (#5529)

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

src/lang/errors.ts

@@ -351,16 +351,16 @@ export function kclErrorsByFilename(
     const sourceRange: SourceRange = error.sourceRange
     const fileIndex = sourceRange[2]
     const modulePath: ModulePath | undefined = filenames[fileIndex]
-    if (modulePath) {
+    if (modulePath && modulePath.type === 'Local') {
-      let stdOrLocalPath = modulePath.value
+      let localPath = modulePath.value
-      if (stdOrLocalPath) {
+      if (localPath) {
         // Build up an array of errors per file name
-        const value = fileNameToError.get(stdOrLocalPath)
+        const value = fileNameToError.get(localPath)
         if (!value) {
-          fileNameToError.set(stdOrLocalPath, [error])
+          fileNameToError.set(localPath, [error])
         } else {
           value.push(error)
-          fileNameToError.set(stdOrLocalPath, [error])
+          fileNameToError.set(localPath, [error])
         }
       }
     }

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

@@ -93,8 +93,14 @@ impl ExecutorContext {
         exec_state: &mut ExecState,
         exec_kind: ExecutionKind,
         preserve_mem: bool,
-    ) -> Result<(Option<KclValue>, EnvironmentRef), KclError> {
+        path: &ModulePath,
+    ) -> Result<(Option<KclValue>, EnvironmentRef, Vec<String>), KclError> {
+        crate::log::log(format!("enter module {path}"));
+
         let old_units = exec_state.length_unit();
+        let original_execution = self.engine.replace_execution_kind(exec_kind).await;
+        let mut local_state = ModuleState::new(&self.settings, path.std_path());
+        std::mem::swap(&mut exec_state.mod_local, &mut local_state);
 
         let no_prelude = self
             .handle_annotations(program.inner_attrs.iter(), crate::execution::BodyType::Root, exec_state)
@@ -103,7 +109,6 @@ impl ExecutorContext {
         if !preserve_mem {
             exec_state.mut_memory().push_new_root_env(!no_prelude);
         }
-        let original_execution = self.engine.replace_execution_kind(exec_kind).await;
 
         let result = self
             .exec_block(program, exec_state, crate::execution::BodyType::Root)
@@ -115,12 +120,16 @@ impl ExecutorContext {
         } else {
             exec_state.mut_memory().pop_env()
         };
+        std::mem::swap(&mut exec_state.mod_local, &mut local_state);
+
         if !exec_kind.is_isolated() && new_units != old_units {
             self.engine.set_units(old_units.into(), Default::default()).await?;
         }
         self.engine.replace_execution_kind(original_execution).await;
 
-        result.map(|result| (result, env_ref))
+        crate::log::log(format!("leave {path}"));
+
+        result.map(|result| (result, env_ref, local_state.module_exports))
     }
 
     /// Execute an AST's program.
@@ -436,31 +445,24 @@ impl ExecutorContext {
         source_range: SourceRange,
     ) -> Result<(Option<KclValue>, EnvironmentRef, Vec<String>), KclError> {
         exec_state.global.mod_loader.enter_module(path);
-        let mut local_state = ModuleState::new(&self.settings, path.std_path());
+        let result = self.exec_module_body(program, exec_state, exec_kind, false, path).await;
-        std::mem::swap(&mut exec_state.mod_local, &mut local_state);
-
-        let result = self.exec_module_body(program, exec_state, exec_kind, false).await;
-
-        std::mem::swap(&mut exec_state.mod_local, &mut local_state);
         exec_state.global.mod_loader.leave_module(path);
 
-        result
+        result.map_err(|err| {
-            .map_err(|err| {
+            if let KclError::ImportCycle(_) = err {
-                if let KclError::ImportCycle(_) = err {
+                // It was an import cycle.  Keep the original message.
-                    // It was an import cycle.  Keep the original message.
+                err.override_source_ranges(vec![source_range])
-                    err.override_source_ranges(vec![source_range])
+            } else {
-                } else {
+                KclError::Semantic(KclErrorDetails {
-                    KclError::Semantic(KclErrorDetails {
+                    message: format!(
-                        message: format!(
+                        "Error loading imported file. Open it to view more details. {}: {}",
-                            "Error loading imported file. Open it to view more details. {}: {}",
+                        path,
-                            path,
+                        err.message()
-                            err.message()
+                    ),
-                        ),
+                    source_ranges: vec![source_range],
-                        source_ranges: vec![source_range],
+                })
-                    })
+            }
-                }
+        })
-            })
-            .map(|(val, env)| (val, env, local_state.module_exports))
     }
 
     #[async_recursion]
@@ -1033,7 +1035,7 @@ impl Node<CallExpressionKw> {
                         // before running, and we will likely want to use the
                         // return value. The call takes ownership of the args,
                         // so we need to build the op before the call.
-                        exec_state.mod_local.operations.push(op);
+                        exec_state.global.operations.push(op);
                     }
                     result
                 }?;
@@ -1055,20 +1057,17 @@ impl Node<CallExpressionKw> {
                     .iter()
                     .map(|(k, arg)| (k.clone(), OpArg::new(OpKclValue::from(&arg.value), arg.source_range)))
                     .collect();
-                exec_state
+                exec_state.global.operations.push(Operation::UserDefinedFunctionCall {
-                    .mod_local
+                    name: Some(fn_name.clone()),
-                    .operations
+                    function_source_range: func.function_def_source_range().unwrap_or_default(),
-                    .push(Operation::UserDefinedFunctionCall {
+                    unlabeled_arg: args
-                        name: Some(fn_name.clone()),
+                        .kw_args
-                        function_source_range: func.function_def_source_range().unwrap_or_default(),
+                        .unlabeled
-                        unlabeled_arg: args
+                        .as_ref()
-                            .kw_args
+                        .map(|arg| OpArg::new(OpKclValue::from(&arg.value), arg.source_range)),
-                            .unlabeled
+                    labeled_args: op_labeled_args,
-                            .as_ref()
+                    source_range: callsite,
-                            .map(|arg| OpArg::new(OpKclValue::from(&arg.value), arg.source_range)),
+                });
-                        labeled_args: op_labeled_args,
-                        source_range: callsite,
-                    });
 
                 let return_value = func
                     .call_fn_kw(args, exec_state, ctx.clone(), callsite)
@@ -1092,10 +1091,7 @@ impl Node<CallExpressionKw> {
                 })?;
 
                 // Track return operation.
-                exec_state
+                exec_state.global.operations.push(Operation::UserDefinedFunctionReturn);
-                    .mod_local
-                    .operations
-                    .push(Operation::UserDefinedFunctionReturn);
 
                 Ok(result)
             }
@@ -1174,7 +1170,7 @@ impl Node<CallExpression> {
                         // before running, and we will likely want to use the
                         // return value. The call takes ownership of the args,
                         // so we need to build the op before the call.
-                        exec_state.mod_local.operations.push(op);
+                        exec_state.global.operations.push(op);
                     }
                     result
                 }?;
@@ -1190,17 +1186,14 @@ impl Node<CallExpression> {
                 let func = exec_state.memory().get(fn_name, source_range)?.clone();
 
                 // Track call operation.
-                exec_state
+                exec_state.global.operations.push(Operation::UserDefinedFunctionCall {
-                    .mod_local
+                    name: Some(fn_name.clone()),
-                    .operations
+                    function_source_range: func.function_def_source_range().unwrap_or_default(),
-                    .push(Operation::UserDefinedFunctionCall {
+                    unlabeled_arg: None,
-                        name: Some(fn_name.clone()),
+                    // TODO: Add the arguments for legacy positional parameters.
-                        function_source_range: func.function_def_source_range().unwrap_or_default(),
+                    labeled_args: Default::default(),
-                        unlabeled_arg: None,
+                    source_range: callsite,
-                        // TODO: Add the arguments for legacy positional parameters.
+                });
-                        labeled_args: Default::default(),
-                        source_range: callsite,
-                    });
 
                 let return_value = func
                     .call_fn(fn_args, exec_state, ctx.clone(), source_range)
@@ -1224,10 +1217,7 @@ impl Node<CallExpression> {
                 })?;
 
                 // Track return operation.
-                exec_state
+                exec_state.global.operations.push(Operation::UserDefinedFunctionReturn);
-                    .mod_local
-                    .operations
-                    .push(Operation::UserDefinedFunctionReturn);
 
                 Ok(result)
             }

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

@@ -724,7 +724,7 @@ impl ExecutorContext {
 
                 KclErrorWithOutputs::new(
                     e,
-                    exec_state.mod_local.operations.clone(),
+                    exec_state.global.operations.clone(),
                     exec_state.global.artifact_commands.clone(),
                     exec_state.global.artifact_graph.clone(),
                     module_id_to_module_path,
@@ -760,7 +760,13 @@ impl ExecutorContext {
             .await?;
 
         let exec_result = self
-            .exec_module_body(program, exec_state, ExecutionKind::Normal, preserve_mem)
+            .exec_module_body(
+                program,
+                exec_state,
+                ExecutionKind::Normal,
+                preserve_mem,
+                &ModulePath::Main,
+            )
             .await;
 
         // Move the artifact commands and responses to simplify cache management
@@ -782,7 +788,7 @@ impl ExecutorContext {
         ) {
             Ok(artifact_graph) => {
                 exec_state.global.artifact_graph = artifact_graph;
-                exec_result.map(|(_, env_ref)| env_ref)
+                exec_result.map(|(_, env_ref, _)| env_ref)
             }
             Err(err) => {
                 // Prefer the exec error.

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

@@ -49,6 +49,9 @@ pub(super) struct GlobalState {
     pub artifact_responses: IndexMap<Uuid, WebSocketResponse>,
     /// Output artifact graph.
     pub artifact_graph: ArtifactGraph,
+    /// Operations that have been performed in execution order, for display in
+    /// the Feature Tree.
+    pub operations: Vec<Operation>,
     /// Module loader.
     pub mod_loader: ModuleLoader,
     /// Errors and warnings.
@@ -62,9 +65,6 @@ pub(super) struct ModuleState {
     pub pipe_value: Option<KclValue>,
     /// Identifiers that have been exported from the current module.
     pub module_exports: Vec<String>,
-    /// Operations that have been performed in execution order, for display in
-    /// the Feature Tree.
-    pub operations: Vec<Operation>,
     /// Settings specified from annotations.
     pub settings: MetaSettings,
 }
@@ -119,7 +119,7 @@ impl ExecState {
                 .find_all_in_env(main_ref, |_| true)
                 .map(|(k, v)| (k.clone(), v.clone()))
                 .collect(),
-            operations: self.mod_local.operations,
+            operations: self.global.operations,
             artifacts: self.global.artifacts,
             artifact_commands: self.global.artifact_commands,
             artifact_graph: self.global.artifact_graph,
@@ -224,6 +224,7 @@ impl GlobalState {
             artifact_commands: Default::default(),
             artifact_responses: Default::default(),
             artifact_graph: Default::default(),
+            operations: Default::default(),
             mod_loader: Default::default(),
             errors: Default::default(),
         };
@@ -252,7 +253,6 @@ impl ModuleState {
         ModuleState {
             pipe_value: Default::default(),
             module_exports: Default::default(),
-            operations: Default::default(),
             settings: MetaSettings {
                 default_length_units: exec_settings.units.into(),
                 default_angle_units: Default::default(),

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

@@ -122,6 +122,8 @@ pub enum ModuleRepr {
 #[derive(Debug, Clone, Eq, PartialEq, Deserialize, Serialize, Hash, ts_rs::TS)]
 #[serde(tag = "type")]
 pub enum ModulePath {
+    // The main file of the project.
+    Main,
     Local { value: PathBuf },
     Std { value: String },
 }
@@ -136,8 +138,8 @@ impl ModulePath {
 
     pub(crate) fn std_path(&self) -> Option<String> {
         match self {
-            ModulePath::Local { value: _ } => None,
             ModulePath::Std { value: p } => Some(p.clone()),
+            _ => None,
         }
     }
 
@@ -152,6 +154,7 @@ impl ModulePath {
                     })
                 })
                 .map(str::to_owned),
+            ModulePath::Main => unreachable!(),
         }
     }
 
@@ -179,6 +182,7 @@ impl ModulePath {
 impl fmt::Display for ModulePath {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         match self {
+            ModulePath::Main => write!(f, "main"),
             ModulePath::Local { value: path } => path.display().fmt(f),
             ModulePath::Std { value: s } => write!(f, "std::{s}"),
         }

src/wasm-lib/kcl/tests/add_lots/ops.snap

@@ -1,7 +1,6 @@
 ---
 source: kcl/src/simulation_tests.rs
 description: Operations executed add_lots.kcl
-snapshot_kind: text
 ---
 [
   {

src/wasm-lib/kcl/tests/assembly_non_default_units/artifact_commands.snap

@@ -305,18 +305,6 @@ description: Artifact commands assembly_non_default_units.kcl
       "unit": "in"
     }
   },
-  {
-    "cmdId": "[uuid]",
-    "range": [
-      0,
-      33,
-      5
-    ],
-    "command": {
-      "type": "set_scene_units",
-      "unit": "in"
-    }
-  },
   {
     "cmdId": "[uuid]",
     "range": [
@@ -435,18 +423,6 @@ description: Artifact commands assembly_non_default_units.kcl
       "path_id": "[uuid]"
     }
   },
-  {
-    "cmdId": "[uuid]",
-    "range": [
-      0,
-      0,
-      0
-    ],
-    "command": {
-      "type": "set_scene_units",
-      "unit": "mm"
-    }
-  },
   {
     "cmdId": "[uuid]",
     "range": [
@@ -576,17 +552,5 @@ description: Artifact commands assembly_non_default_units.kcl
       "type": "close_path",
       "path_id": "[uuid]"
     }
-  },
-  {
-    "cmdId": "[uuid]",
-    "range": [
-      0,
-      0,
-      0
-    ],
-    "command": {
-      "type": "set_scene_units",
-      "unit": "mm"
-    }
   }
 ]

src/wasm-lib/kcl/tests/assembly_non_default_units/ops.snap

@@ -2,4 +2,51 @@
 source: kcl/src/simulation_tests.rs
 description: Operations executed assembly_non_default_units.kcl
 ---
-[]
+[
+  {
+    "labeledArgs": {
+      "data": {
+        "value": {
+          "type": "String",
+          "value": "XZ"
+        },
+        "sourceRange": [
+          186,
+          190,
+          3
+        ]
+      }
+    },
+    "name": "startSketchOn",
+    "sourceRange": [
+      172,
+      191,
+      3
+    ],
+    "type": "StdLibCall",
+    "unlabeledArg": null
+  },
+  {
+    "labeledArgs": {
+      "data": {
+        "value": {
+          "type": "String",
+          "value": "XZ"
+        },
+        "sourceRange": [
+          103,
+          107,
+          4
+        ]
+      }
+    },
+    "name": "startSketchOn",
+    "sourceRange": [
+      89,
+      108,
+      4
+    ],
+    "type": "StdLibCall",
+    "unlabeledArg": null
+  }
+]

src/wasm-lib/kcl/tests/basic_fillet_cube_next_adjacent/ops.snap

@@ -1,7 +1,6 @@
 ---
 source: kcl/src/simulation_tests.rs
 description: Operations executed basic_fillet_cube_next_adjacent.kcl
-snapshot_kind: text
 ---
 [
   {

src/wasm-lib/kcl/tests/fillet-and-shell/ops.snap

@@ -1,7 +1,6 @@
 ---
 source: kcl/src/simulation_tests.rs
 description: Operations executed fillet-and-shell.kcl
-snapshot_kind: text
 ---
 [
   {

src/wasm-lib/kcl/tests/i_shape/ops.snap

@@ -1,7 +1,6 @@
 ---
 source: kcl/src/simulation_tests.rs
 description: Operations executed i_shape.kcl
-snapshot_kind: text
 ---
 [
   {

src/wasm-lib/kcl/tests/import_function_not_sketch/ops.snap

@@ -1,6 +1,70 @@
 ---
 source: kcl/src/simulation_tests.rs
 description: Operations executed import_function_not_sketch.kcl
-snapshot_kind: text
 ---
-[]
+[
+  {
+    "labeledArgs": {
+      "data": {
+        "value": {
+          "type": "String",
+          "value": "XY"
+        },
+        "sourceRange": [
+          66,
+          70,
+          3
+        ]
+      }
+    },
+    "name": "startSketchOn",
+    "sourceRange": [
+      52,
+      71,
+      3
+    ],
+    "type": "StdLibCall",
+    "unlabeledArg": null
+  },
+  {
+    "labeledArgs": {
+      "data": {
+        "value": {
+          "type": "Object",
+          "value": {
+            "axis": {
+              "type": "String",
+              "value": "y"
+            }
+          }
+        },
+        "sourceRange": [
+          312,
+          326,
+          3
+        ]
+      },
+      "sketch": {
+        "value": {
+          "type": "Sketch",
+          "value": {
+            "artifactId": "[uuid]"
+          }
+        },
+        "sourceRange": [
+          328,
+          329,
+          3
+        ]
+      }
+    },
+    "name": "revolve",
+    "sourceRange": [
+      304,
+      330,
+      3
+    ],
+    "type": "StdLibCall",
+    "unlabeledArg": null
+  }
+]

src/wasm-lib/kcl/tests/import_side_effect/ops.snap

@@ -2,4 +2,28 @@
 source: kcl/src/simulation_tests.rs
 description: Operations executed import_side_effect.kcl
 ---
-[]
+[
+  {
+    "labeledArgs": {
+      "data": {
+        "value": {
+          "type": "String",
+          "value": "XY"
+        },
+        "sourceRange": [
+          95,
+          99,
+          3
+        ]
+      }
+    },
+    "name": "startSketchOn",
+    "sourceRange": [
+      81,
+      100,
+      3
+    ],
+    "type": "StdLibCall",
+    "unlabeledArg": null
+  }
+]

src/wasm-lib/kcl/tests/import_whole/ops.snap

@@ -3,6 +3,73 @@ source: kcl/src/simulation_tests.rs
 description: Operations executed import_whole.kcl
 ---
 [
+  {
+    "labeledArgs": {
+      "data": {
+        "value": {
+          "type": "String",
+          "value": "XY"
+        },
+        "sourceRange": [
+          50,
+          54,
+          3
+        ]
+      }
+    },
+    "name": "startSketchOn",
+    "sourceRange": [
+      36,
+      55,
+      3
+    ],
+    "type": "StdLibCall",
+    "unlabeledArg": null
+  },
+  {
+    "labeledArgs": {
+      "length": {
+        "value": {
+          "type": "Number",
+          "value": 10.0,
+          "ty": {
+            "type": "Default",
+            "len": {
+              "type": "Inches"
+            },
+            "angle": {
+              "type": "Degrees"
+            }
+          }
+        },
+        "sourceRange": [
+          127,
+          129,
+          3
+        ]
+      }
+    },
+    "name": "extrude",
+    "sourceRange": [
+      110,
+      130,
+      3
+    ],
+    "type": "StdLibCall",
+    "unlabeledArg": {
+      "value": {
+        "type": "Sketch",
+        "value": {
+          "artifactId": "[uuid]"
+        }
+      },
+      "sourceRange": [
+        0,
+        0,
+        0
+      ]
+    }
+  },
   {
     "labeledArgs": {
       "faces": {