Add minimal support for cargo scripts

Author: HKalbasi

crates/project-model/src/cargo_workspace.rs

@@ -276,6 +276,10 @@ impl CargoWorkspace {
                     .collect(),
             );
         }
+        if cargo_toml.extension().is_some_and(|x| x == "rs") {
+            // TODO: enable `+nightly` for cargo scripts
+            other_options.push("-Zscript".to_owned());
+        }
         meta.other_options(other_options);
 
         // FIXME: Fetching metadata is a slow process, as it might require

crates/project-model/src/lib.rs

@@ -50,7 +50,7 @@ pub use crate::{
     manifest_path::ManifestPath,
     project_json::{ProjectJson, ProjectJsonData},
     sysroot::Sysroot,
-    workspace::{CfgOverrides, PackageRoot, ProjectWorkspace},
+    workspace::{CargoScriptTomls, CfgOverrides, PackageRoot, ProjectWorkspace},
 };
 
 #[derive(Debug, Clone, PartialEq, Eq, Hash, Ord, PartialOrd)]

crates/project-model/src/workspace.rs

@@ -2,7 +2,7 @@
 //! metadata` or `rust-project.json`) into representation stored in the salsa
 //! database -- `CrateGraph`.
 
-use std::{collections::VecDeque, fmt, fs, process::Command, sync};
+use std::{collections::VecDeque, fmt, fs, io::BufRead, process::Command, sync};
 
 use anyhow::{format_err, Context};
 use base_db::{
@@ -96,9 +96,54 @@ pub enum ProjectWorkspace {
         /// Holds cfg flags for the current target. We get those by running
         /// `rustc --print cfg`.
         rustc_cfg: Vec<CfgFlag>,
+        cargo_script: Option<CargoWorkspace>,
     },
 }
 
