User: joe
Date: 22 Jan 24 21:07
Revision: d6289a22bd61582e1097556a315d571c59a6e765
Summary:
Add ability to do custom completion actions on injection point completion


TeamCity URL: https://ci.mcdev.io/viewModification.html?tab=vcsModificationFiles&modId=9021&personal=false

Index: src/main/kotlin/platform/mixin/handlers/injectionPoint/InjectionPoint.kt
===================================================================
--- src/main/kotlin/platform/mixin/handlers/injectionPoint/InjectionPoint.kt	(revision 182ebaaf8edfcbb165debf16110a1053a904dd0f)
+++ src/main/kotlin/platform/mixin/handlers/injectionPoint/InjectionPoint.kt	(revision d6289a22bd61582e1097556a315d571c59a6e765)
@@ -36,6 +36,7 @@
 import com.demonwav.mcdev.util.shortName
 import com.intellij.codeInsight.completion.JavaLookupElementBuilder
 import com.intellij.codeInsight.lookup.LookupElementBuilder
+import com.intellij.openapi.editor.Editor
 import com.intellij.openapi.extensions.RequiredElement
 import com.intellij.openapi.project.Project
 import com.intellij.openapi.util.KeyedExtensionCollector
@@ -48,6 +49,7 @@
 import com.intellij.psi.PsiEnumConstant
 import com.intellij.psi.PsiExpression
 import com.intellij.psi.PsiLambdaExpression
+import com.intellij.psi.PsiLiteral
 import com.intellij.psi.PsiMember
 import com.intellij.psi.PsiMethod
 import com.intellij.psi.PsiMethodReferenceExpression
@@ -76,6 +78,9 @@
 
     open fun usesMemberReference() = false
 
+    open fun onCompleted(editor: Editor, reference: PsiLiteral) {
+    }
+
     abstract fun createNavigationVisitor(
         at: PsiAnnotation,
         target: MixinSelector?,
Index: src/main/kotlin/platform/mixin/reference/InjectionPointReference.kt
===================================================================
--- src/main/kotlin/platform/mixin/reference/InjectionPointReference.kt	(revision 182ebaaf8edfcbb165debf16110a1053a904dd0f)
+++ src/main/kotlin/platform/mixin/reference/InjectionPointReference.kt	(revision d6289a22bd61582e1097556a315d571c59a6e765)
@@ -20,6 +20,7 @@
 
 package com.demonwav.mcdev.platform.mixin.reference
 
+import com.demonwav.mcdev.platform.mixin.handlers.injectionPoint.InjectionPoint
 import com.demonwav.mcdev.platform.mixin.util.MixinConstants.Annotations.AT
 import com.demonwav.mcdev.platform.mixin.util.MixinConstants.Annotations.AT_CODE
 import com.demonwav.mcdev.platform.mixin.util.MixinConstants.Annotations.SLICE
@@ -88,20 +89,28 @@
             getAllAtCodes(context.project).keys.asSequence()
                 .map {
                     PrioritizedLookupElement.withPriority(
-                        LookupElementBuilder.create(it).completeToLiteral(context),
+                        LookupElementBuilder.create(it).completeInjectionPoint(context),
                         1.0,
                     )
                 } +
                 getCustomInjectionPointInheritors(context.project).asSequence()
                     .map {
                         PrioritizedLookupElement.withPriority(
-                            LookupElementBuilder.create(it).completeToLiteral(context),
+                            LookupElementBuilder.create(it).completeInjectionPoint(context),
                             0.0,
                         )
                     }
             ).toTypedArray()
     }
 
