[recommenders] use MLFeatures instead of raw values in feature providers

(cherry picked from commit a7ef61a7580365bad6f59d82b11835076edc8be4)

IJ-MR-27840

GitOrigin-RevId: 58ca8aebb98ca0c97cfccc7f53a881c6e2f2712f

platform/lang-api/src/com/intellij/codeInsight/completion/ml/GenericFeatureProvider.kt

@@ -1,37 +0,0 @@
-// Copyright 2000-2022 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
-package com.intellij.codeInsight.completion.ml
-
-import com.intellij.lang.Language
-import com.intellij.lang.LanguageExtension
-import org.jetbrains.annotations.ApiStatus
-import kotlin.system.measureTimeMillis
-
-@ApiStatus.Internal
-interface GenericFeatureProvider {
-  val name: String
-  fun knownFeatures(): Set<String>
-  fun calculateFeatures(context: MLContext): Map<String, Any>
-  fun isApplicable(context: MLContext): Boolean
-
-  data class FeaturesCalculationResult(val features: Map<String, Any>, val performance: Map<String, Long>)
-
-  companion object {
-    private val EP_NAME = LanguageExtension<GenericFeatureProvider>("com.intellij.completion.ml.genericFeatures")
-
-    fun forLanguage(language: Language): List<GenericFeatureProvider> = EP_NAME.allForLanguageOrAny(language)
-
-    fun calculateFeatures(context: MLContext, features: Set<String>): FeaturesCalculationResult {
-      val featureValues = mutableMapOf<String, Any>()
-      val performance = mutableMapOf<String, Long>()
-      for (provider in forLanguage(context.position.language)) {
-        if (provider.knownFeatures().any { it in features }) {
-          val time = measureTimeMillis {
-            featureValues.putAll(provider.calculateFeatures(context))
-          }
-          performance[provider.name] = time
-        }
-      }
-      return FeaturesCalculationResult(featureValues, performance)
-    }
-  }
-}

platform/lang-api/src/com/intellij/codeInsight/completion/ml/MLFeature.kt

@@ -0,0 +1,11 @@
+// Copyright 2000-2022 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
+package com.intellij.codeInsight.completion.ml
+
+interface MLFeature {
+  val name: String
+  fun compute(): MLFeatureValue
+}
+
+data class SimpleMLFeature(override val name: String, private val value: MLFeatureValue) : MLFeature {
+  override fun compute(): MLFeatureValue = value
+}

platform/lang-api/src/com/intellij/codeInsight/completion/ml/MLFeatureProvider.kt

@@ -0,0 +1,20 @@
+// Copyright 2000-2022 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
+package com.intellij.codeInsight.completion.ml
+
+import com.intellij.lang.LanguageExtension
+import org.jetbrains.annotations.ApiStatus
+
+@ApiStatus.Internal
+interface MLFeatureProvider {
+  val name: String
+  fun knownFeatures(): Set<String>
+  fun calculateFeatures(context: MLContext): List<MLFeature>
+  fun isApplicable(context: MLContext): Boolean
+
+  companion object {
+    private val EP_NAME = LanguageExtension<MLFeatureProvider>("com.intellij.completion.ml.features")
+
+    fun applicableProviders(context: MLContext): List<MLFeatureProvider> =
+      EP_NAME.allForLanguageOrAny(context.position.language).filter { it.isApplicable(context) }
+  }
+}

platform/lang-api/src/com/intellij/codeInsight/completion/ml/MLFeatureValue.kt

@@ -1,4 +1,4 @@
-// Copyright 2000-2021 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license that can be found in the LICENSE file.
+// Copyright 2000-2022 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license that can be found in the LICENSE file.
 package com.intellij.codeInsight.completion.ml
 
 import org.jetbrains.annotations.ApiStatus
@@ -57,15 +57,32 @@ sealed class MLFeatureValue {
     @JvmStatic
     fun <T : Class<*>> className(value: T, useSimpleName: Boolean = true): MLFeatureValue = ClassNameValue(value, useSimpleName)
 
+    @JvmStatic
+    fun string(value: String): MLFeatureValue = StringValue(value)
+
+    @JvmStatic
+    fun list(value: List<*>): MLFeatureValue = ListValue(value)
+
     @JvmStatic
     fun version(value: String): MLFeatureValue = VersionValue(value)
   }
 
   abstract val value: Any
+  abstract val isSafe: Boolean
 
-  data class BinaryValue internal constructor(override val value: Boolean) : MLFeatureValue()
-  data class FloatValue internal constructor(override val value: Double) : MLFeatureValue()
-  data class CategoricalValue internal constructor(override val value: String) : MLFeatureValue()
-  data class ClassNameValue internal constructor(override val value: Class<*>, val useSimpleName: Boolean) : MLFeatureValue()
-  data class VersionValue internal constructor(override val value: String) : MLFeatureValue()
+  abstract class SafeValue : MLFeatureValue() {
+    override val isSafe: Boolean = true
+  }
+
+  abstract class UnSafeValue : MLFeatureValue() {
+    override val isSafe: Boolean = false
+  }
+
+  data class BinaryValue internal constructor(override val value: Boolean) : SafeValue()
+  data class FloatValue internal constructor(override val value: Double) : SafeValue()
+  data class CategoricalValue internal constructor(override val value: String) : SafeValue()
+  data class ClassNameValue internal constructor(override val value: Class<*>, val useSimpleName: Boolean) : UnSafeValue()
+  data class VersionValue internal constructor(override val value: String) : UnSafeValue()
+  data class StringValue internal constructor(override val value: String) : UnSafeValue()
+  data class ListValue internal constructor(override val value: List<*>) : UnSafeValue()
 }

platform/platform-resources/src/META-INF/CompletionExtensionPoints.xml

@@ -15,8 +15,8 @@
       <with attribute="implementationClass" implements="com.intellij.codeInsight.completion.ml.ElementFeatureProvider"/>
     </extensionPoint>
 
-    <extensionPoint name="completion.ml.genericFeatures" beanClass="com.intellij.lang.LanguageExtensionPoint" dynamic="true">
-      <with attribute="implementationClass" implements="com.intellij.codeInsight.completion.ml.GenericFeatureProvider"/>
+    <extensionPoint name="completion.ml.features" beanClass="com.intellij.lang.LanguageExtensionPoint" dynamic="true">
+      <with attribute="implementationClass" implements="com.intellij.codeInsight.completion.ml.MLFeatureProvider"/>
     </extensionPoint>
 
     <extensionPoint name="completion.plainTextSymbol" beanClass="com.intellij.codeInsight.completion.PlainTextSymbolCompletionContributorEP" dynamic="true">

plugins/completion-ml-ranking/src/com/intellij/completion/ml/features/MLFeaturesUtil.kt

@@ -14,6 +14,7 @@ internal object MLFeaturesUtil {
       is MLFeatureValue.CategoricalValue -> featureValue.value
       is MLFeatureValue.ClassNameValue -> getClassNameSafe(featureValue)
       is MLFeatureValue.VersionValue -> getVersionSafe(featureValue)
+      else -> if (featureValue.isSafe) featureValue.value else throw IllegalArgumentException("Feature value $featureValue is unsafe")
     }
   }