+/// Tracks the cargo toml parts in cargo scripts, to detect if they
+/// changed and reload workspace in that case.
+pub struct CargoScriptTomls(pub FxHashMap<AbsPathBuf, String>);
+
+impl CargoScriptTomls {
+    fn extract_toml_part(p: &AbsPath) -> Option<String> {
+        let mut r = String::new();
+        let f = std::fs::File::open(p).ok()?;
+        let f = std::io::BufReader::new(f);
+        let mut started = false;
+        for line in f.lines() {
+            let line = line.ok()?;
+            if started {
+                if line.trim() == "//! ```" {
+                    return Some(r);
+                }
+                r += &line;
+            } else {
+                if line.trim() == "//! ```cargo" {
+                    started = true;
+                }
+            }
+        }
+        None
+    }
+
+    pub fn track_file(&mut self, p: AbsPathBuf) {
+        let toml = CargoScriptTomls::extract_toml_part(&p).unwrap_or_default();
+        self.0.insert(p, toml);
+    }
+
+    pub fn need_reload(&mut self, p: &AbsPath) -> bool {
+        let Some(prev) = self.0.get_mut(p) else {
+            return false; // File is not tracked
+        };
+        let next = CargoScriptTomls::extract_toml_part(p).unwrap_or_default();
+        if *prev == next {
+            return false;
+        }
+        *prev = next;
+        true
+    }
+}
+
 impl fmt::Debug for ProjectWorkspace {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         // Make sure this isn't too verbose.
@@ -136,7 +181,7 @@ impl fmt::Debug for ProjectWorkspace {
                 debug_struct.field("n_rustc_cfg", &rustc_cfg.len());
                 debug_struct.finish()
             }
-            ProjectWorkspace::DetachedFiles { files, sysroot, rustc_cfg } => f
+            ProjectWorkspace::DetachedFiles { files, sysroot, rustc_cfg, cargo_script: _ } => f
                 .debug_struct("DetachedFiles")
                 .field("n_files", &files.len())
                 .field("sysroot", &sysroot.is_ok())
@@ -342,6 +387,7 @@ impl ProjectWorkspace {
     pub fn load_detached_files(
         detached_files: Vec<AbsPathBuf>,
         config: &CargoConfig,
+        cargo_script_tomls: &mut CargoScriptTomls,
     ) -> anyhow::Result<ProjectWorkspace> {
         let sysroot = match &config.sysroot {
             Some(RustLibSource::Path(path)) => Sysroot::with_sysroot_dir(path.clone())
@@ -361,7 +407,24 @@ impl ProjectWorkspace {
             tracing::info!(src_root = %sysroot.src_root(), root = %sysroot.root(), "Using sysroot");
         }
         let rustc_cfg = rustc_cfg::get(None, None, &Default::default());
-        Ok(ProjectWorkspace::DetachedFiles { files: detached_files, sysroot, rustc_cfg })
+        let cargo_toml = ManifestPath::try_from(detached_files[0].clone()).unwrap();
+        let meta =
+            CargoWorkspace::fetch_metadata(&cargo_toml, cargo_toml.parent(), config, &|_| ())
+                .with_context(|| {
+                    format!("Failed to read Cargo metadata from Cargo.toml file {cargo_toml}")
+                })?;
+        let cargo = CargoWorkspace::new(meta);
+
+        for file in &detached_files {
+            cargo_script_tomls.track_file(file.clone());
+        }
+
+        Ok(ProjectWorkspace::DetachedFiles {
+            files: detached_files,
+            sysroot,
+            rustc_cfg,
+            cargo_script: Some(cargo),
+        })
     }
 
     /// Runs the build scripts for this [`ProjectWorkspace`].
@@ -628,14 +691,29 @@ impl ProjectWorkspace {
                 },
                 toolchain.as_ref().and_then(|it| ReleaseChannel::from_str(it.pre.as_str())),
             ),
-            ProjectWorkspace::DetachedFiles { files, sysroot, rustc_cfg } => {
-                detached_files_to_crate_graph(
-                    rustc_cfg.clone(),
-                    load,
-                    files,
-                    sysroot.as_ref().ok(),
-                    Err("detached file projects have no target layout set".into()),
-                )
+            ProjectWorkspace::DetachedFiles { files, sysroot, rustc_cfg, cargo_script } => {
+                if let Some(cargo) = cargo_script {
+                    cargo_to_crate_graph(
+                        load,
+                        None,
+                        cargo,
+                        sysroot.as_ref().ok(),
+                        rustc_cfg.clone(),
+                        &CfgOverrides::default(),
+                        None,
+                        &WorkspaceBuildScripts::default(),
+                        Err("detached file projects have no target layout set".into()),
+                        None,
+                    )
+                } else {
+                    detached_files_to_crate_graph(
+                        rustc_cfg.clone(),
+                        load,
+                        files,
+                        sysroot.as_ref().ok(),
+                        Err("detached file projects have no target layout set".into()),
+                    )
+                }
             }
         };
         if crate_graph.patch_cfg_if() {
@@ -692,9 +770,19 @@ impl ProjectWorkspace {
                     && toolchain == o_toolchain
             }
             (
-                Self::DetachedFiles { files, sysroot, rustc_cfg },
-                Self::DetachedFiles { files: o_files, sysroot: o_sysroot, rustc_cfg: o_rustc_cfg },
-            ) => files == o_files && sysroot == o_sysroot && rustc_cfg == o_rustc_cfg,
+                Self::DetachedFiles { files, sysroot, rustc_cfg, cargo_script },
+                Self::DetachedFiles {
+                    files: o_files,
+                    sysroot: o_sysroot,
+                    rustc_cfg: o_rustc_cfg,
+                    cargo_script: o_cargo_script,
+                },
+            ) => {
+                files == o_files
+                    && sysroot == o_sysroot
+                    && rustc_cfg == o_rustc_cfg
+                    && cargo_script == o_cargo_script
+            }
             _ => false,
         }
     }

crates/rust-analyzer/src/global_state.rs

@@ -14,7 +14,9 @@ use lsp_types::{SemanticTokens, Url};
 use nohash_hasher::IntMap;
 use parking_lot::{Mutex, RwLock};
 use proc_macro_api::ProcMacroServer;
