From d3b9ac41eed36ffecc3d4b047ebb054c429e0b90 Mon Sep 17 00:00:00 2001
From: ortem <ortem00@gmail.com>
Date: Thu, 19 Nov 2020 19:15:57 +0300
Subject: [PATCH] #6411: Add Rust ranking provider based on CatBoost model

---
 build.gradle.kts                                  |  2 ++
 .../kotlin/org/rust/ml/RsMLRankingProvider.kt     | 15 +++++++++++++++
 .../resources/META-INF/ml-completion-only.xml     |  1 +
 .../org/rust/ml/RsModelMetadataConsistencyTest.kt | 13 +++++++++++++
 4 files changed, 31 insertions(+)
 create mode 100644 ml-completion/src/203/main/kotlin/org/rust/ml/RsMLRankingProvider.kt
 create mode 100644 ml-completion/src/203/test/kotlin/org/rust/ml/RsModelMetadataConsistencyTest.kt

diff --git a/build.gradle.kts b/build.gradle.kts
index 58c4bc6e904..5355707dc84 100644
--- a/build.gradle.kts
+++ b/build.gradle.kts
@@ -75,6 +75,7 @@ allprojects {
         jcenter()
         maven("https://dl.bintray.com/jetbrains/markdown")
         maven("http://download.eclipse.org/jgit/maven")
+        maven("https://dl.bintray.com/jetbrains/intellij-third-party-dependencies")
     }
 
     idea {
@@ -464,6 +465,7 @@ project(":ml-completion") {
         }
     }
     dependencies {
+        implementation("org.jetbrains.intellij.deps.completion:completion-ranking-rust:0.0.4")
         implementation(project(":"))
         implementation(project(":common"))
         testImplementation(project(":", "testOutput"))
diff --git a/ml-completion/src/203/main/kotlin/org/rust/ml/RsMLRankingProvider.kt b/ml-completion/src/203/main/kotlin/org/rust/ml/RsMLRankingProvider.kt
new file mode 100644
index 00000000000..8ac0ba1f4d8
--- /dev/null
+++ b/ml-completion/src/203/main/kotlin/org/rust/ml/RsMLRankingProvider.kt
@@ -0,0 +1,15 @@
+/*
+ * Use of this source code is governed by the MIT license that can be
+ * found in the LICENSE file.
+ */
+
+package org.rust.ml
+
+import com.intellij.internal.ml.catboost.CatBoostJarCompletionModelProvider
+import com.intellij.lang.Language
+import org.rust.lang.RsLanguage
+
+@Suppress("UnstableApiUsage")
+class RsMLRankingProvider : CatBoostJarCompletionModelProvider("Rust", "rust_features", "rust_model") {
+    override fun isLanguageSupported(language: Language): Boolean = language == RsLanguage
+}
diff --git a/ml-completion/src/203/main/resources/META-INF/ml-completion-only.xml b/ml-completion/src/203/main/resources/META-INF/ml-completion-only.xml
index bb7850bae68..fb9c5e4471f 100644
--- a/ml-completion/src/203/main/resources/META-INF/ml-completion-only.xml
+++ b/ml-completion/src/203/main/resources/META-INF/ml-completion-only.xml
@@ -3,5 +3,6 @@
     <extensions defaultExtensionNs="com.intellij">
         <completion.ml.ranking.features.policy language="Rust" implementationClass="org.rust.ml.RsCompletionFeaturesPolicy"/>
         <completion.ml.elementFeatures language="Rust" implementationClass="org.rust.ml.RsElementFeatureProvider"/>
+        <completion.ml.model implementation="org.rust.ml.RsMLRankingProvider"/>
     </extensions>
 </idea-plugin>
diff --git a/ml-completion/src/203/test/kotlin/org/rust/ml/RsModelMetadataConsistencyTest.kt b/ml-completion/src/203/test/kotlin/org/rust/ml/RsModelMetadataConsistencyTest.kt
new file mode 100644
index 00000000000..6de8091533d
--- /dev/null
+++ b/ml-completion/src/203/test/kotlin/org/rust/ml/RsModelMetadataConsistencyTest.kt
@@ -0,0 +1,13 @@
+/*
+ * Use of this source code is governed by the MIT license that can be
+ * found in the LICENSE file.
+ */
+
+package org.rust.ml
+
+import org.junit.Test
+
+class RsModelMetadataConsistencyTest {
+    @Test
+    fun `test model metadata consistency`() = RsMLRankingProvider().assertModelMetadataConsistent()
+}