[ml-completion-local] frequency storages refactoring

GitOrigin-RevId: 64c61c774fdf73992de30ca59a19f30170128e0a
This commit is contained in:
Alexey Kalina
2021-02-10 17:43:28 +00:00
committed by intellij-monorepo-bot
parent 74c8007891
commit b4ddea5d82
10 changed files with 54 additions and 36 deletions
@@ -24,21 +24,23 @@ class FrequencyContextFeaturesProvider : ContextFeatureProvider {
override fun calculateFeatures(environment: CompletionEnvironment): MutableMap<String, MLFeatureValue> {
val features = mutableMapOf<String, MLFeatureValue>()
val project = environment.parameters.position.project
LocalModelsManager.getInstance(project).getModel<MethodsFrequencyLocalModel>()?.let { model ->
val methodsModel = LocalModelsManager.getInstance(project).getModel<MethodsFrequencyLocalModel>()
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<ClassesFrequencyLocalModel>()?.let { model ->
features["total_classes"] = MLFeatureValue.numerical(model.totalClassesCount())
features["total_classes_usages"] = MLFeatureValue.numerical(model.totalClassesUsages())
val classesModel = LocalModelsManager.getInstance(project).getModel<ClassesFrequencyLocalModel>()
if (classesModel != null && classesModel.readyToUse()) {
features["total_classes"] = MLFeatureValue.numerical(classesModel.totalClassesCount())
features["total_classes_usages"] = MLFeatureValue.numerical(classesModel.totalClassesUsages())
}
return features
}
@@ -36,7 +36,7 @@ class FrequencyElementFeatureProvider : ElementFeatureProvider {
}
}
val classesModel = LocalModelsManager.getInstance(location.project).getModel<ClassesFrequencyLocalModel>()
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)
@@ -6,4 +6,5 @@ interface LocalModel {
fun fileVisitor(): PsiElementVisitor
fun onStarted()
fun onFinished()
fun readyToUse(): Boolean
}
@@ -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()
}
@@ -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()
}
@@ -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
@@ -4,5 +4,6 @@ interface LocalModelStorage {
fun version(): Int
fun name(): String
fun isValid(): Boolean
fun isEmpty(): Boolean
fun setValid(isValid: Boolean)
}
@@ -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
@@ -1,3 +0,0 @@
package com.intellij.completion.ml.local.models.storage
data class StorageInfo(val version: Int, val isValid: Boolean)
@@ -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<K> clearMap(map: PersistentHashMap<K, *>) {
fun <K> PersistentHashMap<K, *>.isEmpty(): Boolean = this.processKeysWithExistingMapping { false }
fun<K> PersistentHashMap<K, *>.clear() {
val existing = HashSet<K>()
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)
}