From 31315e2c45c00f54725d600b7e44c62e96b5010d Mon Sep 17 00:00:00 2001
From: Roman Golyshev <roman.golyshev@jetbrains.com>
Date: Fri, 31 Jan 2020 16:00:41 +0300
Subject: [PATCH] Add `plugin_version` and `file_type` features for ML
 completion

- This is required to be able to use ML completion infrastructure to monitor code completion quality (even without ML enhancement)
---
 .../completion/KotlinCompletionContributor.kt |  6 +--
 .../ml/KotlinContextFeatureProvider.kt        | 41 +++++++++++++++++++
 .../ml/KotlinContextFeatureProvider.kt.192    |  0
 .../idea/statistics/CompletionFUSCollector.kt | 11 ++++-
 idea/resources/META-INF/plugin.xml            |  1 +
 5 files changed, 53 insertions(+), 6 deletions(-)
 create mode 100644 idea/idea-completion/src/org/jetbrains/kotlin/idea/completion/ml/KotlinContextFeatureProvider.kt
 create mode 100644 idea/idea-completion/src/org/jetbrains/kotlin/idea/completion/ml/KotlinContextFeatureProvider.kt.192

diff --git a/idea/idea-completion/src/org/jetbrains/kotlin/idea/completion/KotlinCompletionContributor.kt b/idea/idea-completion/src/org/jetbrains/kotlin/idea/completion/KotlinCompletionContributor.kt
index 353225e68cf..9ef3fba0414 100644
--- a/idea/idea-completion/src/org/jetbrains/kotlin/idea/completion/KotlinCompletionContributor.kt
+++ b/idea/idea-completion/src/org/jetbrains/kotlin/idea/completion/KotlinCompletionContributor.kt
@@ -281,11 +281,7 @@ class KotlinCompletionContributor : CompletionContributor() {
                 CompletionType.CLASS_NAME -> CompletionTypeStats.BASIC //there is no class name anymore actually
                 CompletionType.SMART -> CompletionTypeStats.SMART
             },
-            fileType = when {
-                name.endsWith(".kt") -> FileTypeStats.KT
-                name.endsWith(".gradle.kts") -> FileTypeStats.GRADLEKTS
-                else -> FileTypeStats.KTS
-            },
+            fileType = FileTypeStats.parseFromFileName(name),
             invocationCount = parameters.invocationCount
         )
         val position = parameters.position
diff --git a/idea/idea-completion/src/org/jetbrains/kotlin/idea/completion/ml/KotlinContextFeatureProvider.kt b/idea/idea-completion/src/org/jetbrains/kotlin/idea/completion/ml/KotlinContextFeatureProvider.kt
new file mode 100644
index 00000000000..1bf90febf1d
--- /dev/null
+++ b/idea/idea-completion/src/org/jetbrains/kotlin/idea/completion/ml/KotlinContextFeatureProvider.kt
@@ -0,0 +1,41 @@
+/*
+ * Copyright 2010-2020 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.idea.completion.ml
+
+import com.intellij.codeInsight.completion.ml.CompletionEnvironment
+import com.intellij.codeInsight.completion.ml.ContextFeatureProvider
+import com.intellij.codeInsight.completion.ml.MLFeatureValue
+import org.jetbrains.kotlin.idea.KotlinPluginUtil
+import org.jetbrains.kotlin.idea.statistics.FileTypeStats
+
+class KotlinContextFeatureProvider : ContextFeatureProvider {
+    override fun getName(): String = "kotlin"
+
+    override fun calculateFeatures(environment: CompletionEnvironment): Map<String, MLFeatureValue> {
+        val features = mutableMapOf("plugin_version" to MLFeatureValue.categorical(KotlinVersionFakeEnum.VERSION))
+
+        val fileType = environment.parameters.originalFile.virtualFile?.name?.let { FileTypeStats.parseFromFileName(it) }
+        if (fileType != null) {
+            features["file_type"] = MLFeatureValue.categorical(fileType)
+        }
+
+        return features
+    }
+}
+
+/**
+ * We do not have a enum to represent Kotlin Plugin version; because of that we cannot
+ * make it a categorical feature for ML completion (`MLFeatureValue.categorical` accepts only enums).
+ *
+ * This fake enum is used as a workaround to this problem.
+ *
+ * TODO As soon as there would be a way to pass Kotlin Plugin version without this enum, it should be removed.
+ */
+private enum class KotlinVersionFakeEnum {
+    VERSION;
+
+    override fun toString(): String = KotlinPluginUtil.getPluginVersion()
+}
diff --git a/idea/idea-completion/src/org/jetbrains/kotlin/idea/completion/ml/KotlinContextFeatureProvider.kt.192 b/idea/idea-completion/src/org/jetbrains/kotlin/idea/completion/ml/KotlinContextFeatureProvider.kt.192
new file mode 100644
index 00000000000..e69de29bb2d
diff --git a/idea/idea-core/src/org/jetbrains/kotlin/idea/statistics/CompletionFUSCollector.kt b/idea/idea-core/src/org/jetbrains/kotlin/idea/statistics/CompletionFUSCollector.kt
index c524c315249..92b120a3b51 100644
--- a/idea/idea-core/src/org/jetbrains/kotlin/idea/statistics/CompletionFUSCollector.kt
+++ b/idea/idea-core/src/org/jetbrains/kotlin/idea/statistics/CompletionFUSCollector.kt
@@ -71,7 +71,16 @@ object CompletionFUSCollector {
 }
 
 enum class FileTypeStats {
-    KT, GRADLEKTS, KTS
+    KT, GRADLEKTS, KTS;
+
+    companion object {
+        fun parseFromFileName(fileName: String): FileTypeStats? = when {
+            fileName.endsWith(".kt") -> KT
+            fileName.endsWith(".gradle.kts") -> GRADLEKTS
+            fileName.endsWith(".kts") -> KTS
+            else -> null
+        }
+    }
 }
 
 enum class CompletionTypeStats {
diff --git a/idea/resources/META-INF/plugin.xml b/idea/resources/META-INF/plugin.xml
index ad77898df68..0510f557b8e 100644
--- a/idea/resources/META-INF/plugin.xml
+++ b/idea/resources/META-INF/plugin.xml
@@ -88,5 +88,6 @@ The Kotlin plugin provides language support in IntelliJ IDEA and Android Studio.
       <fileTypeUsageSchemaDescriptor schema="Gradle Script" implementationClass="org.jetbrains.kotlin.idea.core.script.KotlinGradleScriptFileTypeSchemaDetector"/>
 
       <completion.ml.model implementation="org.jetbrains.kotlin.idea.completion.ml.KotlinMLRankingProvider"/>
+      <completion.ml.contextFeatures language="kotlin" implementationClass="org.jetbrains.kotlin.idea.completion.ml.KotlinContextFeatureProvider"/>
     </extensions>
 </idea-plugin>