intellij-rust#8224: fix type inference if dependencies have names as stdlib crates

Author: Undin

src/main/kotlin/org/rust/cargo/project/workspace/CargoWorkspace.kt

@@ -9,6 +9,7 @@ import com.intellij.openapi.util.UserDataHolderBase
 import com.intellij.openapi.util.UserDataHolderEx
 import com.intellij.openapi.vfs.VirtualFile
 import com.intellij.openapi.vfs.VirtualFileManager
+import com.intellij.util.ThreeState
 import org.jetbrains.annotations.TestOnly
 import org.rust.cargo.CfgOptions
 import org.rust.cargo.project.model.CargoProjectsService
@@ -48,7 +49,14 @@ interface CargoWorkspace {
     val featureGraph: FeatureGraph
 
     fun findPackageById(id: PackageId): Package? = packages.find { it.id == id }
-    fun findPackageByName(name: String): Package? = packages.find { it.name == name || it.normName == name }
+    fun findPackageByName(name: String, isStd: ThreeState = ThreeState.UNSURE): Package? = packages.find {
+        if (it.name != name && it.normName != name) return@find false
+        when (isStd) {
+            ThreeState.YES -> it.origin == STDLIB
+            ThreeState.NO -> it.origin == WORKSPACE || it.origin == DEPENDENCY
+            ThreeState.UNSURE -> true
+        }
+    }
 
     fun findTargetByCrateRoot(root: VirtualFile): Target?
     fun isCrateRoot(root: VirtualFile) = findTargetByCrateRoot(root) != null

src/main/kotlin/org/rust/lang/core/resolve/KnownItems.kt

@@ -10,6 +10,7 @@ import com.intellij.openapi.util.Key
 import com.intellij.psi.util.CachedValue
 import com.intellij.psi.util.CachedValueProvider
 import com.intellij.psi.util.CachedValuesManager
+import com.intellij.util.ThreeState
 import org.rust.cargo.project.model.CargoProject
 import org.rust.cargo.project.workspace.CargoWorkspace
 import org.rust.cargo.util.AutoInjectedCrates.CORE
@@ -51,8 +52,8 @@ class KnownItems(
     fun findLangItemRaw(langAttribute: String, crateName: String) =
         lookup.findLangItem(langAttribute, crateName)
 
-    fun findItemRaw(path: String): RsNamedElement? =
-        lookup.findItem(path)
+    fun findItemRaw(path: String, isStd: Boolean): RsNamedElement? =
+        lookup.findItem(path, isStd)
 
     /**
      * Find some known item by its "lang" attribute
@@ -66,8 +67,8 @@ class KnownItems(
         crateName: String = CORE
     ): T? = findLangItemRaw(langAttribute, crateName) as? T
 
-    inline fun <reified T : RsNamedElement> findItem(path: String): T? =
-        findItemRaw(path) as? T
+    inline fun <reified T : RsNamedElement> findItem(path: String, isStd: Boolean = true): T? =
+        findItemRaw(path, isStd) as? T
 
     val Vec: RsStructOrEnumItemElement? get() = findItem("alloc::vec::Vec")
     val String: RsStructOrEnumItemElement? get() = findItem("alloc::string::String")
@@ -145,12 +146,12 @@ class KnownItems(
 
 interface KnownItemsLookup {
     fun findLangItem(langAttribute: String, crateName: String): RsNamedElement?
-    fun findItem(path: String): RsNamedElement?
+    fun findItem(path: String, isStd: Boolean): RsNamedElement?
 }
 
 private object DummyKnownItemsLookup : KnownItemsLookup {
     override fun findLangItem(langAttribute: String, crateName: String): RsNamedElement? = null
-    override fun findItem(path: String): RsNamedElement? = null
+    override fun findItem(path: String, isStd: Boolean): RsNamedElement? = null
 }
 
 private class RealKnownItemsLookup(
@@ -167,9 +168,9 @@ private class RealKnownItemsLookup(
         }.orElse(null)
     }
 
-    override fun findItem(path: String): RsNamedElement? {
+    override fun findItem(path: String, isStd: Boolean): RsNamedElement? {
         return resolvedItems.getOrPut(path) {
-            Optional.ofNullable(resolveStringPath(path, workspace, project)?.first)
+            Optional.ofNullable(resolveStringPath(path, workspace, project, ThreeState.fromBoolean(isStd))?.first)
         }.orElse(null)
     }
 }
@@ -189,11 +190,11 @@ enum class KnownDerivableTrait(
     PartialOrd(KnownItems::PartialOrd, arrayOf(PartialEq)),
     Ord(KnownItems::Ord, arrayOf(PartialOrd, Eq, PartialEq)),
 
-    Serialize({ it.findItem("serde::Serialize") }, isStd = false),
-    Deserialize({ it.findItem("serde::Deserialize") }, isStd = false),
+    Serialize({ it.findItem("serde::Serialize", isStd = false) }, isStd = false),
+    Deserialize({ it.findItem("serde::Deserialize", isStd = false) }, isStd = false),
 
     // TODO Fail also derives `Display`. Ignore it for now
-    Fail({ it.findItem("failure::Fail") }, arrayOf(Debug), isStd = false),
+    Fail({ it.findItem("failure::Fail", isStd = false) }, arrayOf(Debug), isStd = false),
     ;
 
     fun findTrait(items: KnownItems): RsTraitItem? = resolver(items)

src/main/kotlin/org/rust/lang/core/resolve/NameResolution.kt

@@ -22,6 +22,7 @@ import com.intellij.psi.util.CachedValue
 import com.intellij.psi.util.CachedValueProvider
 import com.intellij.psi.util.CachedValuesManager
 import com.intellij.psi.util.PsiTreeUtil
+import com.intellij.util.ThreeState
 import org.rust.cargo.project.workspace.CargoWorkspace
 import org.rust.cargo.project.workspace.PackageOrigin
 import org.rust.cargo.util.AutoInjectedCrates.CORE
@@ -759,9 +760,14 @@ fun processLocalVariables(place: RsElement, originalProcessor: (RsPatBinding) ->
 /**
  * Resolves an absolute path.
  */
-fun resolveStringPath(path: String, workspace: CargoWorkspace, project: Project): Pair<RsNamedElement, CargoWorkspace.Package>? {
+fun resolveStringPath(
+    path: String,
+    workspace: CargoWorkspace,
+    project: Project,
+    isStd: ThreeState = ThreeState.UNSURE
+): Pair<RsNamedElement, CargoWorkspace.Package>? {
     val (pkgName, crateRelativePath) = splitAbsolutePath(path) ?: return null
-    val pkg = workspace.findPackageByName(pkgName) ?: run {
+    val pkg = workspace.findPackageByName(pkgName, isStd) ?: run {
         return if (isUnitTestMode) {
             // Allows to set a fake path for some item in tests via
             // lang attribute, e.g. `#[lang = "std::iter::Iterator"]`

src/test/kotlin/org/rust/RustProjectDescriptors.kt

@@ -95,6 +95,39 @@ open class RustProjectDescriptorBase : LightProjectDescriptor() {
             CargoWorkspaceData(packages, emptyMap(), emptyMap(), contentRoot), CfgOptions.DEFAULT)
     }
 
+    protected open fun externalPackage(
+        contentRoot: String,
+        source: String?,
+        name: String,
+        targetName: String = name,
+        version: String = "0.0.1",
+        origin: PackageOrigin = PackageOrigin.DEPENDENCY,
+        libKind: LibKind = LibKind.LIB,
+        procMacroArtifact: CargoWorkspaceData.ProcMacroArtifact? = null,
+    ): Package {
+        return Package(
+            id = "$name $version",
+            contentRootUrl = contentRoot,
+            name = name,
+            version = version,
+            targets = listOf(
+                // don't use `FileUtil.join` here because it uses `File.separator`
+                // which is system dependent although all other code uses `/` as separator
+                Target(source?.let { "$contentRoot/$it" } ?: "", targetName,
+                    TargetKind.Lib(libKind), Edition.EDITION_2015, doctest = true, requiredFeatures = emptyList())
+            ),
+            source = source,
+            origin = origin,
+            edition = Edition.EDITION_2015,
+            features = emptyMap(),
+            enabledFeatures = emptySet(),
+            cfgOptions = CfgOptions.EMPTY,
+            env = emptyMap(),
+            outDirUrl = null,
+            procMacroArtifact = procMacroArtifact
+        )
+    }
+
     protected fun testCargoPackage(contentRoot: String, name: String = "test-package") = Package(
         id = "$name 0.0.1",
         contentRootUrl = contentRoot,
@@ -292,38 +325,6 @@ open class WithCustomStdlibRustProjectDescriptor(
 }
 
 object WithDependencyRustProjectDescriptor : RustProjectDescriptorBase() {
-    private fun externalPackage(
-        contentRoot: String,
-        source: String?,
-        name: String,
-        targetName: String = name,
-        version: String = "0.0.1",
-        origin: PackageOrigin = PackageOrigin.DEPENDENCY,
-        libKind: LibKind = LibKind.LIB,
-        procMacroArtifact: CargoWorkspaceData.ProcMacroArtifact? = null,
-    ): Package {
-        return Package(
-            id = "$name $version",
-            contentRootUrl = contentRoot,
-            name = name,
-            version = version,
-            targets = listOf(
-                // don't use `FileUtil.join` here because it uses `File.separator`
-                // which is system dependent although all other code uses `/` as separator
-                Target(source?.let { "$contentRoot/$it" } ?: "", targetName,
-                    TargetKind.Lib(libKind), Edition.EDITION_2015, doctest = true, requiredFeatures = emptyList())
-            ),
-            source = source,
-            origin = origin,
-            edition = Edition.EDITION_2015,
-            features = emptyMap(),
-            enabledFeatures = emptySet(),
-            cfgOptions = CfgOptions.EMPTY,
-            env = emptyMap(),
-            outDirUrl = null,
-            procMacroArtifact = procMacroArtifact
-        )
-    }
 
     override fun setUp(fixture: CodeInsightTestFixture) {
         val root = fixture.findFileInTempDir(".")!!
@@ -395,6 +396,26 @@ object WithDependencyRustProjectDescriptor : RustProjectDescriptorBase() {
     }
 }
 
+private class WithStdlibLikeDependencyRustProjectDescriptor : RustProjectDescriptorBase() {
+    override fun testCargoProject(module: Module, contentRoot: String): CargoWorkspace {
+        val packages = listOf(
+            testCargoPackage(contentRoot),
+            externalPackage("$contentRoot/core", "lib.rs", "core"),
+            externalPackage("$contentRoot/alloc", "lib.rs", "alloc"),
+            externalPackage("$contentRoot/std", "lib.rs", "std")
+        )
+        return CargoWorkspace.deserialize(
+            Paths.get("${Urls.newFromIdea(contentRoot).path}/workspace/Cargo.toml"),
+            CargoWorkspaceData(packages, emptyMap(), emptyMap(), contentRoot), CfgOptions.DEFAULT)
+    }
+}
+
+/**
+ * Provides `core`, `alloc` and `std` workspace dependencies.
+ * It's supposed to be used to check how the plugin works with dependencies that have the same name as stdlib packages
+ */
+object WithStdlibAndStdlibLikeDependencyRustProjectDescriptor : WithRustup(WithStdlibLikeDependencyRustProjectDescriptor())
+
 private fun RsToolchainBase.getRustcInfo(): RustcInfo? {
     val rustc = rustc()
     val sysroot = rustc.getSysroot(Paths.get(".")) ?: return null

src/test/kotlin/org/rust/lang/core/resolve/RsStdlibResolveTest.kt

@@ -410,7 +410,12 @@ class RsStdlibResolveTest : RsResolveTestBase() {
         }
     """)
 
-    fun `test resolve derive traits`() {
+    fun `test resolve derive traits`() = resolveDeriveTraits()
+
+    @ProjectDescriptor(WithStdlibAndStdlibLikeDependencyRustProjectDescriptor::class)
+    fun `test resolve derive traits with stdlib-like dependencies`() = resolveDeriveTraits()
+
+    private fun resolveDeriveTraits() {
         val traitToPath = mapOf(
             "std::marker::Clone" to "clone.rs",
             "std::marker::Copy" to "marker.rs",

src/test/kotlin/org/rust/lang/core/type/RsStdlibExpressionTypeInferenceTest.kt

@@ -5,10 +5,7 @@
 
 package org.rust.lang.core.type
 
-import org.rust.ExpandMacros
-import org.rust.MockEdition
-import org.rust.ProjectDescriptor
-import org.rust.WithStdlibRustProjectDescriptor
+import org.rust.*
 import org.rust.cargo.project.workspace.CargoWorkspace
 import org.rust.lang.core.macros.MacroExpansionScope
 import org.rust.lang.core.psi.ext.*
@@ -126,6 +123,16 @@ class RsStdlibExpressionTypeInferenceTest : RsTypificationTestBase() {
         }
     """)
 
+    @ProjectDescriptor(WithStdlibAndStdlibLikeDependencyRustProjectDescriptor::class)
+    fun `test vec! with stdlib-like dependencies`() = stubOnlyTypeInfer("""
+    //- main.rs
+        fn main() {
+            let x = vec!(1, 2u16, 4, 8);
+            x;
+          //^ Vec<u16> | Vec<u16, Global>
+        }
+    """)
+
     fun `test vec! no_std`() = stubOnlyTypeInfer("""
     //- main.rs
         #![no_std]