Fix stack overflows caused by semantic index recursion

Cargo.toml

@@ -111,6 +111,7 @@ serde_with = "3.19.0"
 sha2 = "0.11.0"
 smallvec = "1.15.1"
 socket2 = "0.6.3"
+stacker = "0.1.21"
 stdext = { path = "crates/stdext" }
 streaming-iterator = "0.1.9"
 strum = "0.28.0"

crates/ark/build.rs

@@ -17,6 +17,25 @@ fn main() {
     println!("cargo:rerun-if-changed=src");
     println!("cargo:rerun-if-changed=resources");
 
+    // Every first-party crate directory name. Used by `logger::internal_crates()`
+    // to `ark`'s `RUST_LOG` level to all our crates.
+    let workspace_crates_dir = Path::new(&std::env::var("CARGO_MANIFEST_DIR").unwrap())
+        .join("..")
+        .canonicalize()
+        .unwrap();
+    println!("cargo:rerun-if-changed={}", workspace_crates_dir.display());
+
+    let internal_crates: Vec<String> = std::fs::read_dir(&workspace_crates_dir)
+        .unwrap()
+        .filter_map(|entry| entry.ok())
+        .filter(|entry| entry.path().is_dir())
+        .filter_map(|entry| entry.file_name().into_string().ok())
+        .collect();
+    println!(
+        "cargo:rustc-env=ARK_INTERNAL_CRATES={}",
+        internal_crates.join(",")
+    );
+
     // Attempt to use `git rev-parse HEAD` to get the current git hash. If this
     // fails, we'll just use the string "<unknown>" to indicate that the git hash
     // could not be determined..

crates/ark/src/logger.rs

@@ -18,6 +18,13 @@ use tracing_subscriber::Layer;
 
 use crate::logger_hprof;
 
+/// Every first-party crate's directory name under `crates/`, baked in by
+/// `build.rs` so [`internal_crates()`] doesn't need a hardcoded list that
+/// goes stale as crates are added.
+fn internal_crates() -> impl Iterator<Item = &'static str> {
+    env!("ARK_INTERNAL_CRATES").split(',')
+}
+
 pub fn init(log_file: Option<&str>, profile_file: Option<&str>) {
     static ONCE: Once = Once::new();
 
@@ -35,7 +42,7 @@ pub fn init(log_file: Option<&str>, profile_file: Option<&str>) {
             .and_then(|c| c.get(1))
             .map(|c| c.as_str())
         {
-            for pkg in ["amalthea", "harp", "stdext"] {
+            for pkg in internal_crates().filter(|&pkg| pkg != "ark") {
                 if let Ok(directive) = format!("{pkg}={level}").parse() {
                     env_filter = env_filter.add_directive(directive);
                 }
@@ -111,3 +118,25 @@ fn non_blocking(file: Option<&str>, cell: &OnceCell<WorkerGuard>) -> BoxMakeWrit
         BoxMakeWriter::new(std::io::stderr)
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::internal_crates;
+
+    #[test]
+    fn test_internal_crates_includes_path_crates() {
+        let crates: Vec<&str> = internal_crates().collect();
+        assert!(crates.contains(&"oak_db"));
+        assert!(crates.contains(&"oak_semantic"));
+        assert!(crates.contains(&"amalthea"));
+        assert!(crates.contains(&"harp"));
+        assert!(crates.contains(&"stdext"));
+    }
+
+    #[test]
+    fn test_internal_crates_excludes_git_dependencies() {
+        let crates: Vec<&str> = internal_crates().collect();
+        assert!(!crates.contains(&"aether_factory"));
+        assert!(!crates.contains(&"aether_syntax"));
+    }
+}

crates/ark/src/lsp/analysis.rs

@@ -22,6 +22,7 @@ pub(crate) use pool::AnalysisPool;
 pub(crate) use refresh::DiagnosticsReady;
 pub(crate) use refresh::DiagnosticsState;
 pub(crate) use snapshot::WorldStateSnapshot;
+pub(crate) use warmup::warm_semantic_indexes;
 pub(crate) use warmup::warm_workspace_index;
 
 /// Run `f`, swallowing a salsa cancellation as `None`. Any other panic propagates.

crates/ark/src/lsp/analysis/warmup.rs

@@ -5,6 +5,9 @@
 //
 //
 
+use oak_db::all_used_files;
+use oak_db::warm_file;
+
 use super::pool::AnalysisPool;
 use crate::lsp;
 use crate::lsp::indexer;
@@ -30,3 +33,21 @@ pub(crate) fn warm_workspace_index(state: &WorldState, pool: &AnalysisPool) {
         lsp::log_info!("Finished workspace index warmup ({:.0?})", now.elapsed());
     })
 }
+
+/// Warm the oak `semantic_index` of every file the workspace depends on, on a
+/// background thread.
+///
+/// Idempotent, so re-running on every revision is cheap once a file's index is
+/// already warm (salsa cache hit). A concurrent write just cancels the
+/// in-flight warm; the next revision re-runs it, which is what carries warmup
+/// through the startup write-storm and warms a freshly-typed `pkg::`
+/// dependency as soon as its sources land.
+pub(crate) fn warm_semantic_indexes(state: &WorldState, pool: &AnalysisPool) {
+    pool.spawn(state.snapshot(), |snapshot| {
+        let now = std::time::Instant::now();
+        for &file in all_used_files(snapshot.db()) {
+            warm_file(snapshot.db(), file);
+        }
+        lsp::log_info!("Warmed semantic indexes ({:.0?})", now.elapsed());
+    })
+}

crates/ark/src/lsp/main_loop.rs

@@ -571,10 +571,9 @@ impl GlobalState {
 
                 dispatch_scan_requests(&self.lsp_state.scan_pool, &self.events_tx, followups);
 
-                // Warm the workspace index once the scan settles. Editor
-                // writes don't need to re-warm: they imply an open document,
-                // and the diagnostics passes they trigger force the same
-                // memos.
+                // Warm the workspace symbol index once the scan settles. The
+                // oak semantic index is warmed separately on every revision (see
+                // the revision-advanced block below).
                 if !self.lsp_state.oak_scheduler.has_pending_scans() {
                     analysis::warm_workspace_index(&self.world, &self.lsp_state.analysis_pool);
                 }
@@ -631,6 +630,12 @@ impl GlobalState {
                 &self.lsp_state.source_pool,
                 &self.events_tx,
             );
+
+            // Re-warm the oak semantic indexes on every revision, counting on
+            // idempotence (warm files are salsa cache hits). Takes care of
+            // warming up the initial workspace as well as any new dependency
+            // introduced by user edits.
+            analysis::warm_semantic_indexes(&self.world, &self.lsp_state.analysis_pool);
         }
 
         Ok(())

crates/ark/src/traps.rs

@@ -39,4 +39,7 @@ pub extern "C-unwind" fn backtrace_handler(signum: libc::c_int) {
     // capture the current thread's backtrace
     let bt = std::backtrace::Backtrace::force_capture();
     log::error!("{}\n{}", header, bt);
+
+    log::logger().flush();
+    std::thread::sleep(std::time::Duration::from_millis(250));
 }

crates/oak_db/Cargo.toml

@@ -29,6 +29,7 @@ oak_semantic = { workspace = true, features = ["salsa"] }
 rustc-hash.workspace = true
 salsa.workspace = true
 salsa-macros.workspace = true
+stacker.workspace = true
 stdext.workspace = true
 url.workspace = true
 

crates/oak_db/src/db.rs

@@ -104,36 +104,57 @@ pub(crate) fn live_roots_query(db: &dyn Db) -> Vec<LiveRoot> {
     roots
 }
 
-/// All files known to the database, in stable order (workspace, library, orphan).
-///
-/// Used as the workspace-wide candidate pool for find-references: callers
-/// apply a textual name filter before building indexes.
+/// Files reachable from workspace roots and orphan buffers, plus library
+/// files that belong to an actual workspace dependency (see
+/// [`crate::all_package_dependencies`]).
 ///
-/// Nested roots overlap on disk, so the same `File` is reachable from several
-/// roots (open `/proj` and `/proj/sub-pkg` and the outer scan walks into
-/// `sub-pkg`). A `seen` set drops the repeats. Unlike `root_by_package` /
-/// `root_by_file`, this query exposes no ownership, just a flat set, so it
-/// doesn't matter which root a duplicate is attributed to. We keep the first
-/// occurrence, which preserves the traversal order above.
+/// Same as [`all_known_files`] but excluding library packages that aren't a
+/// dependency (direct or indirect) of the workspace.
 #[salsa::tracked(returns(ref))]
-pub fn all_files(db: &dyn Db) -> Vec<File> {
+pub fn all_used_files(db: &dyn Db) -> Vec<File> {
+    let dependencies: FxHashSet<Package> = crate::workspace::all_package_dependencies(db)
+        .iter()
+        .copied()
+        .collect();
+
     let mut seen = FxHashSet::default();
     let mut files = Vec::new();
 
     for &root in db.live_roots() {
         match root {
-            LiveRoot::Workspace(r) | LiveRoot::Library(r) => {
-                let root_files = r.scripts(db).iter().chain(
-                    r.packages(db)
-                        .iter()
-                        .flat_map(|&pkg| pkg.files(db).iter().chain(pkg.scripts(db))),
-                );
-                for &file in root_files {
+            LiveRoot::Workspace(root) => push_root_files(db, &mut files, &mut seen, root, None),
+            LiveRoot::Library(root) => {
+                push_root_files(db, &mut files, &mut seen, root, Some(&dependencies))
+            },
+            LiveRoot::Orphan(orphan) => {
+                for &file in orphan.files(db) {
                     if seen.insert(file) {
                         files.push(file);
                     }
                 }
             },
+        }
+    }
+
+    files
+}
+
+/// All files known to the database, in stable order (workspace, library, orphan).
+///
+/// Note that this also contains files from all installed packages, _including
+/// those that are not dependencies of the workspace_. Only use this for very
+/// wide searches. LSP functionality should generally not depend on
+/// non-dependencies, prefer [`all_used_files()`] instead.
+#[salsa::tracked(returns(ref))]
+pub fn all_known_files(db: &dyn Db) -> Vec<File> {
+    let mut seen = FxHashSet::default();
+    let mut files = Vec::new();
+
+    for &root in db.live_roots() {
+        match root {
+            LiveRoot::Workspace(root) | LiveRoot::Library(root) => {
+                push_root_files(db, &mut files, &mut seen, root, None)
+            },
             LiveRoot::Orphan(orphan) => {
                 for &file in orphan.files(db) {
                     if seen.insert(file) {
@@ -147,6 +168,29 @@ pub fn all_files(db: &dyn Db) -> Vec<File> {
     files
 }
 
+/// Pushes `root`'s scripts and package files into `files`, skipping ones
+/// already in `seen`. When `dependencies` is `Some`, packages not in that set
+/// are skipped entirely.
+fn push_root_files(
+    db: &dyn Db,
+    files: &mut Vec<File>,
+    seen: &mut FxHashSet<File>,
+    root: Root,
+    dependencies: Option<&FxHashSet<Package>>,
+) {
+    let root_files = root.scripts(db).iter().chain(
+        root.packages(db)
+            .iter()
+            .filter(|&&pkg| dependencies.is_none_or(|deps| deps.contains(&pkg)))
+            .flat_map(|&pkg| pkg.files(db).iter().chain(pkg.scripts(db))),
+    );
+    for &file in root_files {
+        if seen.insert(file) {
+            files.push(file);
+        }
+    }
+}
+
 /// Files eligible for the workspace symbol index: workspace-root scripts and
 /// package files, plus orphan editor buffers. Library roots are excluded, so
 /// installed package symbols don't leak into e.g. workspace symbols.