+    private fun LookupElementBuilder.completeInjectionPoint(context: PsiElement): LookupElementBuilder {
+        val injectionPoint = InjectionPoint.byAtCode(lookupString) ?: return completeToLiteral(context)
+
+        return completeToLiteral(context) { editor, element ->
+            injectionPoint.onCompleted(editor, element)
+        }
+    }
+
     private val SLICE_SELECTORS_KEY = Key<CachedValue<List<String>>>("mcdev.sliceSelectors")
 
     private fun getSliceSelectors(project: Project): List<String> {
Index: src/main/kotlin/util/psi-utils.kt
===================================================================
--- src/main/kotlin/util/psi-utils.kt	(revision 182ebaaf8edfcbb165debf16110a1053a904dd0f)
+++ src/main/kotlin/util/psi-utils.kt	(revision d6289a22bd61582e1097556a315d571c59a6e765)
@@ -50,6 +50,7 @@
 import com.intellij.psi.PsiMethodReferenceExpression
 import com.intellij.psi.PsiModifier
 import com.intellij.psi.PsiModifier.ModifierConstant
+import com.intellij.psi.PsiModifierList
 import com.intellij.psi.PsiParameter
 import com.intellij.psi.PsiParameterList
 import com.intellij.psi.PsiReference
@@ -79,6 +80,8 @@
 
 fun PsiElement.findContainingMethod(): PsiMethod? = findParent(resolveReferences = false) { it is PsiClass }
 
+fun PsiElement.findContainingModifierList(): PsiModifierList? = findParent(resolveReferences = false) { it is PsiClass }
+
 private val PsiElement.ancestors: Sequence<PsiElement>
     get() = generateSequence(this) { if (it is PsiFile) null else it.parent }
 
Index: src/main/kotlin/util/reference/ReferenceResolver.kt
===================================================================
--- src/main/kotlin/util/reference/ReferenceResolver.kt	(revision 182ebaaf8edfcbb165debf16110a1053a904dd0f)
+++ src/main/kotlin/util/reference/ReferenceResolver.kt	(revision d6289a22bd61582e1097556a315d571c59a6e765)
@@ -25,7 +25,7 @@
 import com.intellij.openapi.command.CommandProcessor
 import com.intellij.openapi.editor.Editor
 import com.intellij.psi.JavaPsiFacade
-import com.intellij.psi.PsiAnnotationMemberValue
+import com.intellij.psi.PsiArrayInitializerMemberValue
 import com.intellij.psi.PsiDocumentManager
 import com.intellij.psi.PsiElement
 import com.intellij.psi.PsiExpression
@@ -109,11 +109,10 @@
     var current: PsiElement
     var parent = this
 
-    @Suppress("KotlinConstantConditions") // kotlin is wrong
     do {
         current = parent
         parent = current.parent
-        if (parent is PsiNameValuePair || parent is PsiAnnotationMemberValue) {
+        if (parent is PsiNameValuePair || parent is PsiArrayInitializerMemberValue) {
             return current
         }
     } while (parent is PsiExpression)
@@ -125,8 +124,11 @@
  * Patches the [LookupElementBuilder] to replace the element with a single
  * [PsiLiteral] when using code completion.
  */
-fun LookupElementBuilder.completeToLiteral(context: PsiElement): LookupElementBuilder {
-    if (context is PsiLiteral) {
+fun LookupElementBuilder.completeToLiteral(
+    context: PsiElement,
+    extraAction: ((Editor, PsiLiteral) -> Unit)? = null
+): LookupElementBuilder {
+    if (context is PsiLiteral && extraAction == null) {
         // Context is already a literal
         return this
     }
@@ -135,7 +137,7 @@
     // not sure how you would keep line breaks after completion
     return withInsertHandler { insertionContext, item ->
         insertionContext.laterRunnable =
-            ReplaceElementWithLiteral(insertionContext.editor, insertionContext.file, item.lookupString)
+            ReplaceElementWithLiteral(insertionContext.editor, insertionContext.file, item.lookupString, extraAction)
     }
 }
 
@@ -143,6 +145,7 @@
     private val editor: Editor,
     private val file: PsiFile,
     private val text: String,
+    private val extraAction: ((Editor, PsiLiteral) -> Unit)?
 ) : Runnable {
 
     override fun run() {
@@ -153,13 +156,17 @@
         CommandProcessor.getInstance().runUndoTransparentAction {
             runWriteAction {
                 val element = file.findElementAt(editor.caretModel.offset)!!.findContextElement()
-                element.replace(
+                val newElement = element.replace(
                     JavaPsiFacade.getElementFactory(element.project).createExpressionFromText(
                         "\"$text\"",
                         element.parent,
                     ),
-                )
+                ) as PsiLiteral
+                val extraAction = this.extraAction
+                if (extraAction != null) {
+                    extraAction(editor, newElement)
-            }
-        }
-    }
-}
+                }
+            }
+        }
+    }
+}