updates

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

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

@@ -147,55 +147,6 @@ 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 {
-            if let BodyItem::ImportStatement(import_stmt) = statement {
-                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>(
@@ -460,7 +411,7 @@ impl ExecutorContext {
         Ok(last_expr)
     }
 
-    pub(super) async fn open_module(
+    pub async fn open_module(
         &self,
         path: &ImportPath,
         attrs: &[Node<Annotation>],
@@ -468,6 +419,7 @@ impl ExecutorContext {
         source_range: SourceRange,
     ) -> Result<ModuleId, KclError> {
         let resolved_path = ModulePath::from_import_path(path, &self.settings.project_directory);
+
         match path {
             ImportPath::Kcl { .. } => {
                 exec_state.global.mod_loader.cycle_check(&resolved_path, source_range)?;
@@ -597,7 +549,7 @@ impl ExecutorContext {
         result
     }
 
-    async fn exec_module_from_ast(
+    pub async fn exec_module_from_ast(
         &self,
         program: &Node<Program>,
         module_id: ModuleId,
@@ -605,6 +557,7 @@ impl ExecutorContext {
         exec_state: &mut ExecState,
         source_range: SourceRange,
     ) -> Result<(Option<KclValue>, EnvironmentRef, Vec<String>), KclError> {
+        println!("exec_module_from_ast {path}");
         exec_state.global.mod_loader.enter_module(path);
         let result = self.exec_module_body(program, exec_state, false, module_id, path).await;
         exec_state.global.mod_loader.leave_module(path);
@@ -2658,6 +2611,7 @@ d = b + c
                     original_file_contents: "".to_owned(),
                 },
                 &mut exec_state,
+                false,
             )
             .await
             .unwrap();

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

@@ -31,14 +31,14 @@ use tokio::task::JoinSet;
 
 use crate::{
     engine::EngineManager,
-    errors::KclError,
+    errors::{KclError, KclErrorDetails},
     execution::{
         artifact::build_artifact_graph,
         cache::{CacheInformation, CacheResult},
         types::{UnitAngle, UnitLen},
     },
     fs::FileManager,
-    modules::{ModuleId, ModulePath},
+    modules::{ModuleId, ModulePath, ModuleRepr},
     parsing::ast::types::{Expr, ImportPath, NodeRef},
     source_range::SourceRange,
     std::StdLib,
@@ -671,7 +671,7 @@ impl ExecutorContext {
             (program, exec_state, false)
         };
 
-        let result = self.inner_run(&program, &mut exec_state, preserve_mem).await;
+        let result = self.run_concurrent(&program, &mut exec_state, preserve_mem).await;
 
         if result.is_err() {
             cache::bust_cache().await;
@@ -704,7 +704,7 @@ impl ExecutorContext {
         program: &crate::Program,
         exec_state: &mut ExecState,
     ) -> Result<(EnvironmentRef, Option<ModelingSessionData>), KclErrorWithOutputs> {
-        self.run_concurrent(program, exec_state).await
+        self.run_concurrent(program, exec_state, false).await
     }
 
     /// Perform the execution of a program using an (experimental!) concurrent
@@ -721,46 +721,112 @@ impl ExecutorContext {
         &self,
         program: &crate::Program,
         exec_state: &mut ExecState,
+        preserve_mem: bool,
     ) -> 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)
+        crate::walk::import_universe(self, &program.ast, &mut universe, exec_state)
             .await
-            .unwrap();
+            .map_err(KclErrorWithOutputs::no_outputs)?;
 
-        for modules in crate::walk::import_graph(&universe).unwrap().into_iter() {
+        for modules in crate::walk::import_graph(&universe, self)
+            .map_err(KclErrorWithOutputs::no_outputs)?
+            .into_iter()
+        {
+            #[cfg(not(target_arch = "wasm32"))]
             let mut set = JoinSet::new();
 
+            #[allow(clippy::type_complexity)]
+            let (results_tx, mut results_rx): (
+                tokio::sync::mpsc::Sender<(
+                    ModuleId,
+                    ModulePath,
+                    Result<(Option<KclValue>, EnvironmentRef, Vec<String>), KclError>,
+                )>,
+                tokio::sync::mpsc::Receiver<_>,
+            ) = tokio::sync::mpsc::channel(1);
+
             for module in modules {
-                let program = universe.get(&module).unwrap().clone();
+                let Some((module_id, module_path, program)) = universe.get(&module) else {
+                    return Err(KclErrorWithOutputs::no_outputs(KclError::Internal(KclErrorDetails {
+                        message: format!("Module {module} not found in universe"),
+                        source_ranges: Default::default(),
+                    })));
+                };
+                let module_id = *module_id;
+                let module_path = module_path.clone();
+                let program = program.clone();
                 let exec_state = exec_state.clone();
                 let exec_ctxt = self.clone();
+                let results_tx = results_tx.clone();
 
+                #[cfg(target_arch = "wasm32")]
+                {
+                    wasm_bindgen_futures::spawn_local(async move {
+                        //set.spawn(async move {
+                        println!("Running module {module} from run_concurrent");
+                        let mut exec_state = exec_state;
+                        let exec_ctxt = exec_ctxt;
+
+                        let result = exec_ctxt
+                            .exec_module_from_ast(
+                                &program,
+                                module_id,
+                                &module_path,
+                                &mut exec_state,
+                                Default::default(),
+                            )
+                            .await;
+
+                        results_tx.send((module_id, module_path, result)).await.unwrap();
+                    });
+                }
+                #[cfg(not(target_arch = "wasm32"))]
+                {
                     set.spawn(async move {
                         println!("Running module {module} from run_concurrent");
                         let mut exec_state = exec_state;
                         let exec_ctxt = exec_ctxt;
-                    let program = program;
 
-                    exec_ctxt
+                        let result = exec_ctxt
-                        .inner_run(
+                            .exec_module_from_ast(
-                            &crate::Program {
+                                &program,
-                                ast: program.clone(),
+                                module_id,
-                                original_file_contents: "".to_owned(),
+                                &module_path,
-                            },
                                 &mut exec_state,
-                            false,
+                                Default::default(),
                             )
-                        .await
+                            .await;
+
+                        results_tx.send((module_id, module_path, result)).await.unwrap();
                     });
                 }
-
-            set.join_all().await;
             }
 
-        self.inner_run(program, exec_state, false).await
+            drop(results_tx);
+
+            while let Some((module_id, _, result)) = results_rx.recv().await {
+                match result {
+                    Ok((_, session_data, variables)) => {
+                        let mut repr = exec_state.global.module_infos[&module_id].take_repr();
+
+                        let ModuleRepr::Kcl(_, cache) = &mut repr else {
+                            continue;
+                        };
+                        *cache = Some((session_data, variables.clone()));
+
+                        exec_state.global.module_infos[&module_id].restore_repr(repr);
+                    }
+                    Err(e) => {
+                        return Err(KclErrorWithOutputs::no_outputs(e));
+                    }
+                }
+            }
+        }
+
+        self.inner_run(program, exec_state, preserve_mem).await
     }
 
     /// Perform the execution of a program.  Accept all possible parameters and
@@ -811,11 +877,11 @@ impl ExecutorContext {
             )
         })?;
 
-        /*        if !self.is_mock() {
+        if !self.is_mock() {
             let mut mem = exec_state.stack().deep_clone();
             mem.restore_env(env_ref);
             cache::write_old_memory((mem, exec_state.global.module_infos.clone())).await;
-        }*/
+        }
         let session_data = self.engine.get_session_data().await;
         Ok((env_ref, session_data))
     }

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

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

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

@@ -6,10 +6,11 @@ use std::{
 use anyhow::Result;
 
 use crate::{
-    fs::FileSystem,
+    errors::KclErrorDetails,
+    modules::{ModulePath, ModuleRepr},
     parsing::ast::types::{ImportPath, Node as AstNode, NodeRef, Program},
     walk::{Node, Visitable},
-    ExecutorContext, SourceRange,
+    ExecState, ExecutorContext, KclError, ModuleId,
 };
 
 /// Specific dependency between two modules. The 0th element of this tuple
@@ -19,20 +20,22 @@ type Dependency = (String, String);
 
 type Graph = Vec<Dependency>;
 
+type Universe = HashMap<String, (ModuleId, ModulePath, AstNode<Program>)>;
+
 /// Process a number of programs, returning the graph of dependencies.
 ///
 /// This will (currently) return a list of lists of IDs that can be safely
 /// 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, AstNode<Program>>) -> Result<Vec<Vec<String>>> {
+pub fn import_graph(progs: &Universe, ctx: &ExecutorContext) -> Result<Vec<Vec<String>>, KclError> {
     let mut graph = Graph::new();
 
-    for (name, program) in progs.iter() {
+    for (name, (_, _, program)) in progs.iter() {
         graph.extend(
-            import_dependencies(program)?
+            import_dependencies(program, ctx)?
                 .into_iter()
-                .map(|dependency| (name.clone(), dependency))
+                .map(|(dependency, _, _)| (name.clone(), dependency))
                 .collect::<Vec<_>>(),
         );
     }
@@ -42,7 +45,10 @@ pub fn import_graph(progs: &HashMap<String, AstNode<Program>>) -> Result<Vec<Vec
 }
 
 #[allow(clippy::iter_over_hash_type)]
-fn topsort(all_modules: &[&str], graph: Graph) -> Result<Vec<Vec<String>>> {
+fn topsort(all_modules: &[&str], graph: Graph) -> Result<Vec<Vec<String>>, KclError> {
+    if all_modules.is_empty() {
+        return Ok(vec![]);
+    }
     let mut dep_map = HashMap::<String, Vec<String>>::new();
 
     for (dependent, dependency) in graph.iter() {
@@ -85,7 +91,10 @@ fn topsort(all_modules: &[&str], graph: Graph) -> Result<Vec<Vec<String>>> {
         }
 
         if stage_modules.is_empty() {
-            anyhow::bail!("imports are acyclic");
+            return Err(KclError::Internal(KclErrorDetails {
+                message: "Circular import detected".to_string(),
+                source_ranges: Default::default(),
+            }));
         }
 
         // not strictly needed here, but perhaps helpful to avoid thinking
@@ -103,58 +112,88 @@ 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>,
+    ctx: &ExecutorContext,
+) -> Result<Vec<(String, ImportPath, ModulePath)>, KclError> {
     let ret = Arc::new(Mutex::new(vec![]));
 
-    fn walk(ret: Arc<Mutex<Vec<String>>>, node: Node<'_>) {
+    fn walk(
-        if let Node::ImportStatement(is) = node {
+        ret: Arc<Mutex<Vec<(String, ImportPath, ModulePath)>>>,
-            let dependency = match &is.path {
+        node: Node<'_>,
-                ImportPath::Kcl { filename } => filename.to_string(),
-                ImportPath::Foreign { path } => path.to_string(),
-                ImportPath::Std { path } => path.join("::"),
-            };
-
-            ret.lock().unwrap().push(dependency);
-        }
-        for child in node.children().iter() {
-            walk(ret.clone(), *child)
-        }
-    }
-
-    walk(ret.clone(), prog.into());
-
-    let ret = ret.lock().unwrap().clone();
-    Ok(ret)
-}
-
-pub(crate) async fn import_universe<'prog>(
         ctx: &ExecutorContext,
-    prog: NodeRef<'prog, Program>,
+    ) -> Result<(), KclError> {
-    out: &mut HashMap<String, AstNode<Program>>,
+        if let Node::ImportStatement(is) = node {
-) -> Result<()> {
+            // We only care about Kcl imports for now.
-    for module in import_dependencies(prog)? {
+            if let ImportPath::Kcl { filename } = &is.path {
-        eprintln!("{:?}", module);
+                let resolved_path = ModulePath::from_import_path(&is.path, &ctx.settings.project_directory);
 
-        if out.contains_key(&module) {
+                // We need to lock the mutex to push the dependency.
+                // This is a bit of a hack, but it works for now.
+                ret.lock()
+                    .map_err(|err| {
+                        KclError::Internal(KclErrorDetails {
+                            message: format!("Failed to lock mutex: {}", err),
+                            source_ranges: Default::default(),
+                        })
+                    })?
+                    .push((filename.to_string(), is.path.clone(), resolved_path));
+            }
+        }
+
+        for child in node.children().iter() {
+            walk(ret.clone(), *child, ctx)?;
+        }
+
+        Ok(())
+    }
+
+    walk(ret.clone(), prog.into(), ctx)?;
+
+    let ret = ret.lock().map_err(|err| {
+        KclError::Internal(KclErrorDetails {
+            message: format!("Failed to lock mutex: {}", err),
+            source_ranges: Default::default(),
+        })
+    })?;
+
+    Ok(ret.clone())
+}
+
+pub(crate) async fn import_universe(
+    ctx: &ExecutorContext,
+    prog: NodeRef<'_, Program>,
+    out: &mut Universe,
+    exec_state: &mut ExecState,
+) -> Result<(), KclError> {
+    let modules = import_dependencies(prog, ctx)?;
+    for (filename, import_path, module_path) in modules {
+        if out.contains_key(&filename) {
             continue;
         }
 
-        // TODO: use open_module and find a way to pass attrs cleanly
+        let module_id = ctx
-        let kcl = ctx
+            .open_module(&import_path, &[], exec_state, Default::default())
-            .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());
+        let program = {
-        Box::pin(import_universe(ctx, &program, out)).await?;
+            let Some(module_info) = exec_state.get_module(module_id) else {
+                return Err(KclError::Internal(KclErrorDetails {
+                    message: format!("Module {} not found", module_id),
+                    source_ranges: Default::default(),
+                }));
+            };
+            let ModuleRepr::Kcl(program, _) = &module_info.repr else {
+                return Err(KclError::Internal(KclErrorDetails {
+                    message: format!("Module {} is not a KCL module", module_id),
+                    source_ranges: Default::default(),
+                }));
+            };
+            program.clone()
+        };
+
+        out.insert(filename.clone(), (module_id, module_path.clone(), program.clone()));
+        Box::pin(import_universe(ctx, &program, out, exec_state)).await?;
     }
 
     Ok(())
@@ -170,26 +209,31 @@ mod tests {
         }};
     }
 
-    #[test]
+    fn into_module_tuple(program: AstNode<Program>) -> (ModuleId, ModulePath, AstNode<Program>) {
-    fn order_imports() {
+        (ModuleId::default(), ModulePath::Local { value: "".into() }, program)
+    }
+
+    #[tokio::test]
+    async fn order_imports() {
         let mut modules = HashMap::new();
 
         let a = kcl!("");
-        modules.insert("a.kcl".to_owned(), a);
+        modules.insert("a.kcl".to_owned(), into_module_tuple(a));
 
         let b = kcl!(
             "
 import \"a.kcl\"
 "
         );
-        modules.insert("b.kcl".to_owned(), b);
+        modules.insert("b.kcl".to_owned(), into_module_tuple(b));
 
-        let order = import_graph(&modules).unwrap();
+        let ctx = ExecutorContext::new_mock().await;
+        let order = import_graph(&modules, &ctx).unwrap();
         assert_eq!(vec![vec!["a.kcl".to_owned()], vec!["b.kcl".to_owned()]], order);
     }
 
-    #[test]
+    #[tokio::test]
-    fn order_imports_none() {
+    async fn order_imports_none() {
         let mut modules = HashMap::new();
 
         let a = kcl!(
@@ -197,49 +241,51 @@ import \"a.kcl\"
 y = 2
 "
         );
-        modules.insert("a.kcl".to_owned(), a);
+        modules.insert("a.kcl".to_owned(), into_module_tuple(a));
 
         let b = kcl!(
             "
 x = 1
 "
         );
-        modules.insert("b.kcl".to_owned(), b);
+        modules.insert("b.kcl".to_owned(), into_module_tuple(b));
 
-        let order = import_graph(&modules).unwrap();
+        let ctx = ExecutorContext::new_mock().await;
+        let order = import_graph(&modules, &ctx).unwrap();
         assert_eq!(vec![vec!["a.kcl".to_owned(), "b.kcl".to_owned()]], order);
     }
 
-    #[test]
+    #[tokio::test]
-    fn order_imports_2() {
+    async fn order_imports_2() {
         let mut modules = HashMap::new();
 
         let a = kcl!("");
-        modules.insert("a.kcl".to_owned(), a);
+        modules.insert("a.kcl".to_owned(), into_module_tuple(a));
 
         let b = kcl!(
             "
 import \"a.kcl\"
 "
         );
-        modules.insert("b.kcl".to_owned(), b);
+        modules.insert("b.kcl".to_owned(), into_module_tuple(b));
 
         let c = kcl!(
             "
 import \"a.kcl\"
 "
         );
-        modules.insert("c.kcl".to_owned(), c);
+        modules.insert("c.kcl".to_owned(), into_module_tuple(c));
 
-        let order = import_graph(&modules).unwrap();
+        let ctx = ExecutorContext::new_mock().await;
+        let order = import_graph(&modules, &ctx).unwrap();
         assert_eq!(
             vec![vec!["a.kcl".to_owned()], vec!["b.kcl".to_owned(), "c.kcl".to_owned()]],
             order
         );
     }
 
-    #[test]
+    #[tokio::test]
-    fn order_imports_cycle() {
+    async fn order_imports_cycle() {
         let mut modules = HashMap::new();
 
         let a = kcl!(
@@ -247,15 +293,16 @@ import \"a.kcl\"
 import \"b.kcl\"
 "
         );
-        modules.insert("a.kcl".to_owned(), a);
+        modules.insert("a.kcl".to_owned(), into_module_tuple(a));
 
         let b = kcl!(
             "
 import \"a.kcl\"
 "
         );
-        modules.insert("b.kcl".to_owned(), b);
+        modules.insert("b.kcl".to_owned(), into_module_tuple(b));
 
-        import_graph(&modules).unwrap_err();
+        let ctx = ExecutorContext::new_mock().await;
+        import_graph(&modules, &ctx).unwrap_err();
     }
 }