ML in SE: Add better tokenizer (LAB-84)

GitOrigin-RevId: ae969ab64990f382836796f944fcd0d8ca7967b0

plugins/search-everywhere-ml/typos/intellij.searchEverywhereMl.typos.iml

@@ -5,6 +5,7 @@
     <content url="file://$MODULE_DIR$">
       <sourceFolder url="file://$MODULE_DIR$/resources" type="java-resource" />
       <sourceFolder url="file://$MODULE_DIR$/src" isTestSource="false" />
+      <sourceFolder url="file://$MODULE_DIR$/testSrc" isTestSource="true" />
     </content>
     <orderEntry type="inheritedJdk" />
     <orderEntry type="sourceFolder" forTests="false" />

plugins/search-everywhere-ml/typos/src/com/intellij/searchEverywhereMl/typos/ActionsLanguageModel.kt

@@ -94,7 +94,11 @@ internal class ActionsLanguageModel(private val actionsDictionary: ActionsDictio
 
     override fun run() {
       (getWordsFromActions() + getWordsFromSettings())
-        .flatMap { it.split(" ", "-", ".") }
+        .flatMap {
+          splitText(it)
+            .filter { token -> token is SearchEverywhereStringToken.Word }
+            .map { token -> token.value }
+        }
         .map { it.filter { c -> c.isLetterOrDigit() } }
         .filterNot { it.isEmpty() || it.isSingleCharacter() }
         .map { it.lowercase() }

plugins/search-everywhere-ml/typos/src/com/intellij/searchEverywhereMl/typos/ActionsTabTypoFixSuggestionProvider.kt

@@ -30,18 +30,23 @@ internal class ActionsTabTypoFixSuggestionProvider(project: Project) {
       return _suggestionRanker
     }
 
-  fun suggestFixFor(query: String): SearchEverywhereSpellCheckResult = query.split(" ")
-    .filter { it.isNotBlank() }
-    .map(::correctWord)
+  fun suggestFixFor(query: String): SearchEverywhereSpellCheckResult = splitText(query)
+    .map { token ->
+      when (token) {
+        is SearchEverywhereStringToken.Delimiter -> WordSpellCheckResult.NoCorrection(token.value)
+        is SearchEverywhereStringToken.Word -> correctWord(token.value)
+      }
+    }
     .toQuerySpellCheckResult()
 
   private fun correctWord(word: String): WordSpellCheckResult = word.takeIf { spellChecker.hasProblem(it) }
     ?.let { misspelledWord ->
-      val correctionSuggestions = spellChecker.getSuggestions(misspelledWord).toLinkedSet()
+      val correctionSuggestions = spellChecker.getSuggestions(misspelledWord.lowercase()).toLinkedSet()
 
-      suggestionRanker?.score(word, correctionSuggestions)
+      suggestionRanker?.score(word.lowercase(), correctionSuggestions)
         ?.asSequence()
         ?.maxBy { it.value }
+        ?.let { (correction, confidence) -> correction.capitalizeBasedOn(word) to confidence }
         ?.let { (word, confidence) -> WordSpellCheckResult.Correction(word, confidence) }
     } ?: WordSpellCheckResult.NoCorrection(word)
 
@@ -50,12 +55,19 @@ internal class ActionsTabTypoFixSuggestionProvider(project: Project) {
     class NoCorrection(originalWord: String) : WordSpellCheckResult(originalWord)
   }
 
-  private fun Collection<WordSpellCheckResult>.toQuerySpellCheckResult(): SearchEverywhereSpellCheckResult {
+  private fun Sequence<WordSpellCheckResult>.toQuerySpellCheckResult(): SearchEverywhereSpellCheckResult {
     if (all { it is WordSpellCheckResult.NoCorrection }) return SearchEverywhereSpellCheckResult.NoCorrection
 
-    val correctedQuery = joinToString(" ") { it.word }
+    val correctedQuery = joinToString("") { it.word }
     val confidence = filterIsInstance<WordSpellCheckResult.Correction>().map { it.confidence }.average()
     return SearchEverywhereSpellCheckResult.Correction(correctedQuery, confidence)
   }
+
+  private fun String.capitalizeBasedOn(other: String): String {
+    if (other.all { it.isLowerCase() }) return this.lowercase()
+    if (other.all { it.isUpperCase() }) return this.uppercase()
+    if (other.first().isUpperCase()) return this.replaceFirstChar { it.uppercase() }
+    else return this
+  }
 }
 

plugins/search-everywhere-ml/typos/src/com/intellij/searchEverywhereMl/typos/SearchEverywhereStringTokenizer.kt

@@ -0,0 +1,65 @@
+package com.intellij.searchEverywhereMl.typos
+
+
+/**
+ * Splits text by non-alphanumeric characters
+ *
+ * See test cases in [com.intellij.searchEverywhereMl.typos.SearchEverywhereStringTokenizerTest] for more examples.
+ */
+internal fun splitText(text: String): Sequence<SearchEverywhereStringToken> {
+  return sequence {
+    val builder = StringBuilder()
+    text.forEachIndexed { index, c ->
+      if (c.isLetterOrDigit().not()) {
+        yieldWordIfNotBlank(builder.toString())
+        builder.clear()
+
+        yieldDelimiter(c.toString())
+        return@forEachIndexed
+      }
+
+      if (c.isUpperCase() && builder.isNotEmpty() && builder.last().isLowerCase()) {
+        yieldWord(builder.toString())
+        builder.clear()
+      }
+
+      builder.append(c)
+
+      if (index == text.lastIndex) {
+        yieldWordIfNotBlank(builder.toString())
+      }
+    }
+  }
+}
+
+internal sealed interface SearchEverywhereStringToken {
+  val value: String
+
+  @JvmInline
+  value class Word(override val value: String) : SearchEverywhereStringToken
+
+  @JvmInline
+  value class Delimiter(override val value: String) : SearchEverywhereStringToken
+}
+
+private suspend fun SequenceScope<SearchEverywhereStringToken>.yieldToken(token: SearchEverywhereStringToken) {
+  yield(token)
+}
+
+private suspend fun SequenceScope<SearchEverywhereStringToken>.yieldWord(word: String) {
+  yieldToken(SearchEverywhereStringToken.Word(word))
+}
+
+private suspend fun SequenceScope<SearchEverywhereStringToken>.yieldWordIfNotBlank(word: String) {
+  if (word.isNotBlank()) {
+    yieldWord(word)
+  }
+}
+
+private suspend fun SequenceScope<SearchEverywhereStringToken>.yieldDelimiter(delimiter: String) {
+  yieldToken(SearchEverywhereStringToken.Delimiter(delimiter))
+}
+
+private fun StringBuilder.clear() {
+  this.delete(0, length)
+}

plugins/search-everywhere-ml/typos/testSrc/com/intellij/searchEverywhereMl/typos/SearchEverywhereStringTokenizerTest.kt

@@ -0,0 +1,83 @@
+package com.intellij.searchEverywhereMl.typos
+
+import com.intellij.searchEverywhereMl.typos.SearchEverywhereStringToken.Delimiter
+import com.intellij.searchEverywhereMl.typos.SearchEverywhereStringToken.Word
+import org.junit.jupiter.api.Assertions.assertEquals
+import org.junit.jupiter.api.Assertions.assertTrue
+import org.junit.jupiter.api.Test
+
+internal class SearchEverywhereStringTokenizerTest {
+  @Test
+  fun `empty string yields no tokens`() {
+    assertTrue(splitText("").toList().isEmpty())
+  }
+
+  @Test
+  fun `blank string yields delimiter tokens`() {
+    val expected = listOf(Delimiter(" "), Delimiter("\t"), Delimiter("\n"))
+    val actual = splitText(" \t\n").toList()
+
+    assertEquals(expected, actual)
+  }
+
+  @Test
+  fun `words are split on case change`() {
+    val expected = listOf(Word("Hello"), Word("Brave"), Word("New"), Word("World"))
+    val actual = splitText("HelloBraveNewWorld").toList()
+
+    assertEquals(expected, actual)
+  }
+
+  @Test
+  fun `words are split by delimiters`() {
+    val expected = listOf(Word("hello"), Delimiter("."),
+                          Word("brave"), Delimiter("_"),
+                          Word("new"), Delimiter("-"),
+                          Word("world"))
+    val actual = splitText("hello.brave_new-world").toList()
+
+    assertEquals(expected, actual)
+  }
+
+  @Test
+  fun `multiple consecutive delimiters are accepted and yielded`() {
+    val expected = listOf(Word("hello"), Delimiter("."), Delimiter("_"),
+                          Word("world"))
+    val actual = splitText("hello._world").toList()
+
+    assertEquals(expected, actual)
+  }
+
+  @Test
+  fun `multiple consecutive spaces are accepted and yielded back in the original form`() {
+    val expected = listOf(Word("hello"), Delimiter(" "), Delimiter(" "),
+                          Word("world"))
+    val actual = splitText("hello  world").toList()
+    assertEquals(expected, actual)
+  }
+
+  @Test
+  fun `all caps are yielded as a single word`() {
+    val expected = listOf(Word("HBNW"))
+    val actual = splitText("HBNW").toList()
+
+    assertEquals(expected, actual)
+  }
+
+  @Test
+  fun `delimiter at the end gets yielded`() {
+    val expected = listOf(Word("hello"), Delimiter("."),
+                          Word("world"), Delimiter("."))
+    val actual = splitText("hello.world.").toList()
+
+    assertEquals(expected, actual)
+  }
+
+  @Test
+  fun `space and capital letter`() {
+    val expected = listOf(Word("hello"), Delimiter(" "), Word("World"))
+    val actual = splitText("hello World").toList()
+
+    assertEquals(expected, actual)
+  }
+}