From b4ddea5d82acf8d974f536cd4d3ff9ff9d22d6e1 Mon Sep 17 00:00:00 2001 From: Alexey Kalina Date: Thu, 14 Jan 2021 11:52:13 +0300 Subject: [PATCH] [ml-completion-local] frequency storages refactoring GitOrigin-RevId: 64c61c774fdf73992de30ca59a19f30170128e0a --- .../FrequencyContextFeaturesProvider.kt | 16 +++++++------ .../FrequencyElementFeatureProvider.kt | 2 +- .../ml/local/models/api/LocalModel.kt | 1 + .../frequency/ClassesFrequencyLocalModel.kt | 2 ++ .../frequency/MethodsFrequencyLocalModel.kt | 2 ++ .../models/storage/ClassesFrequencyStorage.kt | 20 ++++++++++------ .../local/models/storage/LocalModelStorage.kt | 1 + .../models/storage/MethodsFrequencyStorage.kt | 20 ++++++++++------ .../ml/local/models/storage/StorageInfo.kt | 3 --- .../ml/local/models/storage/StorageUtil.kt | 23 ++++++++++--------- 10 files changed, 54 insertions(+), 36 deletions(-) delete mode 100644 plugins/completion-ml-ranking-local/src/com/intellij/completion/ml/local/models/storage/StorageInfo.kt diff --git a/plugins/completion-ml-ranking-local/src/com/intellij/completion/ml/local/features/FrequencyContextFeaturesProvider.kt b/plugins/completion-ml-ranking-local/src/com/intellij/completion/ml/local/features/FrequencyContextFeaturesProvider.kt index 35a6d7057d8a..646572ee3f26 100644 --- a/plugins/completion-ml-ranking-local/src/com/intellij/completion/ml/local/features/FrequencyContextFeaturesProvider.kt +++ b/plugins/completion-ml-ranking-local/src/com/intellij/completion/ml/local/features/FrequencyContextFeaturesProvider.kt @@ -24,21 +24,23 @@ class FrequencyContextFeaturesProvider : ContextFeatureProvider { override fun calculateFeatures(environment: CompletionEnvironment): MutableMap { val features = mutableMapOf() val project = environment.parameters.position.project - LocalModelsManager.getInstance(project).getModel()?.let { model -> + val methodsModel = LocalModelsManager.getInstance(project).getModel() + if (methodsModel != null && methodsModel.readyToUse()) { getReceiverClass(environment.parameters)?.let { cls -> LocalModelsUtil.getClassName(cls)?.let { environment.putUserData(RECEIVER_CLASS_NAME_KEY, it) - model.getMethodsByClass(it)?.let { frequencies -> + methodsModel.getMethodsByClass(it)?.let { frequencies -> environment.putUserData(RECEIVER_CLASS_FREQUENCIES_KEY, frequencies) } } } - features["total_methods"] = MLFeatureValue.numerical(model.totalMethodsCount()) - features["total_methods_usages"] = MLFeatureValue.numerical(model.totalMethodsUsages()) + features["total_methods"] = MLFeatureValue.numerical(methodsModel.totalMethodsCount()) + features["total_methods_usages"] = MLFeatureValue.numerical(methodsModel.totalMethodsUsages()) } - LocalModelsManager.getInstance(project).getModel()?.let { model -> - features["total_classes"] = MLFeatureValue.numerical(model.totalClassesCount()) - features["total_classes_usages"] = MLFeatureValue.numerical(model.totalClassesUsages()) + val classesModel = LocalModelsManager.getInstance(project).getModel() + if (classesModel != null && classesModel.readyToUse()) { + features["total_classes"] = MLFeatureValue.numerical(classesModel.totalClassesCount()) + features["total_classes_usages"] = MLFeatureValue.numerical(classesModel.totalClassesUsages()) } return features } diff --git a/plugins/completion-ml-ranking-local/src/com/intellij/completion/ml/local/features/FrequencyElementFeatureProvider.kt b/plugins/completion-ml-ranking-local/src/com/intellij/completion/ml/local/features/FrequencyElementFeatureProvider.kt index 770ed3cdd07c..a09620adc202 100644 --- a/plugins/completion-ml-ranking-local/src/com/intellij/completion/ml/local/features/FrequencyElementFeatureProvider.kt +++ b/plugins/completion-ml-ranking-local/src/com/intellij/completion/ml/local/features/FrequencyElementFeatureProvider.kt @@ -36,7 +36,7 @@ class FrequencyElementFeatureProvider : ElementFeatureProvider { } } val classesModel = LocalModelsManager.getInstance(location.project).getModel() - if (psi is PsiClass && classesModel != null) { + if (psi is PsiClass && classesModel != null && classesModel.readyToUse()) { LocalModelsUtil.getClassName(psi)?.let { className -> classesModel.getClass(className)?.let { features["absolute_class_frequency"] = MLFeatureValue.numerical(it) diff --git a/plugins/completion-ml-ranking-local/src/com/intellij/completion/ml/local/models/api/LocalModel.kt b/plugins/completion-ml-ranking-local/src/com/intellij/completion/ml/local/models/api/LocalModel.kt index 13afeb061a71..1a4e3de3fc54 100644 --- a/plugins/completion-ml-ranking-local/src/com/intellij/completion/ml/local/models/api/LocalModel.kt +++ b/plugins/completion-ml-ranking-local/src/com/intellij/completion/ml/local/models/api/LocalModel.kt @@ -6,4 +6,5 @@ interface LocalModel { fun fileVisitor(): PsiElementVisitor fun onStarted() fun onFinished() + fun readyToUse(): Boolean } \ No newline at end of file diff --git a/plugins/completion-ml-ranking-local/src/com/intellij/completion/ml/local/models/frequency/ClassesFrequencyLocalModel.kt b/plugins/completion-ml-ranking-local/src/com/intellij/completion/ml/local/models/frequency/ClassesFrequencyLocalModel.kt index d0285c099154..6296572c350d 100644 --- a/plugins/completion-ml-ranking-local/src/com/intellij/completion/ml/local/models/frequency/ClassesFrequencyLocalModel.kt +++ b/plugins/completion-ml-ranking-local/src/com/intellij/completion/ml/local/models/frequency/ClassesFrequencyLocalModel.kt @@ -87,4 +87,6 @@ class ClassesFrequencyLocalModel private constructor(private val storage: Classe override fun onFinished() { storage.setValid(true) } + + override fun readyToUse(): Boolean = storage.isValid() && !storage.isEmpty() } \ No newline at end of file diff --git a/plugins/completion-ml-ranking-local/src/com/intellij/completion/ml/local/models/frequency/MethodsFrequencyLocalModel.kt b/plugins/completion-ml-ranking-local/src/com/intellij/completion/ml/local/models/frequency/MethodsFrequencyLocalModel.kt index 40e9c89115c0..91db6af7df40 100644 --- a/plugins/completion-ml-ranking-local/src/com/intellij/completion/ml/local/models/frequency/MethodsFrequencyLocalModel.kt +++ b/plugins/completion-ml-ranking-local/src/com/intellij/completion/ml/local/models/frequency/MethodsFrequencyLocalModel.kt @@ -46,4 +46,6 @@ class MethodsFrequencyLocalModel private constructor(private val storage: Method override fun onFinished() { storage.setValid(true) } + + override fun readyToUse(): Boolean = storage.isValid() && !storage.isEmpty() } \ No newline at end of file diff --git a/plugins/completion-ml-ranking-local/src/com/intellij/completion/ml/local/models/storage/ClassesFrequencyStorage.kt b/plugins/completion-ml-ranking-local/src/com/intellij/completion/ml/local/models/storage/ClassesFrequencyStorage.kt index fea4ad17b793..040543567baa 100644 --- a/plugins/completion-ml-ranking-local/src/com/intellij/completion/ml/local/models/storage/ClassesFrequencyStorage.kt +++ b/plugins/completion-ml-ranking-local/src/com/intellij/completion/ml/local/models/storage/ClassesFrequencyStorage.kt @@ -1,5 +1,7 @@ package com.intellij.completion.ml.local.models.storage +import com.intellij.completion.ml.local.models.storage.StorageUtil.clear +import com.intellij.completion.ml.local.models.storage.StorageUtil.isEmpty import com.intellij.util.Processor import com.intellij.util.io.EnumeratorStringDescriptor import com.intellij.util.io.IntInlineKeyDescriptor @@ -14,13 +16,11 @@ class ClassesFrequencyStorage internal constructor(private val storageDirectory: fun getStorage(baseDirectory: Path): ClassesFrequencyStorage { val storageDirectory = baseDirectory.resolve(STORAGE_NAME) - val isValid = StorageUtil.prepareStorage(storageDirectory, VERSION) - val storage = ClassesFrequencyStorage(storageDirectory) - storage.setValid(isValid) - return storage + StorageUtil.prepareStorage(storageDirectory, VERSION) + return ClassesFrequencyStorage(storageDirectory) } } - private var isValid: Boolean = false + private var isValid: Boolean = true private val persistentStorage = PersistentHashMap(storageDirectory.resolve(STORAGE_NAME), EnumeratorStringDescriptor(), IntInlineKeyDescriptor()) @@ -32,21 +32,27 @@ class ClassesFrequencyStorage internal constructor(private val storageDirectory: @Volatile var totalClassesUsages = 0 private set + init { + toMemoryStorage() + } + override fun name(): String = STORAGE_NAME override fun version(): Int = VERSION override fun isValid(): Boolean = isValid + override fun isEmpty(): Boolean = persistentStorage.isEmpty() + override fun setValid(isValid: Boolean) { if (isValid) { toMemoryStorage() } else { - StorageUtil.clearMap(persistentStorage) + persistentStorage.clear() memoryStorage.clear() } this.isValid = isValid - StorageUtil.saveInfo(StorageInfo(VERSION, isValid), storageDirectory) + StorageUtil.saveInfo(VERSION, isValid, storageDirectory) } @Synchronized diff --git a/plugins/completion-ml-ranking-local/src/com/intellij/completion/ml/local/models/storage/LocalModelStorage.kt b/plugins/completion-ml-ranking-local/src/com/intellij/completion/ml/local/models/storage/LocalModelStorage.kt index 6b75a677512a..0503ccdbf451 100644 --- a/plugins/completion-ml-ranking-local/src/com/intellij/completion/ml/local/models/storage/LocalModelStorage.kt +++ b/plugins/completion-ml-ranking-local/src/com/intellij/completion/ml/local/models/storage/LocalModelStorage.kt @@ -4,5 +4,6 @@ interface LocalModelStorage { fun version(): Int fun name(): String fun isValid(): Boolean + fun isEmpty(): Boolean fun setValid(isValid: Boolean) } \ No newline at end of file diff --git a/plugins/completion-ml-ranking-local/src/com/intellij/completion/ml/local/models/storage/MethodsFrequencyStorage.kt b/plugins/completion-ml-ranking-local/src/com/intellij/completion/ml/local/models/storage/MethodsFrequencyStorage.kt index 8d038f566d87..be5ba3b12c54 100644 --- a/plugins/completion-ml-ranking-local/src/com/intellij/completion/ml/local/models/storage/MethodsFrequencyStorage.kt +++ b/plugins/completion-ml-ranking-local/src/com/intellij/completion/ml/local/models/storage/MethodsFrequencyStorage.kt @@ -1,6 +1,8 @@ package com.intellij.completion.ml.local.models.storage import com.intellij.completion.ml.local.models.frequency.MethodsFrequencies +import com.intellij.completion.ml.local.models.storage.StorageUtil.clear +import com.intellij.completion.ml.local.models.storage.StorageUtil.isEmpty import com.intellij.util.Processor import com.intellij.util.io.DataExternalizer import com.intellij.util.io.EnumeratorStringDescriptor @@ -16,13 +18,11 @@ class MethodsFrequencyStorage internal constructor(private val storageDirectory: fun getStorage(baseDirectory: Path): MethodsFrequencyStorage { val storageDirectory = baseDirectory.resolve(STORAGE_NAME) - val isValid = StorageUtil.prepareStorage(storageDirectory, VERSION) - val storage = MethodsFrequencyStorage(storageDirectory) - storage.setValid(isValid) - return storage + StorageUtil.prepareStorage(storageDirectory, VERSION) + return MethodsFrequencyStorage(storageDirectory) } } - private var isValid: Boolean = false + private var isValid: Boolean = true private val storage = PersistentHashMap(storageDirectory.resolve(STORAGE_NAME), EnumeratorStringDescriptor(), MyDataExternalizer()) @@ -33,20 +33,26 @@ class MethodsFrequencyStorage internal constructor(private val storageDirectory: @Volatile var totalMethodsUsages = 0 private set + init { + getTotalCounts() + } + override fun name(): String = STORAGE_NAME override fun version(): Int = VERSION override fun isValid(): Boolean = isValid + override fun isEmpty(): Boolean = storage.isEmpty() + override fun setValid(isValid: Boolean) { if (isValid) { getTotalCounts() } else { - StorageUtil.clearMap(storage) + storage.clear() } this.isValid = isValid - StorageUtil.saveInfo(StorageInfo(VERSION, isValid), storageDirectory) + StorageUtil.saveInfo(VERSION, isValid, storageDirectory) } @Synchronized diff --git a/plugins/completion-ml-ranking-local/src/com/intellij/completion/ml/local/models/storage/StorageInfo.kt b/plugins/completion-ml-ranking-local/src/com/intellij/completion/ml/local/models/storage/StorageInfo.kt deleted file mode 100644 index c6ce7ea91097..000000000000 --- a/plugins/completion-ml-ranking-local/src/com/intellij/completion/ml/local/models/storage/StorageInfo.kt +++ /dev/null @@ -1,3 +0,0 @@ -package com.intellij.completion.ml.local.models.storage - -data class StorageInfo(val version: Int, val isValid: Boolean) \ No newline at end of file diff --git a/plugins/completion-ml-ranking-local/src/com/intellij/completion/ml/local/models/storage/StorageUtil.kt b/plugins/completion-ml-ranking-local/src/com/intellij/completion/ml/local/models/storage/StorageUtil.kt index 92d25b9474f8..658ce3d4869d 100644 --- a/plugins/completion-ml-ranking-local/src/com/intellij/completion/ml/local/models/storage/StorageUtil.kt +++ b/plugins/completion-ml-ranking-local/src/com/intellij/completion/ml/local/models/storage/StorageUtil.kt @@ -10,40 +10,41 @@ object StorageUtil { private const val STORAGE_INFO_FILE = "info.json" - fun readInfo(storageDirectory: Path): StorageInfo? { + private fun readInfo(storageDirectory: Path): StorageInfo? { val infoFile = storageDirectory.resolve(STORAGE_INFO_FILE) if (!infoFile.exists()) return null return GSON.fromJson(infoFile.readText(), StorageInfo::class.java) } - fun saveInfo(info: StorageInfo, storageDirectory: Path) { + fun saveInfo(version: Int, isValid: Boolean, storageDirectory: Path) { val infoFile = storageDirectory.resolve(STORAGE_INFO_FILE) - infoFile.writeText(GSON.toJson(info)) + infoFile.writeText(GSON.toJson(StorageInfo(version, isValid))) } - fun clearMap(map: PersistentHashMap) { + fun PersistentHashMap.isEmpty(): Boolean = this.processKeysWithExistingMapping { false } + + fun PersistentHashMap.clear() { val existing = HashSet() - map.processKeysWithExistingMapping { + this.processKeysWithExistingMapping { existing.add(it) true } existing.forEach { fileId -> - map.remove(fileId) + this.remove(fileId) } } - fun prepareStorage(storageDirectory: Path, version: Int): Boolean { - var isValid = false + fun prepareStorage(storageDirectory: Path, version: Int) { if (storageDirectory.exists()) { val info = readInfo(storageDirectory) if (info == null || !info.isValid || info.version != version) { storageDirectory.delete() - } else { - isValid = true } } storageDirectory.createDirectories() - return isValid + saveInfo(version, true, storageDirectory) } + + private data class StorageInfo(val version: Int, val isValid: Boolean) } \ No newline at end of file