Change mock IDs to be stable

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

@@ -488,6 +488,11 @@ impl ArtifactGraph {
     pub fn len(&self) -> usize {
         self.map.len()
     }
+
+    #[cfg(test)]
+    pub(crate) fn iter(&self) -> impl Iterator<Item = (&ArtifactId, &Artifact)> {
+        self.map.iter()
+    }
 }
 
 pub(super) fn build_artifact_graph(

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

@@ -5,6 +5,7 @@ use std::sync::Arc;
 use itertools::{EitherOrBoth, Itertools};
 use tokio::sync::RwLock;
 
+use super::IdGenerator;
 use crate::{
     execution::{annotations, memory::Stack, EnvironmentRef, ExecState, ExecutorSettings},
     parsing::ast::types::{Annotation, Node, Program},
@@ -16,6 +17,8 @@ lazy_static::lazy_static! {
     static ref OLD_AST: Arc<RwLock<Option<OldAstState>>> = Default::default();
     // The last successful run's memory. Not cleared after an unsuccessful run.
     static ref PREV_MEMORY: Arc<RwLock<Option<Stack>>> = Default::default();
+    /// The ID generator for mock execution.
+    static ref MOCK_ID_GENERATOR: Arc<RwLock<Option<IdGenerator>>> = Default::default();
 }
 
 /// Read the old ast memory from the lock.
@@ -49,6 +52,21 @@ pub async fn clear_mem_cache() {
     *old_mem = None;
 }
 
+pub(crate) async fn read_mock_ids() -> Option<IdGenerator> {
+    let cache = MOCK_ID_GENERATOR.read().await;
+    cache.clone()
+}
+
+pub(super) async fn write_mock_ids(id_gen: IdGenerator) {
+    let mut cache = MOCK_ID_GENERATOR.write().await;
+    *cache = Some(id_gen);
+}
+
+pub async fn clear_mock_ids() {
+    let mut cache = MOCK_ID_GENERATOR.write().await;
+    *cache = None;
+}
+
 /// Information for the caching an AST and smartly re-executing it if we can.
 #[derive(Debug, Clone)]
 pub struct CacheInformation<'a> {

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

@@ -7,7 +7,7 @@ pub use artifact::{
     Artifact, ArtifactCommand, ArtifactGraph, ArtifactId, CodeRef, StartSketchOnFace, StartSketchOnPlane,
 };
 use cache::OldAstState;
-pub use cache::{bust_cache, clear_mem_cache};
+pub use cache::{bust_cache, clear_mem_cache, clear_mock_ids};
 pub use cad_op::Operation;
 pub use geometry::*;
 pub(crate) use import::{
@@ -518,7 +518,9 @@ impl ExecutorContext {
     ) -> Result<ExecOutcome, KclErrorWithOutputs> {
         assert!(self.is_mock());
 
-        let mut exec_state = ExecState::new(&self.settings);
+        let mut id_generator = cache::read_mock_ids().await.unwrap_or_default();
+        id_generator.next_id = 0;
+        let mut exec_state = ExecState::with_ids(&self.settings, id_generator);
         if use_prev_memory {
             match cache::read_old_memory().await {
                 Some(mem) => *exec_state.mut_stack() = mem,
@@ -534,6 +536,10 @@ impl ExecutorContext {
 
         let result = self.inner_run(&program, &mut exec_state, true).await?;
 
+        // Mock execution has its own ID generator.  Save it so that multiple executions get the
+        // same IDs.
+        cache::write_mock_ids(exec_state.global.id_generator.clone()).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
         // memory, not to the exec_state which is not cached for mock execution.
@@ -1990,4 +1996,22 @@ let w = f() + f()
         let result = ctx2.run_mock(program2, true).await.unwrap();
         assert_eq!(result.variables.get("z").unwrap().as_f64().unwrap(), 3.0);
     }
+
+    #[tokio::test(flavor = "multi_thread")]
+    async fn mock_has_stable_ids() {
+        let ctx = ExecutorContext::new_mock().await;
+        let code = "sk = startSketchOn(XY)
+        |> startProfileAt([0, 0], %)";
+        let program = crate::Program::parse_no_errs(code).unwrap();
+        let result = ctx.run_mock(program, false).await.unwrap();
+        let ids = result.artifact_graph.iter().map(|(k, _)| *k).collect::<Vec<_>>();
+        assert!(!ids.is_empty(), "IDs should not be empty");
+
+        let ctx2 = ExecutorContext::new_mock().await;
+        let program2 = crate::Program::parse_no_errs(code).unwrap();
+        let result = ctx2.run_mock(program2, false).await.unwrap();
+        let ids2 = result.artifact_graph.iter().map(|(k, _)| *k).collect::<Vec<_>>();
+
+        assert_eq!(ids, ids2, "Generated IDs should match");
+    }
 }

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

@@ -74,8 +74,12 @@ pub(super) struct ModuleState {
 
 impl ExecState {
     pub fn new(exec_settings: &ExecutorSettings) -> Self {
+        Self::with_ids(exec_settings, IdGenerator::default())
+    }
+
+    pub fn with_ids(exec_settings: &ExecutorSettings, id_generator: IdGenerator) -> Self {
         ExecState {
-            global: GlobalState::new(exec_settings),
+            global: GlobalState::new(exec_settings, id_generator),
             mod_local: ModuleState::new(exec_settings, None, ProgramMemory::new()),
         }
     }
@@ -86,8 +90,7 @@ impl ExecState {
         // This is for the front end to keep track of the ids.
         id_generator.next_id = 0;
 
-        let mut global = GlobalState::new(exec_settings);
+        let global = GlobalState::new(exec_settings, id_generator);
-        global.id_generator = id_generator;
 
         *self = ExecState {
             global,
@@ -145,8 +148,8 @@ impl ExecState {
                 .map(|(k, v)| (k.clone(), v.clone()))
                 .collect(),
             operations: Default::default(),
-            artifact_commands: Default::default(),
+            artifact_commands: self.global.artifact_commands,
-            artifact_graph: Default::default(),
+            artifact_graph: self.global.artifact_graph,
             errors: self.global.errors,
             filenames: Default::default(),
         }
@@ -239,9 +242,9 @@ impl ExecState {
 }
 
 impl GlobalState {
-    fn new(settings: &ExecutorSettings) -> Self {
+    fn new(settings: &ExecutorSettings, id_generator: IdGenerator) -> Self {
         let mut global = GlobalState {
-            id_generator: Default::default(),
+            id_generator,
             path_to_source_id: Default::default(),
             module_infos: Default::default(),
             artifacts: Default::default(),

rust/kcl-lib/src/lib.rs

@@ -86,7 +86,8 @@ pub use errors::{
     CompilationError, ConnectionError, ExecError, KclError, KclErrorWithOutputs, Report, ReportWithOutputs,
 };
 pub use execution::{
-    bust_cache, clear_mem_cache, ExecOutcome, ExecState, ExecutorContext, ExecutorSettings, MetaSettings, Point2d,
+    bust_cache, clear_mem_cache, clear_mock_ids, ExecOutcome, ExecState, ExecutorContext, ExecutorSettings,
+    MetaSettings, Point2d,
 };
 pub use lsp::{
     copilot::Backend as CopilotLspBackend,

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

@@ -5,8 +5,8 @@ use std::sync::Arc;
 use futures::stream::TryStreamExt;
 use gloo_utils::format::JsValueSerdeExt;
 use kcl_lib::{
-    bust_cache, clear_mem_cache, exec::IdGenerator, pretty::NumericSuffix, CoreDump, EngineManager, ModuleId, Point2d,
+    bust_cache, clear_mem_cache, clear_mock_ids, exec::IdGenerator, pretty::NumericSuffix, CoreDump, EngineManager,
-    Program,
+    ModuleId, Point2d, Program,
 };
 use tower_lsp::{LspService, Server};
 use wasm_bindgen::prelude::*;
@@ -20,6 +20,7 @@ pub async fn clear_scene_and_bust_cache(
 
     bust_cache().await;
     clear_mem_cache().await;
+    clear_mock_ids().await;
 
     let engine = kcl_lib::wasm_engine::EngineConnection::new(engine_manager)
         .await