Add dynamic selectors to method completion (#2617)

* add dynamic selector to method completion

* Update src/main/kotlin/platform/mixin/reference/MixinSelectors.kt

Co-authored-by: Joseph Burton <burtonjae@hotmail.co.uk>

---------

Co-authored-by: Joseph Burton <burtonjae@hotmail.co.uk>

Author: Bawnorton

src/main/kotlin/platform/mixin/reference/AbstractMethodReference.kt

@@ -21,6 +21,7 @@
 package com.demonwav.mcdev.platform.mixin.reference
 
 import com.demonwav.mcdev.platform.mixin.handlers.MixinAnnotationHandler
+import com.demonwav.mcdev.platform.mixin.handlers.injectionPoint.InjectionPoint
 import com.demonwav.mcdev.platform.mixin.reference.target.TargetReference
 import com.demonwav.mcdev.platform.mixin.util.ClassAndMethodNode
 import com.demonwav.mcdev.platform.mixin.util.bytecode
@@ -39,6 +40,7 @@ import com.demonwav.mcdev.util.reference.completeToLiteral
 import com.demonwav.mcdev.util.toResolveResults
 import com.demonwav.mcdev.util.toTypedArray
 import com.intellij.codeInsight.completion.JavaLookupElementBuilder
+import com.intellij.codeInsight.completion.PrioritizedLookupElement
 import com.intellij.codeInsight.lookup.LookupElementBuilder
 import com.intellij.psi.PsiAnnotation
 import com.intellij.psi.PsiArrayInitializerMemberValue
@@ -193,7 +195,16 @@ abstract class AbstractMethodReference : PolyReferenceResolver(), MixinReference
             }
         }
 
-        return createLookup(context, methods.asSequence().map { ClassAndMethodNode(target, it) }, uniqueMethods)
+        val dynamicSelectors = MixinSelectorParser.EP_NAME.extensionList.asSequence()
+            .filterIsInstance<DynamicSelectorParser>()
+            .map { "@${it.id}" }
+
+        return createLookup(
+            context,
+            methods.asSequence().map { ClassAndMethodNode(target, it) },
+            uniqueMethods,
+            dynamicSelectors
+        )
     }
 
     private fun collectVariants(context: PsiElement, targets: Collection<ClassNode>): Array<Any> {
@@ -227,13 +238,18 @@ abstract class AbstractMethodReference : PolyReferenceResolver(), MixinReference
             }
         }
 
-        return createLookup(context, allMethods.asSequence(), uniqueMethods)
+        val dynamicSelectors = MixinSelectorParser.EP_NAME.extensionList.asSequence()
+            .filterIsInstance<DynamicSelectorParser>()
+            .map { "@${it.id}" }
+
+        return createLookup(context, allMethods.asSequence(), uniqueMethods, dynamicSelectors)
     }
 
     private fun createLookup(
         context: PsiElement,
         methods: Sequence<ClassAndMethodNode>,
         uniqueMethods: Set<String>,
+        dynamicSelectors: Sequence<String>
     ): Array<Any> {
         return methods
             .map { m ->
@@ -256,8 +272,17 @@ abstract class AbstractMethodReference : PolyReferenceResolver(), MixinReference
                     null,
                 )
                     .withPresentableText(m.method.name)
-                addCompletionInfo(builder, context, targetMethodInfo)
-            }.toTypedArray()
+                PrioritizedLookupElement.withPriority(
+                    addCompletionInfo(builder, context, targetMethodInfo),
+                    1.0
+                )
+            }.plus(dynamicSelectors.map {
+                PrioritizedLookupElement.withPriority(
+                    LookupElementBuilder.create(it).completeDynamicSelector(context),
+                    0.0
+                )
+            })
+            .toTypedArray()
     }
 
     open val requireDescriptor = false
@@ -269,4 +294,15 @@ abstract class AbstractMethodReference : PolyReferenceResolver(), MixinReference
     ): LookupElementBuilder {
         return builder.completeToLiteral(context)
     }
+
+    private fun LookupElementBuilder.completeDynamicSelector(context: PsiElement): LookupElementBuilder {
+        val id = lookupString.removePrefix("@")
+        val parser = MixinSelectorParser.EP_NAME.extensionList.asSequence()
+            .filterIsInstance<DynamicSelectorParser>()
+            .firstOrNull { id in it.validIds } ?: return completeToLiteral(context)
+
+        return completeToLiteral(context) { editor, element ->
+            parser.onCompleted(editor, element)
+        }
+    }
 }

