load sources in parallel

roryc89 · roryc89 · commit a1746fe10fbb · 2026-03-09T21:26:23.000+01:00
diff --git a/src/build/mod.rs b/src/build/mod.rs
@@ -317,20 +317,27 @@ fn build_from_sources_impl(
         sources.len() - skip_count
     );
 
-    // Step 3: Parse only non-cached sources in parallel
-    let parse_results: Vec<(usize, Result<(PathBuf, Module), BuildError>)> = sources
-        .par_iter()
-        .enumerate()
-        .filter(|(i, _)| !skip_parse[*i])
-        .map(|(i, &(path_str, source))| {
-            let path = PathBuf::from(path_str);
-            let result = match crate::parser::parse(source) {
-                Ok(module) => Ok((path, module)),
-                Err(e) => Err(BuildError::CompileError { path, error: e }),
-            };
-            (i, result)
-        })
-        .collect();
+    // Step 3: Parse only non-cached sources in parallel (use a pool with large stacks
+    // since the parser can recurse deeply on complex files)
+    let parse_pool = rayon::ThreadPoolBuilder::new()
+        .stack_size(16 * 1024 * 1024)
+        .build()
+        .expect("failed to build parse thread pool");
+    let parse_results: Vec<(usize, Result<(PathBuf, Module), BuildError>)> = parse_pool.install(|| {
+        sources
+            .par_iter()
+            .enumerate()
+            .filter(|(i, _)| !skip_parse[*i])
+            .map(|(i, &(path_str, source))| {
+                let path = PathBuf::from(path_str);
+                let result = match crate::parser::parse(source) {
+                    Ok(module) => Ok((path, module)),
+                    Err(e) => Err(BuildError::CompileError { path, error: e }),
+                };
+                (i, result)
+            })
+            .collect()
+    });
 
     // Step 4: Build parsed vec from both cached stubs and parsed results
     let mut parsed: Vec<ParsedModule> = Vec::new();