-use project_model::{CargoWorkspace, ProjectWorkspace, Target, WorkspaceBuildScripts};
+use project_model::{
+    CargoScriptTomls, CargoWorkspace, ProjectWorkspace, Target, WorkspaceBuildScripts,
+};
 use rustc_hash::{FxHashMap, FxHashSet};
 use triomphe::Arc;
 use vfs::AnchoredPathBuf;
@@ -121,6 +123,7 @@ pub(crate) struct GlobalState {
         OpQueue<(), (Arc<Vec<ProjectWorkspace>>, Vec<anyhow::Result<WorkspaceBuildScripts>>)>,
     pub(crate) fetch_proc_macros_queue: OpQueue<Vec<ProcMacroPaths>, bool>,
     pub(crate) prime_caches_queue: OpQueue,
+    pub(crate) cargo_script_tomls: Arc<Mutex<CargoScriptTomls>>,
 }
 
 /// An immutable snapshot of the world's state at a point in time.
@@ -204,6 +207,7 @@ impl GlobalState {
             fetch_proc_macros_queue: OpQueue::default(),
 
             prime_caches_queue: OpQueue::default(),
+            cargo_script_tomls: Arc::new(Mutex::new(CargoScriptTomls(FxHashMap::default()))),
         };
         // Apply any required database inputs from the config.
         this.update_configuration(config);