src/main/kotlin/platform/mixin/reference/MixinSelectors.kt

@@ -36,6 +36,7 @@ import com.demonwav.mcdev.util.constantStringValue
 import com.demonwav.mcdev.util.descriptor
 import com.demonwav.mcdev.util.findAnnotation
 import com.demonwav.mcdev.util.findContainingClass
+import com.demonwav.mcdev.util.findContainingModifierList
 import com.demonwav.mcdev.util.findField
 import com.demonwav.mcdev.util.findMethods
 import com.demonwav.mcdev.util.findQualifiedClass
@@ -45,6 +46,7 @@ import com.demonwav.mcdev.util.mapToArray
 import com.demonwav.mcdev.util.resolveClass
 import com.demonwav.mcdev.util.resolveType
 import com.demonwav.mcdev.util.resolveTypeArray
+import com.intellij.openapi.editor.Editor
 import com.intellij.openapi.extensions.ExtensionPointName
 import com.intellij.openapi.project.Project
 import com.intellij.openapi.util.RecursionManager
@@ -55,10 +57,14 @@ import com.intellij.psi.PsiCallExpression
 import com.intellij.psi.PsiClass
 import com.intellij.psi.PsiElement
 import com.intellij.psi.PsiField
+import com.intellij.psi.PsiLiteral
 import com.intellij.psi.PsiMember
 import com.intellij.psi.PsiMethod
+import com.intellij.psi.PsiModifierList
 import com.intellij.psi.PsiNameValuePair
 import com.intellij.psi.PsiTypes
+import com.intellij.psi.codeStyle.CodeStyleManager
+import com.intellij.psi.codeStyle.JavaCodeStyleManager
 import com.intellij.psi.search.GlobalSearchScope
 import com.intellij.psi.search.searches.AnnotatedMembersSearch
 import com.intellij.psi.search.searches.MethodReferencesSearch
@@ -418,7 +424,7 @@ private fun findNamespace(
 
 private val DYNAMIC_SELECTOR_PATTERN = "(?i)^@([a-z]+(:[a-z]+)?)(\\((.*)\\))?$".toRegex()
 
-abstract class DynamicSelectorParser(id: String, vararg aliases: String) : MixinSelectorParser {
+abstract class DynamicSelectorParser(val id: String, vararg aliases: String) : MixinSelectorParser {
     val validIds = aliases.toSet() + id
 
     final override fun parse(value: String, context: PsiElement): MixinSelector? {
@@ -431,6 +437,9 @@ abstract class DynamicSelectorParser(id: String, vararg aliases: String) : Mixin
     }
 
     abstract fun parseDynamic(args: String, context: PsiElement): MixinSelector?
+
+    open fun onCompleted(editor: Editor, reference: PsiLiteral) {
+    }
 }
 
 // @Desc
@@ -546,6 +555,32 @@ class DescSelectorParser : DynamicSelectorParser("Desc", "mixin:Desc") {
         }
     }
 
+    override fun onCompleted(editor: Editor, reference: PsiLiteral) {
+        val modifierList = reference.findContainingModifierList() ?: return
+        if (modifierList.hasAnnotation(DESC)) {
+            return
+        }
+
+        val project = reference.project
+
+        val descAnnotation = modifierList.addAfter(
+            JavaPsiFacade.getElementFactory(project)
+                .createAnnotationFromText("@${DESC}(\"\")", reference),
+            null
+        )
+
+        // add imports and reformat
+        JavaCodeStyleManager.getInstance(project).shortenClassReferences(descAnnotation)
+        JavaCodeStyleManager.getInstance(project).optimizeImports(modifierList.containingFile)
+        val formattedModifierList = CodeStyleManager.getInstance(project).reformat(modifierList) as PsiModifierList
+
+        // move the caret to @Desc("<caret>")
+        val formattedDescAnnotation = formattedModifierList.findAnnotation(DESC)
+            ?: return
+        val descLiteral = formattedDescAnnotation.findDeclaredAttributeValue(null) ?: return
+        editor.caretModel.moveToOffset(descLiteral.textRange.startOffset + 1)
+    }
+
     object Util {
         fun descSelectorFromAnnotation(descAnnotation: PsiAnnotation): DescSelector? {
             val explicitOwner = descAnnotation.findAttributeValue("owner")