Merge remote-tracking branch 'origin/main' into paultag/import

Signed-off-by: Jess Frazelle <github@jessfraz.com>

rust/Cargo.lock

@@ -1907,6 +1907,7 @@ dependencies = [
  "serde_json",
  "sha2",
  "tabled",
+ "tempdir",
  "thiserror 2.0.12",
  "tokio",
  "tokio-tungstenite",
@@ -3084,6 +3085,15 @@ version = "0.8.5"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "2b15c43186be67a4fd63bee50d0303afffcef381492ebe2c5d87f324e1b8815c"
 
+[[package]]
+name = "remove_dir_all"
+version = "0.5.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "3acd125665422973a33ac9d3dd2df85edad0f4ae9b00dafb1a05e43a9f5ef8e7"
+dependencies = [
+ "winapi",
+]
+
 [[package]]
 name = "reqwest"
 version = "0.12.15"
@@ -3779,6 +3789,16 @@ version = "0.13.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "e502f78cdbb8ba4718f566c418c52bc729126ffd16baee5baa718cf25dd5a69a"
 
+[[package]]
+name = "tempdir"
+version = "0.3.7"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "15f2b5fb00ccdf689e0149d1b1b3c03fead81c2b37735d812fa8bddbbf41b6d8"
+dependencies = [
+ "rand 0.4.6",
+ "remove_dir_all",
+]
+
 [[package]]
 name = "tempfile"
 version = "3.19.0"

rust/Cargo.toml

@@ -22,7 +22,6 @@ debug = 0
 
 [profile.dev.package]
 insta = { opt-level = 3 }
-similar = { opt-level = 3 }
 
 [profile.test]
 debug = "line-tables-only"

rust/kcl-lib/Cargo.toml

@@ -69,6 +69,7 @@ serde = { workspace = true }
 serde_json = { workspace = true }
 sha2 = "0.10.8"
 tabled = { version = "0.18.0", optional = true }
+tempdir = "0.3.7"
 thiserror = "2.0.0"
 toml = "0.8.19"
 ts-rs = { version = "10.1.0", features = [

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

@@ -153,11 +153,63 @@ impl ExecutorContext {
         result.map(|result| (result, env_ref, local_state.module_exports))
     }
 
+    /// Execute an AST's program.
+    #[async_recursion]
+    pub(super) async fn preload_all_modules<'a>(
+        &'a self,
+        modules: &mut HashMap<String, Program>,
+        program: NodeRef<'a, Program>,
+        exec_state: &mut ExecState,
+    ) -> Result<(), KclError> {
+        for statement in &program.body {
+            match statement {
+                BodyItem::ImportStatement(import_stmt) => {
+                    let path_str = import_stmt.path.to_string();
+
+                    if modules.contains_key(&path_str) {
+                        // don't waste our time if we've already loaded the
+                        // module.
+                        continue;
+                    }
+
+                    let source_range = SourceRange::from(import_stmt);
+                    let attrs = &import_stmt.outer_attrs;
+                    let module_id = self
+                        .open_module(&import_stmt.path, attrs, exec_state, source_range)
+                        .await?;
+
+                    let Some(module) = exec_state.get_module(module_id) else {
+                        crate::log::log("we got back a module id that doesn't exist");
+                        unreachable!();
+                    };
+
+                    let progn = {
+                        // this dance is to avoid taking out a mut borrow
+                        // below on exec_state after borrowing here. As a
+                        // result, we need to clone (ugh) the program for
+                        // now.
+                        let ModuleRepr::Kcl(ref progn, _) = module.repr else {
+                            // not a kcl file, we can skip this
+                            continue;
+                        };
+                        progn.clone()
+                    };
+
+                    modules.insert(path_str, progn.clone().inner);
+
+                    self.preload_all_modules(modules, &progn, exec_state).await?;
+                }
+                _ => {}
+            };
+        }
+        Ok(())
+    }
+
     /// Execute an AST's program.
     #[async_recursion]
     pub(super) async fn exec_block<'a>(
         &'a self,
-        program: NodeRef<'a, crate::parsing::ast::types::Program>,
+        program: NodeRef<'a, Program>,
         exec_state: &mut ExecState,
         body_type: BodyType,
     ) -> Result<Option<KclValue>, KclError> {
@@ -2301,13 +2353,14 @@ impl FunctionSource {
 
 #[cfg(test)]
 mod test {
-    use std::sync::Arc;
-
     use super::*;
     use crate::{
         execution::{memory::Stack, parse_execute, ContextType},
         parsing::ast::types::{DefaultParamVal, Identifier, Parameter},
+        ExecutorSettings,
     };
+    use std::sync::Arc;
+    use tokio::io::AsyncWriteExt;
 
     #[tokio::test(flavor = "multi_thread")]
     async fn test_assign_args_to_params() {
@@ -2416,7 +2469,7 @@ mod test {
             // Run each test.
             let func_expr = &Node::no_src(FunctionExpression {
                 params,
-                body: crate::parsing::ast::types::Program::empty(),
+                body: Program::empty(),
                 return_type: None,
                 digest: None,
             });
@@ -2518,4 +2571,101 @@ a = foo()
         let result = parse_execute(program).await;
         assert!(result.unwrap_err().to_string().contains("return"));
     }
+
+    #[tokio::test(flavor = "multi_thread")]
+    async fn load_all_modules() {
+        let mut universe = HashMap::<String, Node<Program>>::new();
+
+        // program a.kcl
+        let programa_kcl = r#"
+export a = 1
+"#;
+        // program b.kcl
+        let programb_kcl = r#"
+import a from 'a.kcl'
+
+export b = a + 1
+"#;
+        // program c.kcl
+        let programc_kcl = r#"
+import a from 'a.kcl'
+
+export c = a + 2
+"#;
+
+        // program main.kcl
+        let main_kcl = r#"
+import b from 'b.kcl'
+import c from 'c.kcl'
+
+d = b + c
+"#;
+
+        let main = crate::parsing::parse_str(&main_kcl, ModuleId::default())
+            .parse_errs_as_err()
+            .unwrap();
+
+        let tmpdir = tempdir::TempDir::new("zma_kcl_load_all_modules").unwrap();
+
+        tokio::fs::File::create(tmpdir.path().join("main.kcl"))
+            .await
+            .unwrap()
+            .write_all(main_kcl.as_bytes())
+            .await
+            .unwrap();
+
+        tokio::fs::File::create(tmpdir.path().join("a.kcl"))
+            .await
+            .unwrap()
+            .write_all(programa_kcl.as_bytes())
+            .await
+            .unwrap();
+
+        tokio::fs::File::create(tmpdir.path().join("b.kcl"))
+            .await
+            .unwrap()
+            .write_all(programb_kcl.as_bytes())
+            .await
+            .unwrap();
+
+        tokio::fs::File::create(tmpdir.path().join("c.kcl"))
+            .await
+            .unwrap()
+            .write_all(programc_kcl.as_bytes())
+            .await
+            .unwrap();
+
+        let exec_ctxt = ExecutorContext {
+            engine: Arc::new(Box::new(
+                crate::engine::conn_mock::EngineConnection::new()
+                    .await
+                    .map_err(|err| {
+                        KclError::Internal(crate::errors::KclErrorDetails {
+                            message: format!("Failed to create mock engine connection: {}", err),
+                            source_ranges: vec![SourceRange::default()],
+                        })
+                    })
+                    .unwrap(),
+            )),
+            fs: Arc::new(crate::fs::FileManager::new()),
+            settings: ExecutorSettings {
+                project_directory: Some(tmpdir.path().into()),
+                ..Default::default()
+            },
+            stdlib: Arc::new(crate::std::StdLib::new()),
+            context_type: ContextType::Mock,
+        };
+        let mut exec_state = ExecState::new(&exec_ctxt);
+
+        exec_ctxt
+            .run_concurrent(
+                &crate::Program {
+                    ast: main.clone(),
+                    original_file_contents: "".to_owned(),
+                },
+                &mut exec_state,
+            )
+            .await
+            .unwrap();
+    }
 }

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

@@ -27,6 +27,7 @@ pub use memory::EnvironmentRef;
 use schemars::JsonSchema;
 use serde::{Deserialize, Serialize};
 pub use state::{ExecState, MetaSettings};
+use tokio::task::JoinSet;
 
 use crate::{
     engine::EngineManager,
@@ -710,6 +711,62 @@ impl ExecutorContext {
         self.inner_run(program, exec_state, false).await
     }
 
+    /// Perform the execution of a program using an (experimental!) concurrent
+    /// execution model. This has the same signature as [Self::run].
+    ///
+    /// For now -- do not use this unless you're willing to accept some
+    /// breakage.
+    ///
+    /// You can optionally pass in some initialization memory for partial
+    /// execution.
+    ///
+    /// To access non-fatal errors and warnings, extract them from the `ExecState`.
+    pub async fn run_concurrent(
+        &self,
+        program: &crate::Program,
+        exec_state: &mut ExecState,
+    ) -> Result<(EnvironmentRef, Option<ModelingSessionData>), KclErrorWithOutputs> {
+        self.prepare_mem(exec_state).await.unwrap();
+
+        let mut universe = std::collections::HashMap::new();
+
+        crate::walk::import_universe(self, &program.ast, &mut universe)
+            .await
+            .unwrap();
+
+        for modules in crate::walk::import_graph(&universe).unwrap().into_iter() {
+            let mut set = JoinSet::new();
+
+            for module in modules {
+                let program = universe.get(&module).unwrap().clone();
+                let module = module.clone();
+                let mut exec_state = exec_state.clone();
+                let exec_ctxt = self.clone();
+
+                set.spawn(async move {
+                    let module = module;
+                    let mut exec_state = exec_state;
+                    let exec_ctxt = exec_ctxt;
+                    let program = program;
+
+                    exec_ctxt
+                        .run(
+                            &crate::Program {
+                                ast: program.clone(),
+                                original_file_contents: "".to_owned(),
+                            },
+                            &mut exec_state,
+                        )
+                        .await
+                });
+            }
+
+            set.join_all().await;
+        }
+
+        self.run(&program, exec_state).await
+    }
+
     /// Perform the execution of a program.  Accept all possible parameters and
     /// output everything.
     async fn inner_run(

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

@@ -228,6 +228,10 @@ impl ExecState {
         self.global.module_infos.insert(id, module_info);
     }
 
+    pub(super) fn get_module(&mut self, id: ModuleId) -> Option<&ModuleInfo> {
+        self.global.module_infos.get(&id)
+    }
+
     pub fn length_unit(&self) -> UnitLen {
         self.mod_local.settings.default_length_units
     }

rust/kcl-lib/src/lib.rs

@@ -76,7 +76,7 @@ pub mod std;
 pub mod test_server;
 mod thread;
 mod unparser;
-mod walk;
+pub mod walk;
 #[cfg(target_arch = "wasm32")]
 mod wasm;
 

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

@@ -6,8 +6,10 @@ use std::{
 use anyhow::Result;
 
 use crate::{
-    parsing::ast::types::{ImportPath, NodeRef, Program},
+    fs::FileSystem,
+    parsing::ast::types::{ImportPath, Node as AstNode, NodeRef, Program},
     walk::{Node, Visitable},
+    ExecutorContext, SourceRange,
 };
 
 /// Specific dependency between two modules. The 0th element of this tuple
@@ -23,7 +25,7 @@ type Graph = Vec<Dependency>;
 /// 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: HashMap<String, NodeRef<'_, Program>>) -> Result<Vec<Vec<String>>> {
+pub fn import_graph(progs: &HashMap<String, AstNode<Program>>) -> Result<Vec<Vec<String>>> {
     let mut graph = Graph::new();
 
     for (name, program) in progs.iter() {
@@ -101,7 +103,7 @@ fn topsort(all_modules: &[&str], graph: Graph) -> Result<Vec<Vec<String>>> {
     Ok(order)
 }
 
-pub(crate) fn import_dependencies(prog: NodeRef<'_, Program>) -> Result<Vec<String>> {
+pub(crate) fn import_dependencies(prog: NodeRef<Program>) -> Result<Vec<String>> {
     let ret = Arc::new(Mutex::new(vec![]));
 
     fn walk(ret: Arc<Mutex<Vec<String>>>, node: Node<'_>) {
@@ -125,6 +127,39 @@ pub(crate) fn import_dependencies(prog: NodeRef<'_, Program>) -> Result<Vec<Stri
     Ok(ret)
 }
 
+pub(crate) async fn import_universe<'prog>(
+    ctx: &ExecutorContext,
+    prog: NodeRef<'prog, Program>,
+    out: &mut HashMap<String, AstNode<Program>>,
+) -> Result<()> {
+    for module in import_dependencies(prog)? {
+        eprintln!("{:?}", module);
+
+        if out.contains_key(&module) {
+            continue;
+        }
+
+        // TODO: use open_module and find a way to pass attrs cleanly
+        let kcl = ctx
+            .fs
+            .read_to_string(
+                ctx.settings
+                    .project_directory
+                    .clone()
+                    .unwrap_or("".into())
+                    .join(&module),
+                SourceRange::default(),
+            )
+            .await?;
+        let program = crate::parsing::parse_str(&kcl, crate::ModuleId::default()).parse_errs_as_err()?;
+
+        out.insert(module.clone(), program.clone());
+        Box::pin(import_universe(ctx, &program, out)).await?;
+    }
+
+    Ok(())
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -140,16 +175,16 @@ mod tests {
         let mut modules = HashMap::new();
 
         let a = kcl!("");
-        modules.insert("a.kcl".to_owned(), &a);
+        modules.insert("a.kcl".to_owned(), a);
 
         let b = kcl!(
             "
 import \"a.kcl\"
 "
         );
-        modules.insert("b.kcl".to_owned(), &b);
+        modules.insert("b.kcl".to_owned(), b);
 
-        let order = import_graph(modules).unwrap();
+        let order = import_graph(&modules).unwrap();
         assert_eq!(vec![vec!["a.kcl".to_owned()], vec!["b.kcl".to_owned()]], order);
     }
 
@@ -162,16 +197,16 @@ import \"a.kcl\"
 y = 2
 "
         );
-        modules.insert("a.kcl".to_owned(), &a);
+        modules.insert("a.kcl".to_owned(), a);
 
         let b = kcl!(
             "
 x = 1
 "
         );
-        modules.insert("b.kcl".to_owned(), &b);
+        modules.insert("b.kcl".to_owned(), b);
 
-        let order = import_graph(modules).unwrap();
+        let order = import_graph(&modules).unwrap();
         assert_eq!(vec![vec!["a.kcl".to_owned(), "b.kcl".to_owned()]], order);
     }
 
@@ -180,23 +215,23 @@ x = 1
         let mut modules = HashMap::new();
 
         let a = kcl!("");
-        modules.insert("a.kcl".to_owned(), &a);
+        modules.insert("a.kcl".to_owned(), a);
 
         let b = kcl!(
             "
 import \"a.kcl\"
 "
         );
-        modules.insert("b.kcl".to_owned(), &b);
+        modules.insert("b.kcl".to_owned(), b);
 
         let c = kcl!(
             "
 import \"a.kcl\"
 "
         );
-        modules.insert("c.kcl".to_owned(), &c);
+        modules.insert("c.kcl".to_owned(), c);
 
-        let order = import_graph(modules).unwrap();
+        let order = import_graph(&modules).unwrap();
         assert_eq!(
             vec![vec!["a.kcl".to_owned()], vec!["b.kcl".to_owned(), "c.kcl".to_owned()]],
             order
@@ -212,15 +247,15 @@ import \"a.kcl\"
 import \"b.kcl\"
 "
         );
-        modules.insert("a.kcl".to_owned(), &a);
+        modules.insert("a.kcl".to_owned(), a);
 
         let b = kcl!(
             "
 import \"a.kcl\"
 "
         );
-        modules.insert("b.kcl".to_owned(), &b);
+        modules.insert("b.kcl".to_owned(), b);
 
-        import_graph(modules).unwrap_err();
+        import_graph(&modules).unwrap_err();
     }
 }

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

@@ -8,3 +8,5 @@ mod import_graph;
 pub use ast_node::Node;
 pub use ast_visitor::Visitable;
 pub use ast_walk::walk;
+pub use import_graph::import_graph;
+pub(crate) use import_graph::import_universe;

rust/kcl-wasm-lib/Cargo.toml

@@ -24,7 +24,7 @@ kcl-lib = { path = "../kcl-lib" }
 kittycad = { workspace = true }
 kittycad-modeling-cmds = { workspace = true }
 serde_json = { workspace = true }
-tokio = { workspace = true, features = ["sync"] }
+tokio = { workspace = true, features = ["sync", "rt"] }
 toml = "0.8.19"
 tower-lsp = { workspace = true, features = ["runtime-agnostic"] }
 uuid = { workspace = true, features = ["v4", "js", "serde"] }