@@ -276,7 +280,11 @@ impl GlobalState {
                 let vfs_path = &vfs.file_path(file.file_id);
                 if let Some(path) = vfs_path.as_path() {
                     let path = path.to_path_buf();
-                    if reload::should_refresh_for_change(&path, file.change_kind) {
+                    if reload::should_refresh_for_change(
+                        &path,
+                        file.change_kind,
+                        &mut self.cargo_script_tomls.lock(),
+                    ) {
                         workspace_structure_change = Some((path.clone(), false));
                     }
                     if file.is_created_or_deleted() {

crates/rust-analyzer/src/handlers/notification.rs

@@ -124,7 +124,11 @@ pub(crate) fn handle_did_save_text_document(
     if let Ok(vfs_path) = from_proto::vfs_path(&params.text_document.uri) {
         // Re-fetch workspaces if a workspace related file has changed
         if let Some(abs_path) = vfs_path.as_path() {
-            if reload::should_refresh_for_change(abs_path, ChangeKind::Modify) {
+            if reload::should_refresh_for_change(
+                abs_path,
+                ChangeKind::Modify,
+                &mut state.cargo_script_tomls.lock(),
+            ) {
                 state
                     .fetch_workspaces_queue
                     .request_op(format!("DidSaveTextDocument {abs_path}"), false);

crates/rust-analyzer/src/reload.rs

@@ -23,7 +23,7 @@ use ide_db::{
 };
 use load_cargo::{load_proc_macro, ProjectFolders};
 use proc_macro_api::ProcMacroServer;
-use project_model::{ProjectWorkspace, WorkspaceBuildScripts};
+use project_model::{CargoScriptTomls, ProjectWorkspace, WorkspaceBuildScripts};
 use rustc_hash::FxHashSet;
 use stdx::{format_to, thread::ThreadIntent};
 use triomphe::Arc;
@@ -189,6 +189,7 @@ impl GlobalState {
             let linked_projects = self.config.linked_projects();
             let detached_files = self.config.detached_files().to_vec();
             let cargo_config = self.config.cargo();
+            let cargo_script_tomls = self.cargo_script_tomls.clone();
 
             move |sender| {
                 let progress = {
@@ -245,6 +246,7 @@ impl GlobalState {
                     workspaces.push(project_model::ProjectWorkspace::load_detached_files(
                         detached_files,
                         &cargo_config,
+                        &mut cargo_script_tomls.lock(),
                     ));
                 }
 
@@ -622,7 +624,15 @@ impl GlobalState {
     }
 }
 
-pub(crate) fn should_refresh_for_change(path: &AbsPath, change_kind: ChangeKind) -> bool {
+pub(crate) fn should_refresh_for_change(
+    path: &AbsPath,
+    change_kind: ChangeKind,
+    cargo_script_tomls: &mut CargoScriptTomls,
+) -> bool {
+    if cargo_script_tomls.need_reload(path) {
+        return true;
+    }
+
     const IMPLICIT_TARGET_FILES: &[&str] = &["build.rs", "src/main.rs", "src/lib.rs"];
     const IMPLICIT_TARGET_DIRS: &[&str] = &["src/bin", "examples", "tests", "benches"];
 

crates/rust-analyzer/tests/slow-tests/main.rs

@@ -76,6 +76,112 @@ use std::collections::Spam;
     assert!(res.to_string().contains("HashMap"));
 }
 
+#[test]
+fn completes_items_from_standard_library_in_cargo_script() {
+    if skip_slow_tests() {
+        return;
+    }
+
+    let server = Project::with_fixture(
+        r#"
+//- /dependency/Cargo.toml
+[package]
+name = "dependency"
+version = "0.1.0"
+
+//- /dependency/src/lib.rs
+pub struct SpecialHashMap;
+
+//- /dependency2/Cargo.toml
+[package]
+name = "dependency2"
+version = "0.1.0"
+
+//- /dependency2/src/lib.rs
+pub struct SpecialHashMap2;
+
+//- /src/lib.rs
+#!/usr/bin/env -S cargo +nightly -Zscript
+
+//! ```cargo
+//! [dependencies]
+//! dependency = { path = "../dependency" }
+//! ```
+
+use dependency::Spam;
+use dependency2::Spam;
+"#,
+    )
+    .with_config(serde_json::json!({
+        "cargo": { "sysroot": "discover" },
+    }))
+    .server()
+    .wait_until_workspace_is_loaded();
+
+    let res = server.send_request::<Completion>(CompletionParams {
+        text_document_position: TextDocumentPositionParams::new(
+            server.doc_id("src/lib.rs"),
+            Position::new(7, 18),
+        ),
+        context: None,
+        partial_result_params: PartialResultParams::default(),
+        work_done_progress_params: WorkDoneProgressParams::default(),
+    });
+    assert!(res.to_string().contains("SpecialHashMap"));
+
+    let res = server.send_request::<Completion>(CompletionParams {
+        text_document_position: TextDocumentPositionParams::new(
+            server.doc_id("src/lib.rs"),
+            Position::new(8, 18),
+        ),
+        context: None,
+        partial_result_params: PartialResultParams::default(),
+        work_done_progress_params: WorkDoneProgressParams::default(),
+    });
+    assert!(!res.to_string().contains("SpecialHashMap"));
+
+    server.write_file_and_save(
+        "src/lib.rs",
+        r#"#!/usr/bin/env -S cargo +nightly -Zscript
+
+//! ```cargo
+//! [dependencies]
+//! dependency2 = { path = "../dependency2" }
+//! ```
+
+use dependency::Spam;
+use dependency2::Spam;
+"#
+        .to_owned(),
+    );
+
+    let server = server.wait_until_workspace_is_loaded();
+
+    std::thread::sleep(std::time::Duration::from_secs(3));
+
+    let res = server.send_request::<Completion>(CompletionParams {
+        text_document_position: TextDocumentPositionParams::new(
+            server.doc_id("src/lib.rs"),
+            Position::new(7, 18),
+        ),
+        context: None,
+        partial_result_params: PartialResultParams::default(),
+        work_done_progress_params: WorkDoneProgressParams::default(),
+    });
+    assert!(!res.to_string().contains("SpecialHashMap"));
+
+    let res = server.send_request::<Completion>(CompletionParams {
+        text_document_position: TextDocumentPositionParams::new(
+            server.doc_id("src/lib.rs"),
+            Position::new(8, 18),
+        ),
+        context: None,
+        partial_result_params: PartialResultParams::default(),
+        work_done_progress_params: WorkDoneProgressParams::default(),
+    });
+    assert!(res.to_string().contains("SpecialHashMap"));
+}
+
 #[test]
 fn test_runnables_project() {
     if skip_slow_tests() {