diff --git a/platform/ml-embeddings/src/com/intellij/platform/ml/embeddings/indexer/embeddingIndexes.kt b/platform/ml-embeddings/src/com/intellij/platform/ml/embeddings/indexer/embeddingIndexes.kt index 36c2d19756ce..1a567857eeb5 100644 --- a/platform/ml-embeddings/src/com/intellij/platform/ml/embeddings/indexer/embeddingIndexes.kt +++ b/platform/ml-embeddings/src/com/intellij/platform/ml/embeddings/indexer/embeddingIndexes.kt @@ -3,6 +3,8 @@ package com.intellij.platform.ml.embeddings.indexer import com.intellij.openapi.fileTypes.FileType import com.intellij.openapi.vfs.VirtualFile +import com.intellij.platform.ml.embeddings.indexer.entities.IndexableEntity +import com.intellij.platform.ml.embeddings.indexer.entities.IndexableFile import com.intellij.util.indexing.* import com.intellij.util.io.DataExternalizer import com.intellij.util.io.DataInputOutputUtil @@ -14,24 +16,19 @@ internal val FILE_NAME_EMBEDDING_INDEX_NAME = ID.create("F internal val CLASS_NAME_EMBEDDING_INDEX_NAME = ID.create("ClassNameEmbeddingIndex") internal val SYMBOL_NAME_EMBEDDING_INDEX_NAME = ID.create("SymbolNameEmbeddingIndex") -internal class EmbeddingKey(val textHashCode: Int) { - override fun hashCode(): Int = textHashCode.hashCode() - override fun equals(other: Any?): Boolean = other is EmbeddingKey && other.textHashCode == textHashCode - +@JvmInline +internal value class EmbeddingKey(val textHashCode: Int) { fun toLong(fileId: Int): Long { // https://stackoverflow.com/a/12772968 return (fileId.toLong() shl 32) or (0xffffffffL and textHashCode.toLong()) } } -@JvmInline -value class IndexingItem(val text: String) - internal class FileNameEmbeddingIndex : BaseEmbeddingIndex() { override fun getName(): ID = FILE_NAME_EMBEDDING_INDEX_NAME override fun getVersion(): Int = 1 - override fun index(inputData: FileContent): List { - return listOf(IndexingItem(inputData.file.name)) + override fun index(inputData: FileContent): List { + return listOf(IndexableFile(inputData.file)) } override fun getInputFilter(): FileBasedIndex.InputFilter { @@ -52,8 +49,8 @@ internal class ClassNameEmbeddingIndex : PsiBaseEmbeddingIndex() { override fun getName(): ID = CLASS_NAME_EMBEDDING_INDEX_NAME override fun getVersion(): Int = 2 - override fun index(inputData: FileContent): List { - return ClassesProvider.extractClasses(inputData).map { IndexingItem(it.id.id) } + override fun index(inputData: FileContent): List { + return ClassesProvider.extractClasses(inputData) } } @@ -63,8 +60,8 @@ internal class SymbolNameEmbeddingIndex : PsiBaseEmbeddingIndex() { override fun getName(): ID = SYMBOL_NAME_EMBEDDING_INDEX_NAME override fun getVersion(): Int = 2 - override fun index(inputData: FileContent): List { - return SymbolsProvider.extractSymbols(inputData).map { IndexingItem(it.id.id) } + override fun index(inputData: FileContent): List { + return SymbolsProvider.extractSymbols(inputData) } } @@ -73,14 +70,17 @@ internal abstract class BaseEmbeddingIndex : FileBasedIndexExtension { return DataIndexer { inputData -> - index(inputData).associate { item -> - val textHashcode = item.text.hashCode() - EmbeddingKey(textHashcode) to item.text + index(inputData).associate { entity -> + // Hash code should be calculated from indexable representation, not identifier + val textHashcode = entity.indexableRepresentation.hashCode() + // TODO: link to indexable representation (e.g. TextRange in file) should be preserved + // in index values besides id because we calculate vectors based on it + EmbeddingKey(textHashcode) to entity.id.id } } } - abstract fun index(context: FileContent): List + abstract fun index(context: FileContent): List override fun getKeyDescriptor(): KeyDescriptor { return object : KeyDescriptor { diff --git a/platform/ml-embeddings/src/com/intellij/platform/ml/embeddings/indexer/searcher/index/IndexBasedEmbeddingEntitiesIndexer.kt b/platform/ml-embeddings/src/com/intellij/platform/ml/embeddings/indexer/searcher/index/IndexBasedEmbeddingEntitiesIndexer.kt index 50ea6297df3e..577889b86e54 100644 --- a/platform/ml-embeddings/src/com/intellij/platform/ml/embeddings/indexer/searcher/index/IndexBasedEmbeddingEntitiesIndexer.kt +++ b/platform/ml-embeddings/src/com/intellij/platform/ml/embeddings/indexer/searcher/index/IndexBasedEmbeddingEntitiesIndexer.kt @@ -63,6 +63,8 @@ internal class IndexBasedEmbeddingEntitiesIndexer(cs: CoroutineScope) : Embeddin val fileIdsAndNames = smartReadAction(project) { chunk.mapNotNull { key -> var result: Pair? = null + // TODO: besides entity name (id) we should access indexable representation here + // (e.g. via TextRange in index values) because vectors are calculated from it fileBasedIndex.processValues(indexId, key, null, { virtualFile, name -> if (virtualFile is VirtualFileWithId) { result = Pair(key.toLong(virtualFile.id), name)