Prefer core/alloc over std if no_std is conditional

2 files changed, 66 insertions, 4 deletions

diff --git a/crates/hir_def/src/db.rs b/crates/hir_def/src/db.rs
index 74bb7472d57..df6dcb024b5 100644
--- a/crates/hir_def/src/db.rs
+++ b/crates/hir_def/src/db.rs
@@ -18,7 +18,7 @@ use crate::{
     generics::GenericParams,
     import_map::ImportMap,
     intern::Interned,
-    item_tree::ItemTree,
+    item_tree::{AttrOwner, ItemTree},
     lang_item::{LangItemTarget, LangItems},
     nameres::DefMap,
     visibility::{self, Visibility},
@@ -184,6 +184,8 @@ pub trait DefDatabase: InternDatabase + AstDatabase + Upcast<dyn AstDatabase> {
 
     #[salsa::transparent]
     fn crate_limits(&self, crate_id: CrateId) -> CrateLimits;
+
+    fn crate_supports_no_std(&self, crate_id: CrateId) -> bool;
 }
 
 fn crate_def_map_wait(db: &dyn DefDatabase, krate: CrateId) -> Arc<DefMap> {
@@ -204,3 +206,38 @@ fn crate_limits(db: &dyn DefDatabase, crate_id: CrateId) -> CrateLimits {
         recursion_limit: def_map.recursion_limit().unwrap_or(128),
     }
 }
+
+fn crate_supports_no_std(db: &dyn DefDatabase, crate_id: CrateId) -> bool {
+    let file = db.crate_graph()[crate_id].root_file_id;
+    let item_tree = db.file_item_tree(file.into());
+    let attrs = item_tree.raw_attrs(AttrOwner::TopLevel);
+    for attr in &**attrs {
+        match attr.path().as_ident().and_then(|id| id.as_text()) {
+            Some(ident) if ident == "no_std" => return true,
+            Some(ident) if ident == "cfg_attr" => {}
+            _ => continue,
+        }
+
+        // This is a `cfg_attr`; check if it could possibly expand to `no_std`.
+        // Syntax is: `#[cfg_attr(condition(cfg, style), attr0, attr1, <...>)]`
+        let tt = match attr.token_tree_value() {
+            Some(tt) => &tt.token_trees,
+            None => continue,
+        };
+
+        let segments = tt.split(|tt| match tt {
+            tt::TokenTree::Leaf(tt::Leaf::Punct(p)) if p.char == ',' => true,
+            _ => false,
+        });
+        for output in segments.skip(1) {
+            match output {
+                [tt::TokenTree::Leaf(tt::Leaf::Ident(ident))] if ident.text == "no_std" => {
+                    return true
+                }
+                _ => {}
+            }
+        }
+    }
+
+    false
+}
diff --git a/crates/hir_def/src/find_path.rs b/crates/hir_def/src/find_path.rs
index bb89b8cff44..89e961f84fa 100644
--- a/crates/hir_def/src/find_path.rs
+++ b/crates/hir_def/src/find_path.rs
@@ -43,8 +43,7 @@ impl ModPathExt for ModPath {
         self.segments().first() == Some(&known::std)
     }
 
-    // When std library is present, paths starting with `std::`
-    // should be preferred over paths starting with `core::` and `alloc::`
+    // Can we replace the first segment with `std::` and still get a valid, identical path?
     fn can_start_with_std(&self) -> bool {
         let first_segment = self.segments().first();
         first_segment == Some(&known::alloc) || first_segment == Some(&known::core)
@@ -203,7 +202,7 @@ fn find_path_inner_(
     }
 
     // - otherwise, look for modules containing (reexporting) it and import it from one of those
-    let prefer_no_std = db.attrs(crate_root.into()).by_key("no_std").exists();
+    let prefer_no_std = db.crate_supports_no_std(crate_root.krate);
     let mut best_path = None;
     let mut best_path_len = max_len;
 
@@ -839,6 +838,32 @@ pub mod fmt {
             "core::fmt::Error",
             "core::fmt::Error",
         );
+
+        // Should also work (on a best-effort basis) if `no_std` is conditional.
+        check_found_path(
+            r#"
+//- /main.rs crate:main deps:core,std
+#![cfg_attr(not(test), no_std)]
+
+$0
+
+//- /std.rs crate:std deps:core
+
+pub mod fmt {
+    pub use core::fmt::Error;
+}
+
+//- /zzz.rs crate:core
+
+pub mod fmt {
+    pub struct Error;
+}
+        "#,
+            "core::fmt::Error",
+            "core::fmt::Error",
+            "core::fmt::Error",
+            "core::fmt::Error",
+        );
     }
 
     #[test]