From 71bc043aa323330128a6834165e755b68e0d4f5b Mon Sep 17 00:00:00 2001
From: Alexey Kalina <alexey.kalina@jetbrains.com>
Date: Wed, 30 Mar 2022 09:37:51 +0300
Subject: [PATCH] [recommenders] introduce feature providers for recommenders
 features

(cherry picked from commit de5a0200302516736d2d459a56c3ff62c9b5820e)

IJ-MR-27840

GitOrigin-RevId: dcea1afa474c3f4bbf4967cc3258dc21df6779d4
---
 .../completion/ml/GenericFeatureProvider.kt   | 37 +++++++++++++++++++
 .../codeInsight/completion/ml/MLContext.kt    | 10 +++++
 .../META-INF/CompletionExtensionPoints.xml    |  4 ++
 .../CompletionPerformanceTracker.kt           |  6 ++-
 4 files changed, 56 insertions(+), 1 deletion(-)
 create mode 100644 platform/lang-api/src/com/intellij/codeInsight/completion/ml/GenericFeatureProvider.kt
 create mode 100644 platform/lang-api/src/com/intellij/codeInsight/completion/ml/MLContext.kt

diff --git a/platform/lang-api/src/com/intellij/codeInsight/completion/ml/GenericFeatureProvider.kt b/platform/lang-api/src/com/intellij/codeInsight/completion/ml/GenericFeatureProvider.kt
new file mode 100644
index 000000000000..853a2e6329de
--- /dev/null
+++ b/platform/lang-api/src/com/intellij/codeInsight/completion/ml/GenericFeatureProvider.kt
@@ -0,0 +1,37 @@
+// 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)
+    }
+  }
+}
diff --git a/platform/lang-api/src/com/intellij/codeInsight/completion/ml/MLContext.kt b/platform/lang-api/src/com/intellij/codeInsight/completion/ml/MLContext.kt
new file mode 100644
index 000000000000..67c0e00c6cd7
--- /dev/null
+++ b/platform/lang-api/src/com/intellij/codeInsight/completion/ml/MLContext.kt
@@ -0,0 +1,10 @@
+// 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.psi.PsiElement
+import org.jetbrains.annotations.ApiStatus
+
+@ApiStatus.Internal
+interface MLContext {
+  val position: PsiElement
+}
diff --git a/platform/platform-resources/src/META-INF/CompletionExtensionPoints.xml b/platform/platform-resources/src/META-INF/CompletionExtensionPoints.xml
index b72d024291ae..96b6d2f6f5e4 100644
--- a/platform/platform-resources/src/META-INF/CompletionExtensionPoints.xml
+++ b/platform/platform-resources/src/META-INF/CompletionExtensionPoints.xml
@@ -15,6 +15,10 @@
       <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>
+
     <extensionPoint name="completion.plainTextSymbol" beanClass="com.intellij.codeInsight.completion.PlainTextSymbolCompletionContributorEP" dynamic="true">
       <with attribute="implementationClass" implements="com.intellij.codeInsight.completion.PlainTextSymbolCompletionContributor"/>
     </extensionPoint>
diff --git a/plugins/completion-ml-ranking/src/com/intellij/completion/ml/performance/CompletionPerformanceTracker.kt b/plugins/completion-ml-ranking/src/com/intellij/completion/ml/performance/CompletionPerformanceTracker.kt
index df7e96b91c66..6c9495073a4b 100644
--- a/plugins/completion-ml-ranking/src/com/intellij/completion/ml/performance/CompletionPerformanceTracker.kt
+++ b/plugins/completion-ml-ranking/src/com/intellij/completion/ml/performance/CompletionPerformanceTracker.kt
@@ -1,4 +1,4 @@
-// Copyright 2000-2020 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.completion.ml.performance
 
 import java.util.concurrent.ConcurrentHashMap
@@ -31,6 +31,10 @@ class CompletionPerformanceTracker {
     addByKey("context.features.$providerName", timeSpent)
   }
 
+  fun recommendersFeaturesCalculated(providerName: String, timeSpent: Long) {
+    addByKey("recommenders.features.$providerName", timeSpent)
+  }
+
   fun itemsScored(itemsCount: Int, timeSpent: Long) {
     addByKey("model.items.$sortingCount", itemsCount.toLong())
     addByKey("model.time.$sortingCount", timeSpent)