Insert equals and invoke completion automatically when completing tom…

…l key

src/main/kotlin/toml/platform/forge/completion/ModsTomlCompletionContributor.kt

@@ -26,6 +26,7 @@ import com.demonwav.mcdev.toml.TomlStringValueInsertionHandler
 import com.demonwav.mcdev.toml.inModsTomlKey
 import com.demonwav.mcdev.toml.inModsTomlValueWithKey
 import com.demonwav.mcdev.toml.platform.forge.ModsTomlSchema
+import com.demonwav.mcdev.toml.toml.TomlKeyInsertionHandler
 import com.demonwav.mcdev.util.isAncestorOf
 import com.intellij.codeInsight.completion.CompletionContributor
 import com.intellij.codeInsight.completion.CompletionParameters
@@ -88,6 +89,7 @@ object ModsTomlKeyCompletionProvider : CompletionProvider<CompletionParameters>(
 
         val keySegment = parameters.position.parent as? TomlKeySegment ?: return
         val key = keySegment.parent as? TomlKey ?: return
+        val keyValue = key.parent as? TomlKeyValue ?: return
         val table = key.parentOfType<TomlKeyValueOwner>()
         val variants: Collection<TomlSchemaEntry> = when (val parent = key.parent) {
             is TomlTableHeader -> {
@@ -118,7 +120,9 @@ object ModsTomlKeyCompletionProvider : CompletionProvider<CompletionParameters>(
             else -> return
         }
 
-        result.addAllElements(variants.map { entry -> LookupElementBuilder.create(entry, entry.key) })
+        result.addAllElements(variants.map { entry ->
+            LookupElementBuilder.create(entry, entry.key).withInsertHandler(TomlKeyInsertionHandler(keyValue))
+        })
     }
 }
 

src/main/kotlin/toml/toml/TomlKeyInsertionHandler.kt

@@ -0,0 +1,23 @@
+package com.demonwav.mcdev.toml.toml
+
+import com.intellij.codeInsight.AutoPopupController
+import com.intellij.codeInsight.completion.InsertHandler
+import com.intellij.codeInsight.completion.InsertionContext
+import com.intellij.codeInsight.lookup.LookupElement
+import com.intellij.psi.PsiDocumentManager
+import org.toml.lang.psi.TomlElementTypes
+import org.toml.lang.psi.TomlKeyValue
+import org.toml.lang.psi.ext.elementType
+
+/** Inserts `=` after the completed key if missing and invokes the completion popup for the value automatically */
+class TomlKeyInsertionHandler(private val keyValue: TomlKeyValue) : InsertHandler<LookupElement> {
+    override fun handleInsert(context: InsertionContext, item: LookupElement) {
+        val hasEq = keyValue.children.any { it.elementType == TomlElementTypes.EQ }
+        if (!hasEq) {
+            context.document.insertString(context.tailOffset, " = ")
+            PsiDocumentManager.getInstance(context.project).commitDocument(context.document)
+            context.editor.caretModel.moveToOffset(context.tailOffset) // The tail offset is tracked automatically
+            AutoPopupController.getInstance(context.project).scheduleAutoPopup(context.editor);
+        }
+    }
+}