diff --git a/src/lsp/handlers/diagnostics.rs b/src/lsp/handlers/diagnostics.rs
@@ -63,7 +63,7 @@ impl Backend {
         let module_name = interner::resolve_module_name(&module.name.value.parts);
         let module_parts: Vec<interner::Symbol> = module.name.value.parts.clone();
 
-        // Type-check against the registry
+        // Type-check against the registry (use stacker to extend stack for deep recursion)
         let mut registry = self.registry.write().await;
         let check_result = crate::typechecker::check_module_with_registry(&module, &registry);
 
diff --git a/src/lsp/handlers/load_sources.rs b/src/lsp/handlers/load_sources.rs
@@ -1,6 +1,7 @@
 use std::collections::HashMap;
 use std::sync::atomic::Ordering;
 
+use rayon::prelude::*;
 use tower_lsp::lsp_types::*;
 
 use crate::build::BuildOptions;
@@ -51,7 +52,12 @@ impl Backend {
         let ready = self.ready.clone();
         let progress_token = token.clone();
 
-        tokio::task::spawn_blocking(move || {
+        let rt_handle = tokio::runtime::Handle::current();
+        std::thread::Builder::new()
+            .name("pfc-load-sources".to_string())
+            .stack_size(16 * 1024 * 1024) // 16 MB — typechecker needs deep recursion
+            .spawn(move || {
+            let _guard = rt_handle.enter();
             // Run the shell command to get source globs
             let output = match std::process::Command::new("sh")
                 .arg("-c")
@@ -101,34 +107,40 @@ impl Backend {
                     .await;
             });
 
-            // Resolve globs to file paths
-            let mut sources: Vec<(String, String)> = Vec::new();
+            // Resolve globs to file paths (collect paths first, then read in parallel)
+            let mut file_paths: Vec<std::path::PathBuf> = Vec::new();
             for pattern in &globs {
                 match glob::glob(pattern) {
                     Ok(entries) => {
                         for entry in entries.flatten() {
                             if entry.extension().map_or(false, |ext| ext == "purs") {
-                                match std::fs::read_to_string(&entry) {
-                                    Ok(source) => {
-                                        let abs_path = entry
-                                            .canonicalize()
-                                            .unwrap_or_else(|_| entry.clone());
-                                        sources.push((
-                                            abs_path.to_string_lossy().into_owned(),
-                                            source,
-                                        ));
-                                    }
-                                    Err(e) => {
-                                        log::warn!("Failed to read {}: {e}", entry.display())
-                                    }
-                                }
+                                file_paths.push(entry);
                             }
                         }
                     }
                     Err(e) => log::warn!("Invalid glob pattern {pattern}: {e}"),
                 }
             }
 
+            // Read all files in parallel
+            let sources: Vec<(String, String)> = file_paths
+                .par_iter()
+                .filter_map(|entry| {
+                    match std::fs::read_to_string(entry) {
+                        Ok(source) => {
+                            let abs_path = entry
+                                .canonicalize()
+                                .unwrap_or_else(|_| entry.clone());
+                            Some((abs_path.to_string_lossy().into_owned(), source))
+                        }
+                        Err(e) => {
+                            log::warn!("Failed to read {}: {e}", entry.display());
+                            None
+                        }
+                    }
+                })
+                .collect();
+
             // Report progress: building
             rt.block_on(async {
                 client
@@ -179,29 +191,37 @@ impl Backend {
                 .filter(|m| !m.type_errors.is_empty())
                 .count();
 
-            // Build definition index and resolution exports from parsed sources
-            let mut index = DefinitionIndex::new();
-            let mut smap = HashMap::new();
-            let mut mfmap = HashMap::new();
-            let mut parsed_modules = Vec::new();
-            for (path, source) in &sources {
-                if let Ok(module) = crate::parser::parse(source) {
-                    index.add_module(&module, path);
-                    let mod_name = format!("{}", module.name.value);
+            // Parse all sources in parallel for definition index
+            let parse_results: Vec<_> = sources
+                .par_iter()
+                .map(|(path, source)| {
                     let file_uri = Url::from_file_path(path)
                         .map(|u| u.to_string())
                         .unwrap_or_default();
+                    match crate::parser::parse(source) {
+                        Ok(module) => {
+                            let mod_name = format!("{}", module.name.value);
+                            (path.clone(), file_uri, source.clone(), Some((module, mod_name)))
+                        }
+                        Err(_) => {
+                            (path.clone(), file_uri, source.clone(), None)
+                        }
+                    }
+                })
+                .collect();
+
+            // Merge results sequentially (add_module takes &mut self)
+            let mut index = DefinitionIndex::new();
+            let mut smap = HashMap::with_capacity(parse_results.len());
+            let mut mfmap = HashMap::new();
+            let mut parsed_modules = Vec::new();
+            for (path, file_uri, source, parsed) in parse_results {
+                if let Some((module, mod_name)) = parsed {
+                    index.add_module(&module, &path);
                     mfmap.insert(mod_name, file_uri.clone());
                     parsed_modules.push(module);
-                    smap.insert(file_uri, source.clone());
-                } else {
-                    smap.insert(
-                        Url::from_file_path(path)
-                            .map(|u| u.to_string())
-                            .unwrap_or_default(),
-                        source.clone(),
-                    );
                 }
+                smap.insert(file_uri, source);
             }
 
             let exports = crate::lsp::utils::resolve::ResolutionExports::new(&parsed_modules);
@@ -234,6 +254,7 @@ impl Backend {
                     })
                     .await;
             });
-        });
+        })
+        .expect("failed to spawn load-sources thread");
     }
 }
diff --git a/src/lsp/mod.rs b/src/lsp/mod.rs
@@ -154,7 +154,11 @@ impl Backend {
 }
 
 pub fn run_server(sources_cmd: Option<String>) {
-    let rt = tokio::runtime::Runtime::new().expect("failed to create tokio runtime");
+    let rt = tokio::runtime::Builder::new_multi_thread()
+        .enable_all()
+        .thread_stack_size(16 * 1024 * 1024) // 16 MB — typechecker needs deep recursion
+        .build()
+        .expect("failed to create tokio runtime");
     rt.block_on(async {
         let stdin = tokio::io::stdin();
         let stdout = tokio::io::stdout();
