From 91c1f6cdec55b560f36f3d229c371cbb928e93b8 Mon Sep 17 00:00:00 2001
From: Alexey Kalina <alexey.kalina@jetbrains.com>
Date: Thu, 14 Jul 2022 15:45:16 +0200
Subject: [PATCH] [stats-collector] move checks on features logging in
 experiment to ExperimentInfo

(cherry picked from commit 31cfee81bca3ce23cdeef1c9625a2670cc13be2c)

IJ-CR-26550

GitOrigin-RevId: 66caba1ab2748d8489120e155d184fe5a9922551
---
 .../completion/ml/experiment/ExperimentInfo.kt        | 11 ++++++++++-
 .../completion/tracker/CompletionLoggerInitializer.kt | 11 +++--------
 2 files changed, 13 insertions(+), 9 deletions(-)

diff --git a/plugins/completion-ml-ranking/src/com/intellij/completion/ml/experiment/ExperimentInfo.kt b/plugins/completion-ml-ranking/src/com/intellij/completion/ml/experiment/ExperimentInfo.kt
index 5645640eb0c7..540a4486c17c 100644
--- a/plugins/completion-ml-ranking/src/com/intellij/completion/ml/experiment/ExperimentInfo.kt
+++ b/plugins/completion-ml-ranking/src/com/intellij/completion/ml/experiment/ExperimentInfo.kt
@@ -1,9 +1,18 @@
 // 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.experiment
 
+import com.intellij.completion.ml.common.CurrentProjectInfo
+import com.intellij.openapi.project.Project
+
 data class ExperimentInfo(val inExperiment: Boolean,
                           val version: Int,
                           val shouldRank: Boolean = false,
                           val shouldShowArrows: Boolean = false,
                           val shouldCalculateFeatures: Boolean = false,
-                          val shouldLogElementFeatures: Boolean = false)
\ No newline at end of file
+                          val shouldLogElementFeatures: Boolean = false) {
+
+  fun shouldLogSessions(project: Project): Boolean = inExperiment || CurrentProjectInfo.getInstance(project).isIdeaProject
+
+  fun shouldLogElementFeatures(project: Project): Boolean =
+    !inExperiment || shouldCalculateFeatures && (shouldLogElementFeatures || CurrentProjectInfo.getInstance(project).isIdeaProject)
+}
\ No newline at end of file
diff --git a/plugins/stats-collector/src/com/intellij/stats/completion/tracker/CompletionLoggerInitializer.kt b/plugins/stats-collector/src/com/intellij/stats/completion/tracker/CompletionLoggerInitializer.kt
index 5ccfa04d68a3..2557ea672cb9 100644
--- a/plugins/stats-collector/src/com/intellij/stats/completion/tracker/CompletionLoggerInitializer.kt
+++ b/plugins/stats-collector/src/com/intellij/stats/completion/tracker/CompletionLoggerInitializer.kt
@@ -2,7 +2,6 @@
 package com.intellij.stats.completion.tracker
 
 import com.intellij.codeInsight.lookup.impl.LookupImpl
-import com.intellij.completion.ml.common.CurrentProjectInfo
 import com.intellij.completion.ml.experiment.ExperimentInfo
 import com.intellij.completion.ml.experiment.ExperimentStatus
 import com.intellij.completion.ml.storage.MutableLookupStorage
@@ -78,17 +77,13 @@ class CompletionLoggerInitializer : LookupTracker() {
     return LoggerPerformanceTracker(actionsTracker, storage.performanceTracker)
   }
 
-  private fun shouldLogElementFeatures(language: Language, project: Project): Boolean {
-    val experimentInfo = ExperimentStatus.getInstance().forLanguage(language)
-    return !experimentInfo.inExperiment ||
-           (experimentInfo.shouldCalculateFeatures && (experimentInfo.shouldLogElementFeatures ||
-                                                       CurrentProjectInfo.getInstance(project).isIdeaProject))
-  }
+  private fun shouldLogElementFeatures(language: Language, project: Project): Boolean =
+    ExperimentStatus.getInstance().forLanguage(language).shouldLogElementFeatures(project)
 
   private fun sessionShouldBeLogged(experimentInfo: ExperimentInfo, language: Language, project: Project): Boolean {
     if (CompletionStatsPolicy.isStatsLogDisabled(language) || !getPluginInfo(language::class.java).isSafeToReport()) return false
     val application = ApplicationManager.getApplication()
-    if (application.isUnitTestMode || experimentInfo.inExperiment || CurrentProjectInfo.getInstance(project).isIdeaProject) return true
+    if (application.isUnitTestMode || experimentInfo.shouldLogSessions(project)) return true
 
     if (!isCompletionLogsSendAllowed()) {
       return false