<commit_before>package org.rust.lang.core.resolve.indexes

import com.intellij.openapi.project.Project
import com.intellij.psi.search.GlobalSearchScope
import com.intellij.psi.stubs.AbstractStubIndex
import com.intellij.psi.stubs.IndexSink
import com.intellij.psi.stubs.StubIndex
import com.intellij.psi.stubs.StubIndexKey
import com.intellij.util.io.KeyDescriptor
import org.rust.lang.core.psi.RsFunction
import org.rust.lang.core.psi.RsImplItem
import org.rust.lang.core.psi.ext.isAssocFn
import org.rust.lang.core.psi.ext.resolveToTrait
import org.rust.lang.core.stubs.RsFileStub
import org.rust.lang.core.stubs.RsImplItemStub
import org.rust.lang.core.types.RustType
import org.rust.lang.core.types.RustTypeFingerprint
import org.rust.lang.core.types.type

class RsImplIndex : AbstractStubIndex<RustTypeFingerprint, RsImplItem>() {
    override fun getVersion(): Int = RsFileStub.Type.stubVersion
    override fun getKey(): StubIndexKey<RustTypeFingerprint, RsImplItem> = KEY
    override fun getKeyDescriptor(): KeyDescriptor<RustTypeFingerprint> = RustTypeFingerprint.KeyDescriptor

    companion object {
        fun findMethodsFor(target: RustType, project: Project): Collection<RsFunction> =
            findMethodsAndAssociatedFunctionsFor(target, project)
                .filter { !it.isAssocFn }

        fun findMethodsAndAssociatedFunctionsFor(target: RustType, project: Project): Collection<RsFunction> =
            findImplsFor(target, project)
                .flatMap { it.allMethodsAndAssocFunctions }

        fun findImplsFor(target: RustType, project: Project): Collection<RsImplItem> {
            val fingerprint = RustTypeFingerprint.create(target)
                ?: return emptyList()

            return StubIndex.getElements(
                KEY,
                fingerprint,
                project,
                GlobalSearchScope.allScope(project),
                RsImplItem::class.java
            ).filter { impl ->
                val ty = impl.typeReference?.type
                // Addition class check is a temporal solution to filter impls for type parameter
                // with the same name
                // struct S; impl<S: Tr1> Tr2 for S {}
                ty != null && ty.javaClass == target.javaClass && ty.canUnifyWith(target, project)
            }
        }

        fun index(stub: RsImplItemStub, sink: IndexSink) {
            val type = stub.psi.typeReference ?: return
            val key = RustTypeFingerprint.create(type)
            if (key != null) {
                sink.occurrence(KEY, key)
            }
        }

        private val KEY: StubIndexKey<RustTypeFingerprint, RsImplItem> =
            StubIndexKey.createIndexKey("org.rust.lang.core.stubs.index.RustImplIndex.TraitImpls")
    }
}

private val RsImplItem.allMethodsAndAssocFunctions: Collection<RsFunction> get() {
    val directlyImplemented = functionList.map { it.name }.toSet()
    val defaulted = traitRef?.resolveToTrait?.functionList.orEmpty().asSequence().filter {
        it.name !in directlyImplemented
    }

    return functionList + defaulted
}<commit_msg>TY: cache impls per type<commit_after>package org.rust.lang.core.resolve.indexes

import com.intellij.openapi.project.Project
import com.intellij.psi.search.GlobalSearchScope
import com.intellij.psi.stubs.AbstractStubIndex
import com.intellij.psi.stubs.IndexSink
import com.intellij.psi.stubs.StubIndex
import com.intellij.psi.stubs.StubIndexKey
import com.intellij.psi.util.CachedValueProvider
import com.intellij.psi.util.CachedValuesManager
import com.intellij.psi.util.PsiModificationTracker
import com.intellij.util.containers.ContainerUtil
import com.intellij.util.io.KeyDescriptor
import org.rust.lang.core.psi.RsFunction
import org.rust.lang.core.psi.RsImplItem
import org.rust.lang.core.psi.ext.isAssocFn
import org.rust.lang.core.psi.ext.resolveToTrait
import org.rust.lang.core.stubs.RsFileStub
import org.rust.lang.core.stubs.RsImplItemStub
import org.rust.lang.core.types.RustType
import org.rust.lang.core.types.RustTypeFingerprint
import org.rust.lang.core.types.type
import java.util.concurrent.ConcurrentMap

class RsImplIndex : AbstractStubIndex<RustTypeFingerprint, RsImplItem>() {
    override fun getVersion(): Int = RsFileStub.Type.stubVersion
    override fun getKey(): StubIndexKey<RustTypeFingerprint, RsImplItem> = KEY
    override fun getKeyDescriptor(): KeyDescriptor<RustTypeFingerprint> = RustTypeFingerprint.KeyDescriptor

    companion object {
        fun findMethodsFor(target: RustType, project: Project): Collection<RsFunction> =
            findMethodsAndAssociatedFunctionsFor(target, project)
                .filter { !it.isAssocFn }

        fun findMethodsAndAssociatedFunctionsFor(target: RustType, project: Project): Collection<RsFunction> =
            findImplsFor(target, project)
                .flatMap { it.allMethodsAndAssocFunctions }

        fun findImplsFor(target: RustType, project: Project): Collection<RsImplItem> =
            project.implsCache.getOrPut(target) { doFindImplsFor(target, project) }

        private fun doFindImplsFor(target: RustType, project: Project): Collection<RsImplItem> {
            val fingerprint = RustTypeFingerprint.create(target)
                ?: return emptyList()

            return StubIndex.getElements(
                KEY,
                fingerprint,
                project,
                GlobalSearchScope.allScope(project),
                RsImplItem::class.java
            ).filter { impl ->
                val ty = impl.typeReference?.type
                // Addition class check is a temporal solution to filter impls for type parameter
                // with the same name
                // struct S; impl<S: Tr1> Tr2 for S {}
                ty != null && ty.javaClass == target.javaClass && ty.canUnifyWith(target, project)
            }
        }

        fun index(stub: RsImplItemStub, sink: IndexSink) {
            val type = stub.psi.typeReference ?: return
            val key = RustTypeFingerprint.create(type)
            if (key != null) {
                sink.occurrence(KEY, key)
            }
        }

        private val KEY: StubIndexKey<RustTypeFingerprint, RsImplItem> =
            StubIndexKey.createIndexKey("org.rust.lang.core.stubs.index.RustImplIndex.TraitImpls")
    }
}

private val RsImplItem.allMethodsAndAssocFunctions: Collection<RsFunction> get() {
    val directlyImplemented = functionList.map { it.name }.toSet()
    val defaulted = traitRef?.resolveToTrait?.functionList.orEmpty().asSequence().filter {
        it.name !in directlyImplemented
    }

    return functionList + defaulted
}

private val Project.implsCache: ConcurrentMap<RustType, Collection<RsImplItem>>
    get() = CachedValuesManager.getManager(this)
        .getCachedValue(this, {
            CachedValueProvider.Result.create(
                ContainerUtil.newConcurrentMap<RustType, Collection<RsImplItem>>(),
                PsiModificationTracker.MODIFICATION_COUNT
            )
        })
{'fraction_non_alphanumeric': 0.09212587213441549, 'fraction_numerical': 0.0005695571693008686, 'mean_word_length': 2.9284116331096195, 'pattern_counts': {'":': 0, '<': 26, '<?xml version=': 0, '>': 28, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6647425', 'n_tokens_mistral': 2072, 'n_tokens_neox': 2013, 'n_words': 378}