parallelized modules from kcl (#6206)

* wip * sketch a bit more; going to pull this out of tests next * wip * lock start things * this was a bad idea * Revert "this was a bad idea" This reverts commit a2092e7ed6. * prepare prelude before spawning * error * poop * yike * :( * ok * Reapply "this was a bad idea" This reverts commit fafdf41093. * chip away more * man this is bad * fix rebase add feature flag Signed-off-by: Jess Frazelle <github@jessfraz.com> * get rid of execution kind Signed-off-by: Jess Frazelle <github@jessfraz.com> * clippy Signed-off-by: Jess Frazelle <github@jessfraz.com> * logs Signed-off-by: Jess Frazelle <github@jessfraz.com> * updates Signed-off-by: Jess Frazelle <github@jessfraz.com> * no extra executes Signed-off-by: Jess Frazelle <github@jessfraz.com> * race w batch Signed-off-by: Jess Frazelle <github@jessfraz.com> * cluppy Signed-off-by: Jess Frazelle <github@jessfraz.com> * no printlns Signed-off-by: Jess Frazelle <github@jessfraz.com> * no printlns Signed-off-by: Jess Frazelle <github@jessfraz.com> * fix source ranges Signed-off-by: Jess Frazelle <github@jessfraz.com> * batch shit Signed-off-by: Jess Frazelle <github@jessfraz.com> * fixes Signed-off-by: Jess Frazelle <github@jessfraz.com> * updates Signed-off-by: Jess Frazelle <github@jessfraz.com> * fix Signed-off-by: Jess Frazelle <github@jessfraz.com> * fix some bugs Signed-off-by: Jess Frazelle <github@jessfraz.com> * fix error Signed-off-by: Jess Frazelle <github@jessfraz.com> * cut 1 * preserve mem * re-ad deep_clone the helper we were calling was pushing a new call, which was hanging out. we can skip the middleman since we already have something properly prepared, just without a stdlib in some cases. * skip non-kcl * clean up source range bug * error message changed the uuids also changed because the error is hit before execute even starts. * typo * rensnapshot a few * order things * MAYBE REVERT LATER: attempt at an ordering * snapsnap * Revert "snapsnap" This reverts commit 7350b32c7d. * Revert "MAYBE REVERT LATER:" This reverts commit ab49f3e85f. * ugh * poop * poop2 * lint * tranche 1 * more * more snaps * snap * more * update * MAYBE REVERT THIS * cache multi-file Signed-off-by: Jess Frazelle <github@jessfraz.com> * addd tests Signed-off-by: Jess Frazelle <github@jessfraz.com> * set to false Signed-off-by: Jess Frazelle <github@jessfraz.com> * add test outputs Signed-off-by: Jess Frazelle <github@jessfraz.com> * clippy Signed-off-by: Jess Frazelle <github@jessfraz.com> * kcl-py-bindings uses carwheel Signed-off-by: Jess Frazelle <github@jessfraz.com> * update snapshots Signed-off-by: Jess Frazelle <github@jessfraz.com> * updates Signed-off-by: Jess Frazelle <github@jessfraz.com> --------- Signed-off-by: Jess Frazelle <github@jessfraz.com> Co-authored-by: Paul R. Tagliamonte <paul@zoo.dev> Co-authored-by: Paul Tagliamonte <paultag@gmail.com>
2025-04-16 11:52:14 -07:00
parent 5586646a60
commit d9fe78171f
195 changed files with 15585 additions and 8887 deletions
--- a/rust/kcl-lib/src/walk/import_graph.rs
+++ b/rust/kcl-lib/src/walk/import_graph.rs
@ -6,31 +6,37 @@ use std::{
 use anyhow::Result;

 use crate::{
-    parsing::ast::types::{ImportPath, NodeRef, Program},
+    errors::KclErrorDetails,
+    modules::{ModulePath, ModuleRepr},
+    parsing::ast::types::{ImportPath, ImportStatement, Node as AstNode, NodeRef, Program},
    walk::{Node, Visitable},
+    ExecState, ExecutorContext, KclError, ModuleId, SourceRange,
 };

-/// Specific dependency between two modules. The 0th element of this tuple
+/// Specific dependency between two modules. The 0th element of this info
 /// is the "importing" module, the 1st is the "imported" module. The 0th
 /// module *depends on* the 1st module.
 type Dependency = (String, String);

 type Graph = Vec<Dependency>;

+type DependencyInfo = (AstNode<ImportStatement>, ModuleId, ModulePath, AstNode<Program>);
+type Universe = HashMap<String, DependencyInfo>;
+
 /// 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, NodeRef<'_, 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<_>>(),
        );
    }
@ -40,7 +46,10 @@ pub fn import_graph(progs: HashMap<String, NodeRef<'_, Program>>) -> Result<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() {
@ -83,7 +92,13 @@ fn topsort(all_modules: &[&str], graph: Graph) -> Result<Vec<Vec<String>>> {
        }

        if stage_modules.is_empty() {
-            anyhow::bail!("imports are acyclic");
+            waiting_modules.sort();
+
+            return Err(KclError::ImportCycle(KclErrorDetails {
+                message: format!("circular import of modules not allowed: {}", waiting_modules.join(", ")),
+                // TODO: we can get the right import lines from the AST, but we don't
+                source_ranges: vec![SourceRange::default()],
+            }));
        }

        // not strictly needed here, but perhaps helpful to avoid thinking
@ -101,33 +116,97 @@ fn topsort(all_modules: &[&str], graph: Graph) -> Result<Vec<Vec<String>>> {
    Ok(order)
 }

-pub(crate) fn import_dependencies(prog: NodeRef<'_, Program>) -> Result<Vec<String>> {
+type ImportDependencies = Vec<(String, AstNode<ImportStatement>, ModulePath)>;
+
+pub(crate) fn import_dependencies(
+    prog: NodeRef<Program>,
+    ctx: &ExecutorContext,
+) -> Result<ImportDependencies, KclError> {
    let ret = Arc::new(Mutex::new(vec![]));

-    fn walk(ret: Arc<Mutex<Vec<String>>>, node: Node<'_>) {
+    fn walk(ret: Arc<Mutex<ImportDependencies>>, node: Node<'_>, ctx: &ExecutorContext) -> Result<(), KclError> {
        if let Node::ImportStatement(is) = node {
-            let dependency = match &is.path {
-                ImportPath::Kcl { filename } => filename.to_string(),
-                ImportPath::Foreign { path } => path.to_string(),
-                ImportPath::Std { path } => path.join("::"),
-            };
+            // We only care about Kcl imports for now.
+            if let ImportPath::Kcl { filename } = &is.path {
+                let resolved_path = ModulePath::from_import_path(&is.path, &ctx.settings.project_directory);

-            ret.lock().unwrap().push(dependency);
+                // 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.clone(), resolved_path));
+            }
        }
+
        for child in node.children().iter() {
-            walk(ret.clone(), *child)
+            walk(ret.clone(), *child, ctx)?;
        }
+
+        Ok(())
    }

-    walk(ret.clone(), prog.into());
+    walk(ret.clone(), prog.into(), ctx)?;

-    let ret = ret.lock().unwrap().clone();
-    Ok(ret)
+    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_stmt, module_path) in modules {
+        if out.contains_key(&filename) {
+            continue;
+        }
+
+        let module_id = ctx
+            .open_module(&import_stmt.path, &[], exec_state, Default::default())
+            .await?;
+
+        let program = {
+            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: vec![import_stmt.into()],
+                }));
+            };
+            let ModuleRepr::Kcl(program, _) = &module_info.repr else {
+                // if it's not a KCL module we can skip it since it has no
+                // dependencies.
+                continue;
+            };
+            program.clone()
+        };
+
+        out.insert(
+            filename.clone(),
+            (import_stmt.clone(), module_id, module_path.clone(), program.clone()),
+        );
+        Box::pin(import_universe(ctx, &program, out, exec_state)).await?;
+    }
+
+    Ok(())
 }

 #[cfg(test)]
 mod tests {
    use super::*;
+    use crate::parsing::ast::types::ImportSelector;

    macro_rules! kcl {
        ( $kcl:expr ) => {{
@ -135,26 +214,41 @@ mod tests {
        }};
    }

-    #[test]
-    fn order_imports() {
+    fn into_module_info(program: AstNode<Program>) -> DependencyInfo {
+        (
+            AstNode::no_src(ImportStatement {
+                selector: ImportSelector::None { alias: None },
+                path: ImportPath::Kcl { filename: "".into() },
+                visibility: Default::default(),
+                digest: None,
+            }),
+            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_info(a));

        let b = kcl!(
            "
 import \"a.kcl\"
 "
        );
-        modules.insert("b.kcl".to_owned(), &b);
+        modules.insert("b.kcl".to_owned(), into_module_info(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]
-    fn order_imports_none() {
+    #[tokio::test]
+    async fn order_imports_none() {
        let mut modules = HashMap::new();

        let a = kcl!(
@ -162,49 +256,51 @@ import \"a.kcl\"
 y = 2
 "
        );
-        modules.insert("a.kcl".to_owned(), &a);
+        modules.insert("a.kcl".to_owned(), into_module_info(a));

        let b = kcl!(
            "
 x = 1
 "
        );
-        modules.insert("b.kcl".to_owned(), &b);
+        modules.insert("b.kcl".to_owned(), into_module_info(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]
-    fn order_imports_2() {
+    #[tokio::test]
+    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_info(a));

        let b = kcl!(
            "
 import \"a.kcl\"
 "
        );
-        modules.insert("b.kcl".to_owned(), &b);
+        modules.insert("b.kcl".to_owned(), into_module_info(b));

        let c = kcl!(
            "
 import \"a.kcl\"
 "
        );
-        modules.insert("c.kcl".to_owned(), &c);
+        modules.insert("c.kcl".to_owned(), into_module_info(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]
-    fn order_imports_cycle() {
+    #[tokio::test]
+    async fn order_imports_cycle() {
        let mut modules = HashMap::new();

        let a = kcl!(
@ -212,15 +308,16 @@ import \"a.kcl\"
 import \"b.kcl\"
 "
        );
-        modules.insert("a.kcl".to_owned(), &a);
+        modules.insert("a.kcl".to_owned(), into_module_info(a));

        let b = kcl!(
            "
 import \"a.kcl\"
 "
        );
-        modules.insert("b.kcl".to_owned(), &b);
+        modules.insert("b.kcl".to_owned(), into_module_info(b));

-        import_graph(modules).unwrap_err();
+        let ctx = ExecutorContext::new_mock().await;
+        import_graph(&modules, &ctx).unwrap_err();
    }
 }