TY&RES: ignore impls from non-dependency crates

src/main/kotlin/org/rust/lang/core/crate/Crate.kt

@@ -112,3 +112,6 @@ interface Crate : UserDataHolderEx {
 
 fun Crate.findDependency(normName: String): Crate? =
     dependencies.find { it.normName == normName }?.crate
+
+fun Crate.hasTransitiveDependencyOrSelf(other: Crate): Boolean =
+    other == this || other in flatDependencies

src/main/kotlin/org/rust/lang/core/psi/RsFile.kt

@@ -42,6 +42,7 @@ import org.rust.lang.core.macros.macroExpansionManager
 import org.rust.lang.core.psi.ext.*
 import org.rust.lang.core.resolve.DEFAULT_RECURSION_LIMIT
 import org.rust.lang.core.resolve.ref.RsReference
+import org.rust.lang.core.resolve2.ModData
 import org.rust.lang.core.resolve2.findModDataFor
 import org.rust.lang.core.resolve2.toRsMod
 import org.rust.lang.core.stubs.RsFileStub
@@ -75,6 +76,7 @@ class RsFile(
     val cargoProject: CargoProject? get() = cachedData.cargoProject
     val cargoWorkspace: CargoWorkspace? get() = cachedData.cargoWorkspace
     val crate: Crate? get() = cachedData.crate
+    val crates: List<Crate> get() = cachedData.crates
     override val crateRoot: RsMod? get() = cachedData.crateRoot
     val isDeeplyEnabledByCfg: Boolean get() = cachedData.isDeeplyEnabledByCfg
     val isIncludedByIncludeMacro: Boolean get() = cachedData.isIncludedByIncludeMacro
@@ -114,20 +116,25 @@ class RsFile(
                 crate.cargoWorkspace,
                 crate.rootMod,
                 crate,
+                crates = listOf(crate),
                 isDeeplyEnabledByCfg = true, // Macros are ony expanded in cfg-enabled mods
                 isIncludedByIncludeMacro = false, // An expansion file obviously can't be included
             )
         }
 
         // Note: `this` file can be not a module (can be included with `include!()` macro)
-        val modData = findModDataFor(this)
+        val allModData = findModDataFor(this)
+        val modData = allModData.pickSingleModData()
         if (modData != null) {
-            val crate = project.crateGraph.findCrateById(modData.crate) ?: return EMPTY_CACHED_DATA
+            val crateGraph = project.crateGraph
+            val crates = allModData.mapNotNull { crateGraph.findCrateById(it.crate) }
+            val crate = crates.find { it.id == modData.crate } ?: return EMPTY_CACHED_DATA
             return CachedData(
                 crate.cargoProject,
                 crate.cargoWorkspace,
                 crate.rootMod,
                 crate,
+                crates,
                 modData.isDeeplyEnabledByCfg,
                 isIncludedByIncludeMacro = virtualFile is VirtualFileWithId && virtualFile.id != modData.fileId,
             )
@@ -173,7 +180,7 @@ class RsFile(
 
     override val `super`: RsMod?
         get() {
-            val modData = findModDataFor(this) ?: return null
+            val modData = findModDataFor(this).pickSingleModData() ?: return null
             val parenModData = modData.parent ?: return null
             return parenModData.toRsMod(project).firstOrNull()
         }
@@ -296,12 +303,23 @@ private data class CachedData(
     val cargoWorkspace: CargoWorkspace? = null,
     val crateRoot: RsFile? = null,
     val crate: Crate? = null,
+    val crates: List<Crate> = emptyList(),
     val isDeeplyEnabledByCfg: Boolean = true,
     val isIncludedByIncludeMacro: Boolean = false,
 )
 
 private val EMPTY_CACHED_DATA: CachedData = CachedData()
 
+// A rust file can be included in multiple places, but currently IntelliJ Rust works properly only with one
+// inclusion point, so we have to choose one
+private fun List<ModData>.pickSingleModData(): ModData? {
+    if (isEmpty()) return null
+    singleOrNull()?.let { return it }
+
+    // If there are still multiple options, choose one deterministically
+    return minByOrNull { it.crate }
+}
+
 private fun VirtualFile.getInjectedFromIfDoctestInjection(project: Project): RsFile? {
     if (!isDoctestInjection(project)) return null
     return ((this as? VirtualFileWindow)?.delegate?.toPsiFile(project) as? RsFile)
@@ -324,14 +342,14 @@ private val CACHED_DATA_KEY: Key<CachedValue<CachedData>> = Key.create("CACHED_D
 val RsElement.isValidProjectMember: Boolean
     get() = isValidProjectMemberAndContainingCrate.first
 
-val RsElement.isValidProjectMemberAndContainingCrate: Pair<Boolean, Crate?>
+val RsElement.isValidProjectMemberAndContainingCrate: Triple<Boolean, Crate?, List<Crate>>
     get() {
-        val file = containingRsFileSkippingCodeFragments ?: return true to null
-        if (!file.isDeeplyEnabledByCfg) return false to null
-        val crate = file.crate ?: return false to null
-        if (!existsAfterExpansion(crate)) return false to null
+        val file = containingRsFileSkippingCodeFragments ?: return Triple(true, null, emptyList())
+        if (!file.isDeeplyEnabledByCfg) return Triple(false, null, emptyList())
+        val crate = file.crate ?: return Triple(false, null, emptyList())
+        if (!existsAfterExpansion(crate)) return Triple(false, null, emptyList())
 
-        return true to crate
+        return Triple(true, crate, file.crates)
     }
 
 /** Usually used to filter out test/bench non-workspace crates */

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

@@ -8,12 +8,11 @@ package org.rust.lang.core.resolve
 import com.intellij.openapi.project.Project
 import com.intellij.util.SmartList
 import gnu.trove.THashMap
-import org.rust.cargo.project.model.CargoProject
+import org.rust.lang.core.crate.Crate
+import org.rust.lang.core.crate.hasTransitiveDependencyOrSelf
 import org.rust.lang.core.psi.*
 import org.rust.lang.core.psi.ext.*
 import org.rust.lang.core.resolve.SelectionCandidate.*
-import org.rust.lang.core.resolve.indexes.RsImplIndex
-import org.rust.lang.core.resolve.indexes.RsTypeAliasIndex
 import org.rust.lang.core.types.*
 import org.rust.lang.core.types.consts.CtConstParameter
 import org.rust.lang.core.types.consts.CtInferVar
@@ -22,10 +21,8 @@ import org.rust.lang.core.types.infer.*
 import org.rust.lang.core.types.infer.TypeInferenceMarks.WinnowParamCandidateLoses
 import org.rust.lang.core.types.infer.TypeInferenceMarks.WinnowParamCandidateWins
 import org.rust.lang.core.types.ty.*
-import org.rust.lang.utils.CargoProjectCache
 import org.rust.openapiext.hitOnTrue
 import org.rust.openapiext.testAssert
-import org.rust.stdext.Cache
 import org.rust.stdext.buildList
 import org.rust.stdext.exhaustive
 import org.rust.stdext.swapRemoveAt
@@ -183,7 +180,7 @@ data class ParamEnv(val callerBounds: List<TraitRef>) {
                 1 -> return ParamEnv(rawBounds)
             }
 
-            val lookup = ImplLookup(decl.project, decl.cargoProject, decl.knownItems, ParamEnv(rawBounds))
+            val lookup = ImplLookup(decl.project, decl.containingCrate, decl.knownItems, ParamEnv(rawBounds))
             val ctx = lookup.ctx
             val bounds2 = rawBounds.map {
                 val (bound, obligations) = ctx.normalizeAssociatedTypesIn(it)
@@ -199,27 +196,14 @@ data class ParamEnv(val callerBounds: List<TraitRef>) {
 
 class ImplLookup(
     private val project: Project,
-    cargoProject: CargoProject?,
+    private val containingCrate: Crate?,
     val items: KnownItems,
     private val paramEnv: ParamEnv = ParamEnv.EMPTY
 ) {
     // Non-concurrent HashMap and lazy(NONE) are safe here because this class isn't shared between threads
-    private val traitSelectionCache: Cache<TraitRef, SelectionResult<SelectionCandidate>> =
-        if (paramEnv.isEmpty() && cargoProject != null) {
-            cargoProjectGlobalTraitSelectionCache.getCache(cargoProject)
-        } else {
-            // function-local cache is used when [paramEnv] is not empty, i.e. if there are trait bounds
-            // that affect trait selection
-            Cache.new()
-        }
-    private val findImplsAndTraitsCache: Cache<Ty, List<TraitImplSource>> =
-        if (cargoProject != null) {
-            cargoProjectGlobalFindImplsAndTraitsCache.getCache(cargoProject)
-        } else {
-            Cache.new()
-        }
-    private val implIndexCache: Cache<TyFingerprint, List<RsCachedImplItem>> = Cache.new()
-    private val typeAliasIndexCache: Cache<TyFingerprint, List<RsCachedTypeAlias>> = Cache.new()
+    private val traitSelectionCache: MutableMap<TraitRef, SelectionResult<SelectionCandidate>> = hashMapOf()
+    private val findImplsAndTraitsCache: MutableMap<Ty, List<TraitImplSource>> = hashMapOf()
+    private val indexCache = RsImplIndexAndTypeAliasCache.getInstance(project)
     private val fnTraits = listOfNotNull(items.Fn, items.FnMut, items.FnOnce)
     private val fnOnceOutput: RsTypeAlias? by lazy(NONE) {
         val trait = items.FnOnce ?: return@lazy null
@@ -321,7 +305,7 @@ class ImplLookup(
     }
 
     private fun findExplicitImplsWithoutAliases(selfTy: Ty, tyf: TyFingerprint, processor: RsProcessor<RsCachedImplItem>): Boolean {
-        val impls = implIndexCache.getOrPut(tyf) { RsImplIndex.findPotentialImpls(project, tyf) }
+        val impls = findPotentialImpls(tyf)
         return impls.any { cachedImpl ->
             if (cachedImpl.isNegativeImpl) return@any false
             val (type, generics, constGenerics) = cachedImpl.typeAndGenerics ?: return@any false
@@ -338,9 +322,7 @@ class ImplLookup(
         if (fingerprint != null) {
             val set = mutableSetOf(fingerprint)
             if (processor(fingerprint)) return true
-            val aliases = typeAliasIndexCache.getOrPut(fingerprint) {
-                RsTypeAliasIndex.findPotentialAliases(project, fingerprint)
-            }
+            val aliases = findPotentialAliases(fingerprint)
             val result = aliases.any {
                 val name = it.name ?: return@any false
                 val aliasFingerprint = TyFingerprint(name)
@@ -355,6 +337,19 @@ class ImplLookup(
         return processor(TyFingerprint.TYPE_PARAMETER_OR_MACRO_FINGERPRINT)
     }
 
+    private fun findPotentialImpls(tyf: TyFingerprint): Sequence<RsCachedImplItem> =
+        indexCache.findPotentialImpls(tyf)
+            .asSequence()
+            .filter { useImplsFromCrate(it.containingCrates) }
+
+    private fun findPotentialAliases(tyf: TyFingerprint) =
+        indexCache.findPotentialAliases(tyf)
+            .asSequence()
+            .filter { useImplsFromCrate(it.containingCrates) }
+
+    private fun useImplsFromCrate(crates: List<Crate>): Boolean =
+        containingCrate == null || crates.any { containingCrate.hasTransitiveDependencyOrSelf(it) }
+
     private fun canCombineTypes(
         ty1: Ty,
         ty2: Ty,
@@ -381,10 +376,8 @@ class ImplLookup(
     }
 
     /** return impls for a generic type `impl<T> Trait for T {}` */
-    private fun findBlanketImpls(): List<RsCachedImplItem> {
-        return implIndexCache.getOrPut(TyFingerprint.TYPE_PARAMETER_OR_MACRO_FINGERPRINT) {
-            RsImplIndex.findPotentialImpls(project, TyFingerprint.TYPE_PARAMETER_OR_MACRO_FINGERPRINT)
-        }
+    private fun findBlanketImpls(): Sequence<RsCachedImplItem> {
+        return findPotentialImpls(TyFingerprint.TYPE_PARAMETER_OR_MACRO_FINGERPRINT)
     }
 
     /**
@@ -767,7 +760,7 @@ class ImplLookup(
     }
 
     private fun assembleImplCandidatesWithoutAliases(ref: TraitRef, tyf: TyFingerprint, processor: RsProcessor<SelectionCandidate>): Boolean {
-        val impls = implIndexCache.getOrPut(tyf) { RsImplIndex.findPotentialImpls(project, tyf) }
+        val impls = findPotentialImpls(tyf)
         return impls.any {
             val candidate = it.trySelectCandidate(ref)
             candidate != null && processor(candidate)
@@ -1288,14 +1281,8 @@ class ImplLookup(
             } else {
                 ParamEnv.EMPTY
             }
-            return ImplLookup(psi.project, psi.cargoProject, psi.knownItems, paramEnv)
+            return ImplLookup(psi.project, psi.containingCrate, psi.knownItems, paramEnv)
         }
-
-        private val cargoProjectGlobalFindImplsAndTraitsCache =
-            CargoProjectCache<Ty, List<TraitImplSource>>("cargoProjectGlobalFindImplsAndTraitsCache")
-
-        private val cargoProjectGlobalTraitSelectionCache =
-            CargoProjectCache<TraitRef, SelectionResult<SelectionCandidate>>("cargoProjectGlobalTraitSelectionCache")
     }
 }
 

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

@@ -32,12 +32,14 @@ class RsCachedImplItem(
 ) {
     private val traitRef: RsTraitRef? = impl.traitRef
     val containingCrate: Crate?
+    val containingCrates: List<Crate>
     val isValid: Boolean
     val isNegativeImpl: Boolean = impl.isNegativeImpl
 
     init {
-        val (isValid, crate) = impl.isValidProjectMemberAndContainingCrate
+        val (isValid, crate, crates) = impl.isValidProjectMemberAndContainingCrate
         this.containingCrate = crate
+        this.containingCrates = crates
         this.isValid = isValid && !impl.isReservationImpl
     }
 

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

@@ -30,10 +30,12 @@ class RsCachedTypeAlias(
     val isFreeAndValid: Boolean
 
     val containingCrate: Crate?
+    val containingCrates: List<Crate>
 
     init {
-        val (isValid, crate) = alias.isValidProjectMemberAndContainingCrate
+        val (isValid, crate, crates) = alias.isValidProjectMemberAndContainingCrate
         this.containingCrate = crate
+        this.containingCrates = crates
         this.isFreeAndValid = isValid
             && name != null
             && alias.owner is RsAbstractableOwner.Free

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

@@ -0,0 +1,95 @@
+/*
+ * Use of this source code is governed by the MIT license that can be
+ * found in the LICENSE file.
+ */
+
+package org.rust.lang.core.resolve
+
+import com.intellij.openapi.Disposable
+import com.intellij.openapi.components.Service
+import com.intellij.openapi.components.service
+import com.intellij.openapi.project.Project
+import com.intellij.psi.PsiElement
+import com.intellij.psi.PsiFile
+import com.intellij.util.containers.ContainerUtil
+import org.rust.lang.core.psi.RustStructureChangeListener
+import org.rust.lang.core.psi.rustPsiManager
+import org.rust.lang.core.resolve.indexes.RsImplIndex
+import org.rust.lang.core.resolve.indexes.RsTypeAliasIndex
+import org.rust.lang.core.types.TyFingerprint
+import java.util.concurrent.ConcurrentMap
+import java.util.concurrent.atomic.AtomicReference
+
+@Service
+class RsImplIndexAndTypeAliasCache(private val project: Project) : Disposable {
+    // strong key -> soft value maps
+    private val _implIndexCache: AtomicReference<ConcurrentMap<TyFingerprint, List<RsCachedImplItem>>?> = AtomicReference(null)
+    private val _typeAliasIndexCache: AtomicReference<ConcurrentMap<TyFingerprint, List<RsCachedTypeAlias>>?> = AtomicReference(null)
+
+    private val implIndexCache: ConcurrentMap<TyFingerprint, List<RsCachedImplItem>>
+        get() = _implIndexCache.getOrCreateMap()
+
+    private val typeAliasIndexCache: ConcurrentMap<TyFingerprint, List<RsCachedTypeAlias>>
+        get() = _typeAliasIndexCache.getOrCreateMap()
+
+    /**
+     * This map is actually used is a [Set] (the value is always [placeholder]).
+     * The only purpose of this set is holding links to [PsiFile]s, so retain them in memory.
+     * Without this set [PsiFile]s are retained only by [java.lang.ref.WeakReference], hence they are
+     * quickly collected by GC and then further index lookups work slower.
+     *
+     * Note: keys in this map are referenced via [java.lang.ref.SoftReference], so they're also
+     * can be collected by GC, hence there isn't a memory leak
+     */
+    private val usedPsiFiles: ConcurrentMap<PsiFile, Any> = ContainerUtil.createConcurrentSoftMap()
+    private val placeholder = Any()
+
+    init {
+        val rustPsiManager = project.rustPsiManager
+        val connection = project.messageBus.connect(this)
+        rustPsiManager.subscribeRustStructureChange(connection, object : RustStructureChangeListener {
+            override fun rustStructureChanged(file: PsiFile?, changedElement: PsiElement?) {
+                _implIndexCache.getAndSet(null)
+                _typeAliasIndexCache.getAndSet(null)
+            }
+        })
+    }
+
+    fun findPotentialImpls(tyf: TyFingerprint): List<RsCachedImplItem> {
+        return implIndexCache.getOrPut(tyf) {
+            RsImplIndex.findPotentialImpls(project, tyf).filter {
+                retainPsi(it.impl.containingFile)
+                it.isValid
+            }
+        }
+    }
+
+    fun findPotentialAliases(tyf: TyFingerprint): List<RsCachedTypeAlias> {
+        return typeAliasIndexCache.getOrPut(tyf) {
+            RsTypeAliasIndex.findPotentialAliases(project, tyf).filter {
+                retainPsi(it.alias.containingFile)
+                it.isFreeAndValid
+            }
+        }
+    }
+
+    private fun retainPsi(containingFile: PsiFile) {
+        usedPsiFiles[containingFile] = placeholder
+    }
+
+    override fun dispose() {}
+
+    companion object {
+        fun getInstance(project: Project): RsImplIndexAndTypeAliasCache =
+            project.service()
+
+        @JvmStatic
+        private fun <T : Any> AtomicReference<ConcurrentMap<TyFingerprint, T>?>.getOrCreateMap(): ConcurrentMap<TyFingerprint, T> {
+            while (true) {
+                get()?.let { return it }
+                val map = ContainerUtil.createConcurrentSoftValueMap<TyFingerprint, T>()
+                if (compareAndSet(null, map)) return map
+            }
+        }
+    }
+}

src/main/kotlin/org/rust/lang/core/resolve/indexes/RsImplIndex.kt

@@ -39,10 +39,7 @@ class RsImplIndex : AbstractStubIndex<TyFingerprint, RsImplItem>() {
 
             // Note that `getElements` is intentionally used with intermediate collection instead of
             // `StubIndex.processElements` in order to simplify profiling
-            return impls.mapNotNull { impl ->
-                val cachedImpl = RsCachedImplItem.forImpl(impl)
-                cachedImpl.takeIf { it.isValid }
-            }
+            return impls.map { RsCachedImplItem.forImpl(it) }
         }
 
         fun index(stub: RsImplItemStub, sink: IndexSink) {

src/main/kotlin/org/rust/lang/core/resolve/indexes/RsTypeAliasIndex.kt

@@ -35,7 +35,6 @@ class RsTypeAliasIndex : AbstractStubIndex<TyFingerprint, RsTypeAlias>() {
             // `StubIndex.processElements` in order to simplify profiling
             val aliases = getElements(KEY, tyf, project, RsWithMacrosProjectScope(project))
                 .map { RsCachedTypeAlias.forAlias(it) }
-                .filter { it.isFreeAndValid }
 
             // This is basically a hack to make some crates (winapi 0.2) work in a reasonable amount of time.
             // If the number of aliases exceeds the threshold, we prefer ones from stdlib and a workspace over