diff --git a/platform/lang-api/src/com/intellij/codeInsight/completion/ml/ContextFeatures.java b/platform/lang-api/src/com/intellij/codeInsight/completion/ml/ContextFeatures.java
new file mode 100644
index 000000000000..673322bbbf2a
--- /dev/null
+++ b/platform/lang-api/src/com/intellij/codeInsight/completion/ml/ContextFeatures.java
@@ -0,0 +1,16 @@
+// Copyright 2000-2019 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.NotNull;
+import org.jetbrains.annotations.Nullable;
+
+public interface ContextFeatures {
+  @Nullable
+  Boolean binaryValue(@NotNull String name);
+
+  @Nullable
+  Double floatValue(@NotNull String name);
+
+  @Nullable
+  String categoricalValue(@NotNull String name);
+}
diff --git a/platform/lang-api/src/com/intellij/codeInsight/completion/ml/ContextFeatures.kt b/platform/lang-api/src/com/intellij/codeInsight/completion/ml/ContextFeatures.kt
deleted file mode 100644
index d2479c6c57d7..000000000000
--- a/platform/lang-api/src/com/intellij/codeInsight/completion/ml/ContextFeatures.kt
+++ /dev/null
@@ -1,32 +0,0 @@
-// Copyright 2000-2019 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 com.intellij.openapi.util.Key
-import com.intellij.psi.PsiElement
-import java.util.*
-
-class ContextFeatures(private val featuresSnapshot: Map<String, MLFeatureValue>) {
-  companion object {
-    private val CONTEXT_FEATURES_KEY = Key.create<ContextFeatures>("com.intellij.completion.ml.context_features")
-    private val EMPTY = ContextFeatures(emptyMap())
-
-    fun extract(element: PsiElement): ContextFeatures {
-      return element.getUserData(CONTEXT_FEATURES_KEY) ?: EMPTY
-    }
-
-    fun setContextFeatures(element: PsiElement, features: Map<String, MLFeatureValue>) {
-      element.putUserData(CONTEXT_FEATURES_KEY, ContextFeatures(Collections.unmodifiableMap(features)))
-    }
-
-    fun clear(element: PsiElement) {
-      element.putUserData(CONTEXT_FEATURES_KEY, null)
-    }
-  }
-
-  fun binaryValue(name: String): Boolean? = (featuresSnapshot[name])?.asBinary()
-
-  fun floatValue(name: String): Double? = (featuresSnapshot[name])?.asFloat()
-
-  fun categoricalValue(name: String): String? = (featuresSnapshot[name])?.asCategorical()
-}
-
diff --git a/plugins/stats-collector/src/com/intellij/completion/ml/ContextFeaturesStorage.kt b/plugins/stats-collector/src/com/intellij/completion/ml/ContextFeaturesStorage.kt
new file mode 100644
index 000000000000..0293b7b48704
--- /dev/null
+++ b/plugins/stats-collector/src/com/intellij/completion/ml/ContextFeaturesStorage.kt
@@ -0,0 +1,35 @@
+// Copyright 2000-2019 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
+
+import com.intellij.codeInsight.completion.ml.ContextFeatures
+import com.intellij.codeInsight.completion.ml.MLFeatureValue
+import com.intellij.openapi.util.Key
+import com.intellij.psi.PsiElement
+import java.util.*
+
+class ContextFeaturesStorage(private val featuresSnapshot: Map<String, MLFeatureValue>): ContextFeatures {
+  companion object {
+    private val CONTEXT_FEATURES_KEY = Key.create<ContextFeaturesStorage>("com.intellij.completion.ml.context_features")
+    private val EMPTY = ContextFeaturesStorage(emptyMap())
+
+    fun extract(element: PsiElement): ContextFeaturesStorage {
+      return element.getUserData(CONTEXT_FEATURES_KEY) ?: EMPTY
+    }
+
+    fun setContextFeatures(element: PsiElement, features: Map<String, MLFeatureValue>) {
+      element.putUserData(CONTEXT_FEATURES_KEY,
+                          ContextFeaturesStorage(Collections.unmodifiableMap(features)))
+    }
+
+    fun clear(element: PsiElement) {
+      element.putUserData(CONTEXT_FEATURES_KEY, null)
+    }
+  }
+
+  override fun binaryValue(name: String): Boolean? = (featuresSnapshot[name])?.asBinary()
+
+  override fun floatValue(name: String): Double? = (featuresSnapshot[name])?.asFloat()
+
+  override fun categoricalValue(name: String): String? = (featuresSnapshot[name])?.asCategorical()
+}
+
diff --git a/plugins/stats-collector/src/com/intellij/completion/ml/MLCompletionWeigher.kt b/plugins/stats-collector/src/com/intellij/completion/ml/MLCompletionWeigher.kt
index 94770dfa8122..fea5f041c538 100644
--- a/plugins/stats-collector/src/com/intellij/completion/ml/MLCompletionWeigher.kt
+++ b/plugins/stats-collector/src/com/intellij/completion/ml/MLCompletionWeigher.kt
@@ -3,14 +3,13 @@ package com.intellij.completion.ml
 
 import com.intellij.codeInsight.completion.CompletionLocation
 import com.intellij.codeInsight.completion.CompletionWeigher
-import com.intellij.codeInsight.completion.ml.ContextFeatures
 import com.intellij.codeInsight.completion.ml.ElementFeatureProvider
 import com.intellij.codeInsight.lookup.LookupElement
 
 class MLCompletionWeigher : CompletionWeigher() {
   override fun weigh(element: LookupElement, location: CompletionLocation): Comparable<Nothing>? {
     val psiFile = location.completionParameters.originalFile
-    val contextFeatures = ContextFeatures.extract(psiFile)
+    val contextFeatures = ContextFeaturesStorage.extract(psiFile)
     val result = mutableMapOf<String, Any>()
     for (provider in ElementFeatureProvider.forLanguage(psiFile.language)) {
       val name = provider.name
diff --git a/plugins/stats-collector/src/com/intellij/stats/completion/CompletionTrackerInitializer.kt b/plugins/stats-collector/src/com/intellij/stats/completion/CompletionTrackerInitializer.kt
index 660789ccef83..cc8b2b2a4c31 100644
--- a/plugins/stats-collector/src/com/intellij/stats/completion/CompletionTrackerInitializer.kt
+++ b/plugins/stats-collector/src/com/intellij/stats/completion/CompletionTrackerInitializer.kt
@@ -4,7 +4,7 @@ package com.intellij.stats.completion
 import com.intellij.codeInsight.lookup.LookupManager
 import com.intellij.codeInsight.lookup.impl.LookupImpl
 import com.intellij.codeInsight.completion.ml.ContextFeatureProvider
-import com.intellij.codeInsight.completion.ml.ContextFeatures
+import com.intellij.completion.ml.ContextFeaturesStorage
 import com.intellij.codeInsight.completion.ml.MLFeatureValue
 import com.intellij.completion.settings.CompletionMLRankingSettings
 import com.intellij.completion.tracker.PositionTrackingListener
@@ -107,8 +107,8 @@ class CompletionTrackerInitializer(experimentHelper: WebServiceStatus) {
         }
       }
 
-      ContextFeatures.setContextFeatures(file, result)
-      Disposer.register(lookup, Disposable { ContextFeatures.clear(file) })
+      ContextFeaturesStorage.setContextFeatures(file, result)
+      Disposer.register(lookup, Disposable { ContextFeaturesStorage.clear(file) })
       lookupStorage.contextFactors = result.mapValues { it.value.toString() }
     }
   }