User: llamalad7
Date: 23 Sep 23 16:20
Revision: 449a4b79ff378b209b2be82c6dc8e53d8a7f2c8a
Summary:
MixinExtras support (#2127)

* New: Support for MixinExtras injectors.

* Fix: Enable mixin autocomplete popup in all handled annotations even if they're not from the Mixin package.

* Fix: Support class constants in ConstantInjectionPoint.

* Improvement: Unify Mixin `at` handling and allow for different `at` keys.

* Checkstyle and License

* Improvement: Extract injection point annotations to an extension point.

* Implement feedback.

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

Index: src/main/kotlin/platform/mixin/completion/MixinCompletionConfidence.kt
===================================================================
--- src/main/kotlin/platform/mixin/completion/MixinCompletionConfidence.kt	(revision 9ae67da17c7109ef7ce9fd12290798cab82e1d02)
+++ src/main/kotlin/platform/mixin/completion/MixinCompletionConfidence.kt	(revision 449a4b79ff378b209b2be82c6dc8e53d8a7f2c8a)
@@ -20,6 +20,7 @@
 
 package com.demonwav.mcdev.platform.mixin.completion
 
+import com.demonwav.mcdev.platform.mixin.handlers.MixinAnnotationHandler
 import com.demonwav.mcdev.platform.mixin.util.MixinConstants
 import com.intellij.codeInsight.completion.CompletionConfidence
 import com.intellij.codeInsight.completion.SkipAutopopupInStrings
@@ -35,7 +36,13 @@
     private val mixinAnnotation = PlatformPatterns.psiElement()
         .inside(
             false,
-            PsiJavaPatterns.psiAnnotation().qName(StandardPatterns.string().startsWith(MixinConstants.PACKAGE)),
+            PsiJavaPatterns.psiAnnotation().qName(
+                StandardPatterns.or(
+                    StandardPatterns.string().startsWith(MixinConstants.PACKAGE),
+                    StandardPatterns.string()
+                        .oneOf(MixinAnnotationHandler.getBuiltinHandlers().map { it.first }.toList()),
+                )
+            ),
             PlatformPatterns.psiFile(),
         )!!
 
Index: src/main/kotlin/platform/mixin/handlers/InjectorAnnotationHandler.kt
===================================================================
--- src/main/kotlin/platform/mixin/handlers/InjectorAnnotationHandler.kt	(revision 9ae67da17c7109ef7ce9fd12290798cab82e1d02)
+++ src/main/kotlin/platform/mixin/handlers/InjectorAnnotationHandler.kt	(revision 449a4b79ff378b209b2be82c6dc8e53d8a7f2c8a)
@@ -81,12 +81,14 @@
         }.reduceOrNull(InsnResolutionInfo.Failure::combine) ?: InsnResolutionInfo.Failure()
     }
 
+    open fun getAtKey(annotation: PsiAnnotation): String = "at"
+
     protected open fun isUnresolved(
         annotation: PsiAnnotation,
         targetClass: ClassNode,
         targetMethod: MethodNode,
     ): InsnResolutionInfo.Failure? {
-        return annotation.findAttributeValue("at")?.findAnnotations()
+        return annotation.findAttributeValue(getAtKey(annotation))?.findAnnotations()
             .ifNullOrEmpty { return InsnResolutionInfo.Failure() }!!
             .firstNotNullOfOrNull { AtResolver(it, targetClass, targetMethod).isUnresolved() }
     }
@@ -103,7 +105,7 @@
         targetClass: ClassNode,
         targetMethod: MethodNode,
     ): List<PsiElement> {
-        return annotation.findAttributeValue("at")?.findAnnotations()
+        return annotation.findAttributeValue(getAtKey(annotation))?.findAnnotations()
             .ifNullOrEmpty { return emptyList() }!!
             .flatMap { AtResolver(it, targetClass, targetMethod).resolveNavigationTargets() }
     }
@@ -129,7 +131,7 @@
         targetMethod: MethodNode,
         mode: CollectVisitor.Mode = CollectVisitor.Mode.MATCH_ALL,
     ): List<CollectVisitor.Result<*>> {
-        return annotation.findAttributeValue("at")?.findAnnotations()
+        return annotation.findAttributeValue(getAtKey(annotation))?.findAnnotations()
             .ifNullOrEmpty { return emptyList() }!!
             .flatMap { AtResolver(it, targetClass, targetMethod).resolveInstructions(mode) }
     }
Index: src/main/kotlin/platform/mixin/handlers/ModifyConstantHandler.kt
===================================================================
--- src/main/kotlin/platform/mixin/handlers/ModifyConstantHandler.kt	(revision 9ae67da17c7109ef7ce9fd12290798cab82e1d02)
+++ src/main/kotlin/platform/mixin/handlers/ModifyConstantHandler.kt	(revision 449a4b79ff378b209b2be82c6dc8e53d8a7f2c8a)
@@ -20,28 +20,15 @@
 
 package com.demonwav.mcdev.platform.mixin.handlers
 
-import com.demonwav.mcdev.platform.mixin.handlers.injectionPoint.CollectVisitor
 import com.demonwav.mcdev.platform.mixin.handlers.injectionPoint.ConstantInjectionPoint
 import com.demonwav.mcdev.platform.mixin.handlers.injectionPoint.InjectionPoint
-import com.demonwav.mcdev.platform.mixin.handlers.injectionPoint.InsnResolutionInfo
 import com.demonwav.mcdev.platform.mixin.inspection.injector.MethodSignature
 import com.demonwav.mcdev.platform.mixin.inspection.injector.ParameterGroup
-import com.demonwav.mcdev.platform.mixin.util.MixinConstants.Classes.CONSTANT_CONDITION
-import com.demonwav.mcdev.platform.mixin.util.findSourceElement
-import com.demonwav.mcdev.util.constantValue
-import com.demonwav.mcdev.util.descriptor
 import com.demonwav.mcdev.util.findAnnotations
-import com.demonwav.mcdev.util.fullQualifiedName
-import com.demonwav.mcdev.util.parseArray
-import com.demonwav.mcdev.util.resolveClass
 import com.intellij.psi.PsiAnnotation
-import com.intellij.psi.PsiElement
-import com.intellij.psi.PsiEnumConstant
 import com.intellij.psi.PsiManager
 import com.intellij.psi.PsiMethod
-import com.intellij.psi.PsiReferenceExpression
 import com.intellij.psi.PsiType
-import com.intellij.psi.search.GlobalSearchScope
 import com.intellij.psi.util.parentOfType
 import org.objectweb.asm.Opcodes
 import org.objectweb.asm.Type
@@ -74,70 +61,18 @@
         Opcodes.IFLE,
     )
 
-    private class ModifyConstantInfo(
-        val constantInfo: ConstantInjectionPoint.ConstantInfo,
-        val constantAnnotation: PsiAnnotation,
-    )
-
-    private fun getConstantInfos(modifyConstant: PsiAnnotation): List<ModifyConstantInfo>? {
+    private fun getConstantInfos(modifyConstant: PsiAnnotation): List<ConstantInjectionPoint.ConstantInfo>? {
         val constants = modifyConstant.findDeclaredAttributeValue("constant")
             ?.findAnnotations()
             ?.takeIf { it.isNotEmpty() }
             ?: return null
         return constants.map { constant ->
-            val nullValue = constant.findDeclaredAttributeValue("nullValue")?.constantValue as? Boolean ?: false
-            val intValue = (constant.findDeclaredAttributeValue("intValue")?.constantValue as? Number)?.toInt()
-            val floatValue = (constant.findDeclaredAttributeValue("floatValue")?.constantValue as? Number)?.toFloat()
-            val longValue = (constant.findDeclaredAttributeValue("longValue")?.constantValue as? Number)?.toLong()
-            val doubleValue = (constant.findDeclaredAttributeValue("doubleValue")?.constantValue as? Number)?.toDouble()
-            val stringValue = constant.findDeclaredAttributeValue("stringValue")?.constantValue as? String
-            val classValue = constant.findDeclaredAttributeValue("classValue")?.resolveClass()?.descriptor?.let {
-                Type.getType(
-                    it,
-                )
+            (InjectionPoint.byAtCode("CONSTANT") as ConstantInjectionPoint).getConstantInfo(constant) ?: return null
-            }
+        }
-
-            fun Boolean.toInt(): Int {
-                return if (this) 1 else 0
-            }
+    }
 
-            val count = nullValue.toInt() +
-                (intValue != null).toInt() +
-                (floatValue != null).toInt() +
-                (longValue != null).toInt() +
-                (doubleValue != null).toInt() +
-                (stringValue != null).toInt() +
-                (classValue != null).toInt()
-            if (count != 1) {
-                return null
-            }
+    override fun getAtKey(annotation: PsiAnnotation) = "constant"
 
-            val value = if (nullValue) {
-                null
-            } else {
-                intValue ?: floatValue ?: longValue ?: doubleValue ?: stringValue ?: classValue
-            }
-
-            val expandConditions = constant.findDeclaredAttributeValue("expandZeroConditions")?.parseArray {
-                if (it !is PsiReferenceExpression) {
-                    return@parseArray null
-                }
-                val field = it.resolve() as? PsiEnumConstant ?: return@parseArray null
-                val enumClass = field.containingClass ?: return@parseArray null
-                if (enumClass.fullQualifiedName != CONSTANT_CONDITION) {
-                    return@parseArray null
-                }
-                try {
-                    ConstantInjectionPoint.ExpandCondition.valueOf(field.name)
-                } catch (e: IllegalArgumentException) {
-                    null
-                }
-            }?.toSet() ?: emptySet()
-
-            ModifyConstantInfo(ConstantInjectionPoint.ConstantInfo(value, expandConditions), constant)
-        }
-    }
-
     override fun expectedMethodSignature(
         annotation: PsiAnnotation,
         targetClass: ClassNode,
@@ -167,7 +102,7 @@
 
         val psiManager = PsiManager.getInstance(annotation.project)
         return constantInfos.asSequence().map {
-            when (it.constantInfo.constant) {
+            when (it.constant) {
                 null -> PsiType.getJavaLangObject(psiManager, annotation.resolveScope)
                 is Int -> PsiType.INT
                 is Float -> PsiType.FLOAT
@@ -175,7 +110,7 @@
                 is Double -> PsiType.DOUBLE
                 is String -> PsiType.getJavaLangString(psiManager, annotation.resolveScope)
                 is Type -> PsiType.getJavaLangClass(psiManager, annotation.resolveScope)
-                else -> throw IllegalStateException("Unknown constant type: ${it.constantInfo.constant.javaClass.name}")
+                else -> throw IllegalStateException("Unknown constant type: ${it.constant.javaClass.name}")
             }
         }.distinct().map { type ->
             MethodSignature(
@@ -192,154 +127,6 @@
         }.toList()
     }
 
-    override fun resolveForNavigation(
-        annotation: PsiAnnotation,
-        targetClass: ClassNode,
-        targetMethod: MethodNode,
-    ): List<PsiElement> {
-        val targetElement = targetMethod.findSourceElement(
-            targetClass,
-            annotation.project,
-            GlobalSearchScope.allScope(annotation.project),
-            canDecompile = true,
-        ) ?: return emptyList()
-
-        val constantInfos = getConstantInfos(annotation)
-        if (constantInfos == null) {
-            val returnType = annotation.parentOfType<PsiMethod>()?.returnType
-                ?: return emptyList()
-
-            val collectVisitor = ConstantInjectionPoint.MyCollectVisitor(
-                annotation.project,
-                CollectVisitor.Mode.MATCH_ALL,
-                null,
-                Type.getType(returnType.descriptor)
-            )
-            collectVisitor.visit(targetMethod)
-            val bytecodeResults = collectVisitor.result
-
-            val navigationVisitor = ConstantInjectionPoint.MyNavigationVisitor(
-                null,
-                Type.getType(returnType.descriptor)
-            )
-            targetElement.accept(navigationVisitor)
-
-            return bytecodeResults.asSequence().mapNotNull { bytecodeResult ->
-                navigationVisitor.result.getOrNull(bytecodeResult.index)
-            }.sortedBy { it.textOffset }.toList()
-        }
-
-        val constantInjectionPoint = InjectionPoint.byAtCode("CONSTANT") as? ConstantInjectionPoint
-            ?: return emptyList()
-
-        return constantInfos.asSequence().flatMap { modifyConstantInfo ->
-            val collectVisitor = ConstantInjectionPoint.MyCollectVisitor(
-                annotation.project,
-                CollectVisitor.Mode.MATCH_ALL,
-                modifyConstantInfo.constantInfo,
-            )
-            constantInjectionPoint.addStandardFilters(
-                modifyConstantInfo.constantAnnotation,
-                targetClass,
-                collectVisitor,
-            )
-            collectVisitor.visit(targetMethod)
-            val bytecodeResults = collectVisitor.result
-
-            val navigationVisitor = ConstantInjectionPoint.MyNavigationVisitor(modifyConstantInfo.constantInfo)
-            targetElement.accept(navigationVisitor)
-
-            bytecodeResults.asSequence().mapNotNull { bytecodeResult ->
-                navigationVisitor.result.getOrNull(bytecodeResult.index)
-            }
-        }.sortedBy { it.textOffset }.toList()
-    }
-
-    override fun resolveInstructions(
-        annotation: PsiAnnotation,
-        targetClass: ClassNode,
-        targetMethod: MethodNode,
-        mode: CollectVisitor.Mode,
-    ): List<CollectVisitor.Result<*>> {
-        val constantInfos = getConstantInfos(annotation)
-        if (constantInfos == null) {
-            val returnType = annotation.parentOfType<PsiMethod>()?.returnType
-                ?: return emptyList()
-
-            val collectVisitor = ConstantInjectionPoint.MyCollectVisitor(
-                annotation.project,
-                mode,
-                null,
-                Type.getType(returnType.descriptor)
-            )
-            collectVisitor.visit(targetMethod)
-            return collectVisitor.result.sortedBy { targetMethod.instructions.indexOf(it.insn) }
-        }
-
-        val constantInjectionPoint = InjectionPoint.byAtCode("CONSTANT") as? ConstantInjectionPoint
-            ?: return emptyList()
-        return constantInfos.asSequence().flatMap { modifyConstantInfo ->
-            val collectVisitor = ConstantInjectionPoint.MyCollectVisitor(
-                annotation.project,
-                mode,
-                modifyConstantInfo.constantInfo,
-            )
-            constantInjectionPoint.addStandardFilters(
-                modifyConstantInfo.constantAnnotation,
-                targetClass,
-                collectVisitor,
-            )
-            collectVisitor.visit(targetMethod)
-            collectVisitor.result.asSequence()
-        }.sortedBy { targetMethod.instructions.indexOf(it.insn) }.toList()
-    }
-
-    override fun isUnresolved(
-        annotation: PsiAnnotation,
-        targetClass: ClassNode,
-        targetMethod: MethodNode,
-    ): InsnResolutionInfo.Failure? {
-        val constantInfos = getConstantInfos(annotation)
-        if (constantInfos == null) {
-            val returnType = annotation.parentOfType<PsiMethod>()?.returnType
-                ?: return InsnResolutionInfo.Failure()
-
-            val collectVisitor = ConstantInjectionPoint.MyCollectVisitor(
-                annotation.project,
-                CollectVisitor.Mode.MATCH_FIRST,
-                null,
-                Type.getType(returnType.descriptor)
-            )
-            collectVisitor.visit(targetMethod)
-            return if (collectVisitor.result.isEmpty()) {
-                InsnResolutionInfo.Failure(collectVisitor.filterToBlame)
-            } else {
-                null
-            }
-        }
-
-        val constantInjectionPoint = InjectionPoint.byAtCode("CONSTANT") as? ConstantInjectionPoint
-            ?: return null
-        return constantInfos.firstNotNullOfOrNull { modifyConstantInfo ->
-            val collectVisitor = ConstantInjectionPoint.MyCollectVisitor(
-                annotation.project,
-                CollectVisitor.Mode.MATCH_FIRST,
-                modifyConstantInfo.constantInfo,
-            )
-            constantInjectionPoint.addStandardFilters(
-                modifyConstantInfo.constantAnnotation,
-                targetClass,
-                collectVisitor,
-            )
-            collectVisitor.visit(targetMethod)
-            if (collectVisitor.result.isEmpty()) {
-                InsnResolutionInfo.Failure(collectVisitor.filterToBlame)
-            } else {
-                null
-            }
-        }
-    }
-
     override fun isInsnAllowed(insn: AbstractInsnNode): Boolean {
         return insn.opcode in allowedOpcodes
     }
Index: src/main/kotlin/platform/mixin/handlers/injectionPoint/AtResolver.kt
===================================================================
--- src/main/kotlin/platform/mixin/handlers/injectionPoint/AtResolver.kt	(revision 9ae67da17c7109ef7ce9fd12290798cab82e1d02)
+++ src/main/kotlin/platform/mixin/handlers/injectionPoint/AtResolver.kt	(revision 449a4b79ff378b209b2be82c6dc8e53d8a7f2c8a)
@@ -27,14 +27,19 @@
 import com.demonwav.mcdev.platform.mixin.util.findSourceElement
 import com.demonwav.mcdev.util.computeStringArray
 import com.demonwav.mcdev.util.constantStringValue
+import com.demonwav.mcdev.util.constantValue
+import com.demonwav.mcdev.util.fullQualifiedName
 import com.intellij.codeInsight.lookup.LookupElementBuilder
 import com.intellij.psi.PsiAnnotation
+import com.intellij.psi.PsiAnnotationMemberValue
+import com.intellij.psi.PsiArrayInitializerMemberValue
 import com.intellij.psi.PsiClass
 import com.intellij.psi.PsiClassType
 import com.intellij.psi.PsiElement
 import com.intellij.psi.PsiExpression
 import com.intellij.psi.PsiQualifiedReference
 import com.intellij.psi.PsiReference
+import com.intellij.psi.PsiReferenceExpression
 import com.intellij.psi.search.GlobalSearchScope
 import com.intellij.psi.util.parentOfType
 import org.objectweb.asm.tree.ClassNode
@@ -69,7 +74,8 @@
 ) {
     companion object {
         private fun getInjectionPoint(at: PsiAnnotation): InjectionPoint<*>? {
-            var atCode = at.findDeclaredAttributeValue("value")?.constantStringValue ?: return null
+            var atCode = at.qualifiedName?.let { InjectionPointAnnotation.atCodeFor(it) }
+                ?: at.findDeclaredAttributeValue("value")?.constantStringValue ?: return null
 
             // remove slice selector
             val isInSlice = at.parentOfType<PsiAnnotation>()?.hasQualifiedName(SLICE) ?: false
@@ -88,8 +94,8 @@
         }
 
         fun getArgs(at: PsiAnnotation): Map<String, String> {
-            val args = at.findAttributeValue("args")?.computeStringArray() ?: return emptyMap()
-            return args.asSequence()
+            val args = at.findAttributeValue("args")?.computeStringArray().orEmpty()
+            val explicitArgs = args.asSequence()
                 .map {
                     val parts = it.split('=', limit = 2)
                     if (parts.size == 1) {
@@ -99,9 +105,35 @@
                     }
                 }
                 .toMap()
+            return getInherentArgs(at) + explicitArgs
         }
+
+        private fun getInherentArgs(at: PsiAnnotation): Map<String, String> {
+            return at.attributes.asSequence()
+                .mapNotNull {
+                    val name = it.attributeName
+                    val value = at.findAttributeValue(name) ?: return@mapNotNull null
+                    val string = valueToString(value) ?: return@mapNotNull null
+                    name to string
-    }
+                }
+                .toMap()
+        }
 
+        private fun valueToString(value: PsiAnnotationMemberValue): String? {
+            if (value is PsiArrayInitializerMemberValue) {
+                return value.initializers.map { valueToString(it) ?: return null }.joinToString(",")
+            }
+            return when (val constant = value.constantValue) {
+                is PsiClassType -> constant.fullQualifiedName?.replace('.', '/')
+                null -> when (value) {
+                    is PsiReferenceExpression -> value.referenceName
+                    else -> null
+                }
+                else -> constant.toString()
+            }
+        }
+    }
+
     fun isUnresolved(): InsnResolutionInfo.Failure? {
         val injectionPoint = getInjectionPoint(at)
             ?: return null // we don't know what to do with custom handlers, assume ok
Index: src/main/kotlin/platform/mixin/handlers/injectionPoint/ConstantInjectionPoint.kt
===================================================================
--- src/main/kotlin/platform/mixin/handlers/injectionPoint/ConstantInjectionPoint.kt	(revision 9ae67da17c7109ef7ce9fd12290798cab82e1d02)
+++ src/main/kotlin/platform/mixin/handlers/injectionPoint/ConstantInjectionPoint.kt	(revision 449a4b79ff378b209b2be82c6dc8e53d8a7f2c8a)
@@ -52,9 +52,10 @@
 import org.objectweb.asm.tree.LabelNode
 import org.objectweb.asm.tree.LdcInsnNode
 import org.objectweb.asm.tree.MethodNode
+import org.objectweb.asm.tree.TypeInsnNode
 
 class ConstantInjectionPoint : InjectionPoint<PsiElement>() {
-    private fun getConstantInfo(at: PsiAnnotation): ConstantInfo? {
+    fun getConstantInfo(at: PsiAnnotation): ConstantInfo? {
         val args = AtResolver.getArgs(at)
         val nullValue = args["nullValue"]?.let(java.lang.Boolean::parseBoolean) ?: false
         val intValue = args["intValue"]?.toIntOrNull()
@@ -283,6 +284,15 @@
                         0
                     }
 
+                    is TypeInsnNode -> {
+                        if (insn.opcode < Opcodes.CHECKCAST) {
+                            // Don't treat NEW and ANEWARRAY as constants
+                            // Matches Mixin's handling
+                            return@forEachRemaining
+                        }
+                        Type.getObjectType(insn.desc)
+                    }
+
                     else -> return@forEachRemaining
                 }
 
Index: src/main/kotlin/platform/mixin/handlers/injectionPoint/InjectionPointAnnotation.kt
===================================================================
--- src/main/kotlin/platform/mixin/handlers/injectionPoint/InjectionPointAnnotation.kt	(revision 449a4b79ff378b209b2be82c6dc8e53d8a7f2c8a)
+++ src/main/kotlin/platform/mixin/handlers/injectionPoint/InjectionPointAnnotation.kt	(revision 449a4b79ff378b209b2be82c6dc8e53d8a7f2c8a)
@@ -0,0 +1,50 @@
+/*
+ * Minecraft Development for IntelliJ
+ *
+ * https://mcdev.io/
+ *
+ * Copyright (C) 2023 minecraft-dev
+ *
+ * This program is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU Lesser General Public License as published
+ * by the Free Software Foundation, version 3.0 only.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU Lesser General Public License
+ * along with this program.  If not, see <https://www.gnu.org/licenses/>.
+ */
+
+package com.demonwav.mcdev.platform.mixin.handlers.injectionPoint
+
+import com.intellij.openapi.extensions.ExtensionPointName
+import com.intellij.openapi.extensions.RequiredElement
+import com.intellij.openapi.util.KeyedExtensionCollector
+import com.intellij.util.KeyedLazyInstance
+import com.intellij.util.xmlb.annotations.Attribute
+
+class InjectionPointAnnotation : KeyedLazyInstance<String> {
+    @Attribute("annotation")
+    @RequiredElement
+    lateinit var annotation: String
+
+    @Attribute("atCode")
+    @RequiredElement
+    lateinit var atCode: String
+
+    override fun getKey() = annotation
+
+    override fun getInstance() = atCode
+
+    companion object {
+        private val EP_NAME = ExtensionPointName<InjectionPointAnnotation>(
+            "com.demonwav.minecraft-dev.injectionPointAnnotation"
+        )
+        private val COLLECTOR = KeyedExtensionCollector<String, String>(EP_NAME)
+
+        fun atCodeFor(qualifiedName: String): String? = COLLECTOR.findSingle(qualifiedName)
+    }
+}
Index: src/main/kotlin/platform/mixin/handlers/mixinextras/MixinExtrasInjectorAnnotationHandler.kt
===================================================================
--- src/main/kotlin/platform/mixin/handlers/mixinextras/MixinExtrasInjectorAnnotationHandler.kt	(revision 449a4b79ff378b209b2be82c6dc8e53d8a7f2c8a)
+++ src/main/kotlin/platform/mixin/handlers/mixinextras/MixinExtrasInjectorAnnotationHandler.kt	(revision 449a4b79ff378b209b2be82c6dc8e53d8a7f2c8a)
@@ -0,0 +1,338 @@
+/*
+ * Minecraft Development for IntelliJ
+ *
+ * https://mcdev.io/
+ *
+ * Copyright (C) 2023 minecraft-dev
+ *
+ * This program is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU Lesser General Public License as published
+ * by the Free Software Foundation, version 3.0 only.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU Lesser General Public License
+ * along with this program.  If not, see <https://www.gnu.org/licenses/>.
+ */
+
+package com.demonwav.mcdev.platform.mixin.handlers.mixinextras
+
+import com.demonwav.mcdev.platform.mixin.handlers.InjectorAnnotationHandler
+import com.demonwav.mcdev.platform.mixin.inspection.injector.MethodSignature
+import com.demonwav.mcdev.platform.mixin.inspection.injector.ParameterGroup
+import com.demonwav.mcdev.platform.mixin.util.FieldTargetMember
+import com.demonwav.mcdev.platform.mixin.util.MethodTargetMember
+import com.demonwav.mcdev.platform.mixin.util.getGenericParameterTypes
+import com.demonwav.mcdev.platform.mixin.util.getGenericReturnType
+import com.demonwav.mcdev.platform.mixin.util.getGenericType
+import com.demonwav.mcdev.platform.mixin.util.toPsiType
+import com.demonwav.mcdev.util.MemberReference
+import com.demonwav.mcdev.util.Parameter
+import com.demonwav.mcdev.util.toJavaIdentifier
+import com.intellij.psi.JavaPsiFacade
+import com.intellij.psi.PsiAnnotation
+import com.intellij.psi.PsiType
+import org.objectweb.asm.Opcodes
+import org.objectweb.asm.Type
+import org.objectweb.asm.tree.AbstractInsnNode
+import org.objectweb.asm.tree.ClassNode
+import org.objectweb.asm.tree.FieldInsnNode
+import org.objectweb.asm.tree.LdcInsnNode
+import org.objectweb.asm.tree.MethodInsnNode
+import org.objectweb.asm.tree.MethodNode
+import org.objectweb.asm.tree.TypeInsnNode
+
+abstract class MixinExtrasInjectorAnnotationHandler : InjectorAnnotationHandler() {
+    open val oldSuperBehavior = false
+
+    enum class InstructionType {
+        METHOD_CALL {
+            override fun matches(insn: AbstractInsnNode) = insn is MethodInsnNode && insn.name != "<init>"
+        },
+        FIELD_GET {
+            override fun matches(insn: AbstractInsnNode) =
+                insn.opcode == Opcodes.GETFIELD || insn.opcode == Opcodes.GETSTATIC
+        },
+        FIELD_SET {
+            override fun matches(insn: AbstractInsnNode) =
+                insn.opcode == Opcodes.PUTFIELD || insn.opcode == Opcodes.PUTSTATIC
+        },
+        INSTANTIATION {
+            override fun matches(insn: AbstractInsnNode) = insn.opcode == Opcodes.NEW
+        },
+        INSTANCEOF {
+            override fun matches(insn: AbstractInsnNode) = insn.opcode == Opcodes.INSTANCEOF
+        },
+        CONSTANT {
+            override fun matches(insn: AbstractInsnNode) = isConstant(insn)
+        },
+        RETURN {
+            override fun matches(insn: AbstractInsnNode) = insn.opcode in Opcodes.IRETURN..Opcodes.ARETURN
+        };
+
+        abstract fun matches(insn: AbstractInsnNode): Boolean
+    }
+
+    abstract val supportedInstructionTypes: Collection<InstructionType>
+
+    open fun extraTargetRestrictions(insn: AbstractInsnNode): Boolean = true
+
+    abstract fun expectedMethodSignature(
+        annotation: PsiAnnotation,
+        targetClass: ClassNode,
+        targetMethod: MethodNode,
+        insn: AbstractInsnNode
+    ): Pair<ParameterGroup, PsiType>?
+
+    override val allowCoerce = true
+
+    override fun expectedMethodSignature(
+        annotation: PsiAnnotation,
+        targetClass: ClassNode,
+        targetMethod: MethodNode
+    ): List<MethodSignature>? {
+        val insns = resolveInstructions(annotation, targetClass, targetMethod)
+            .ifEmpty { return emptyList() }
+            .map { it.insn }
+        if (insns.any { insn -> supportedInstructionTypes.none { it.matches(insn) } }) return emptyList()
+        val signatures = insns.map { expectedMethodSignature(annotation, targetClass, targetMethod, it) }
+        val firstMatch = signatures[0] ?: return emptyList()
+        if (signatures.drop(1).any { it != firstMatch }) return emptyList()
+        return listOf(
+            MethodSignature(
+                listOf(
+                    firstMatch.first,
+                    ParameterGroup(
+                        collectTargetMethodParameters(annotation.project, targetClass, targetMethod),
+                        required = ParameterGroup.RequiredLevel.OPTIONAL,
+                        isVarargs = true,
+                    ),
+                ),
+                firstMatch.second
+            )
+        )
+    }
+
+    protected fun getInsnReturnType(insn: AbstractInsnNode): Type? {
+        return when {
+            insn is MethodInsnNode -> Type.getReturnType(insn.desc)
+            insn is FieldInsnNode -> when (insn.opcode) {
+                Opcodes.GETFIELD, Opcodes.GETSTATIC -> Type.getType(insn.desc)
+                else -> Type.VOID_TYPE
+            }
+
+            insn is TypeInsnNode -> when (insn.opcode) {
+                Opcodes.NEW -> Type.getObjectType(insn.desc)
+                Opcodes.INSTANCEOF -> Type.BOOLEAN_TYPE
+                else -> null
+            }
+
+            isConstant(insn) -> getConstantType(insn)
+            else -> null
+        }
+    }
+
+    protected fun getPsiReturnType(
+        insn: AbstractInsnNode,
+        annotation: PsiAnnotation
+    ): PsiType? {
+        val elementFactory = JavaPsiFacade.getElementFactory(annotation.project)
+
+        return when (insn) {
+            is MethodInsnNode -> {
+                val sourceClassAndMethod = (
+                    MemberReference(insn.name, insn.desc, insn.owner.replace('/', '.'))
+                        .resolveAsm(annotation.project) as? MethodTargetMember
+                    )?.classAndMethod
+                sourceClassAndMethod?.method?.getGenericReturnType(sourceClassAndMethod.clazz, annotation.project)
+                    ?: Type.getType(insn.desc).toPsiType(elementFactory)
+            }
+
+            is FieldInsnNode -> {
+                val sourceClassAndField = (
+                    MemberReference(insn.name, insn.desc, insn.owner.replace('/', '.'))
+                        .resolveAsm(annotation.project) as? FieldTargetMember
+                    )?.classAndField
+                sourceClassAndField?.field?.getGenericType(sourceClassAndField.clazz, annotation.project)
+                    ?: Type.getType(insn.desc).toPsiType(elementFactory)
+            }
+
+            else -> getInsnReturnType(insn)?.toPsiType(elementFactory)
+        }
+    }
+
+    protected fun getInsnArgTypes(
+        insn: AbstractInsnNode,
+        targetClass: ClassNode
+    ): List<Type>? {
+        return when (insn) {
+            is MethodInsnNode -> {
+                val args = Type.getArgumentTypes(insn.desc).toMutableList()
+                when (insn.opcode) {
+                    Opcodes.INVOKESTATIC -> {}
+                    Opcodes.INVOKESPECIAL -> {
+                        args.add(0, Type.getObjectType(if (oldSuperBehavior) insn.owner else targetClass.name))
+                    }
+
+                    else -> {
+                        args.add(0, Type.getObjectType(insn.owner))
+                    }
+                }
+                args
+            }
+
+            is FieldInsnNode -> {
+                when (insn.opcode) {
+                    Opcodes.GETFIELD -> listOf(Type.getObjectType(insn.owner))
+                    Opcodes.PUTFIELD -> listOf(Type.getObjectType(insn.owner), Type.getType(insn.desc))
+                    Opcodes.GETSTATIC -> emptyList()
+                    Opcodes.PUTSTATIC -> listOf(Type.getType(insn.desc))
+                    else -> null
+                }
+            }
+
+            is TypeInsnNode -> {
+                when (insn.opcode) {
+                    Opcodes.INSTANCEOF -> listOf(Type.getType(Any::class.java))
+                    else -> null
+                }
+            }
+
+            else -> null
+        }
+    }
+
+    protected fun getPsiParameters(
+        insn: AbstractInsnNode,
+        targetClass: ClassNode,
+        annotation: PsiAnnotation
+    ): List<Parameter>? {
+        val elementFactory = JavaPsiFacade.getElementFactory(annotation.project)
+
+        return when (insn) {
+            is MethodInsnNode -> {
+                val sourceClassAndMethod = (
+                    MemberReference(insn.name, insn.desc, insn.owner.replace('/', '.'))
+                        .resolveAsm(annotation.project) as? MethodTargetMember
+                    )?.classAndMethod
+                val parameters = mutableListOf<Parameter>()
+                if (insn.opcode != Opcodes.INVOKESTATIC) {
+                    val receiver = if (insn.opcode == Opcodes.INVOKESPECIAL && !oldSuperBehavior) {
+                        targetClass.name
+                    } else {
+                        insn.owner
+                    }
+                    parameters += Parameter("instance", Type.getObjectType(receiver).toPsiType(elementFactory))
+                }
+                val genericParams = sourceClassAndMethod?.method?.getGenericParameterTypes(
+                    sourceClassAndMethod.clazz,
+                    annotation.project,
+                )
+                if (genericParams != null) {
+                    genericParams.withIndex().mapTo(parameters) { (index, type) ->
+                        val i = if (insn.opcode == Opcodes.INVOKESTATIC) index else index + 1
+                        val name = sourceClassAndMethod.method.localVariables?.getOrNull(i)?.name?.toJavaIdentifier()
+                        sanitizedParameter(type, name)
+                    }
+                } else {
+                    Type.getArgumentTypes(insn.desc).mapTo(parameters) {
+                        sanitizedParameter(it.toPsiType(elementFactory), null)
+                    }
+                }
+                parameters
+            }
+
+            is FieldInsnNode -> {
+                val ownerType = Type.getObjectType(insn.owner).toPsiType(elementFactory)
+                val sourceClassAndField = (
+                    MemberReference(insn.name, insn.desc, insn.owner.replace('/', '.'))
+                        .resolveAsm(annotation.project) as? FieldTargetMember
+                    )?.classAndField
+                val valueType =
+                    sourceClassAndField?.field?.getGenericType(sourceClassAndField.clazz, annotation.project)
+                        ?: Type.getType(insn.desc).toPsiType(elementFactory)
+                val owner = Parameter("instance", ownerType)
+                val value = Parameter("value", valueType)
+                when (insn.opcode) {
+                    Opcodes.GETFIELD -> listOf(owner)
+                    Opcodes.PUTFIELD -> listOf(owner, value)
+                    Opcodes.GETSTATIC -> emptyList()
+                    Opcodes.PUTSTATIC -> listOf(value)
+                    else -> null
+                }
+            }
+
+            is TypeInsnNode -> {
+                when (insn.opcode) {
+                    Opcodes.INSTANCEOF ->
+                        listOf(Parameter("object", Type.getType(Any::class.java).toPsiType(elementFactory)))
+
+                    else -> null
+                }
+            }
+
+            else -> null
+        } ?: getInsnArgTypes(insn, targetClass)?.map { Parameter(null, it.toPsiType(elementFactory)) }
+    }
+}
+
+private val CONSTANTS_ALL = intArrayOf(
+    Opcodes.ACONST_NULL,
+    Opcodes.ICONST_M1,
+    Opcodes.ICONST_0, Opcodes.ICONST_1, Opcodes.ICONST_2, Opcodes.ICONST_3, Opcodes.ICONST_4, Opcodes.ICONST_5,
+    Opcodes.LCONST_0, Opcodes.LCONST_1,
+    Opcodes.FCONST_0, Opcodes.FCONST_1, Opcodes.FCONST_2,
+    Opcodes.DCONST_0, Opcodes.DCONST_1,
+    Opcodes.BIPUSH, // 15
+    Opcodes.SIPUSH, // 16
+    Opcodes.LDC, // 17
+    Opcodes.CHECKCAST, // 18
+    Opcodes.INSTANCEOF // 19
+)
+
+private val CONSTANTS_TYPES = arrayOf(
+    "V", // null is returned as Type.VOID_TYPE
+    "I",
+    "I", "I", "I", "I", "I", "I",
+    "J", "J",
+    "F", "F", "F",
+    "D", "D",
+    "I", // "B",
+    "I"
+)
+
+private fun isConstant(insn: AbstractInsnNode?): Boolean {
+    return insn != null && insn.opcode in CONSTANTS_ALL
+}
+
+private fun getConstantType(insn: AbstractInsnNode?): Type? {
+    return when (insn) {
+        null -> null
+        is LdcInsnNode -> {
+            val cst = insn.cst
+            when (cst) {
+                is Int -> return Type.INT_TYPE
+                is Float -> return Type.FLOAT_TYPE
+                is Long -> return Type.LONG_TYPE
+                is Double -> return Type.DOUBLE_TYPE
+                is String -> return Type.getType(String::class.java)
+                is Type -> return Type.getType(Class::class.java)
+                else -> null
+            }
+        }
+
+        is TypeInsnNode -> {
+            return if (insn.getOpcode() < Opcodes.CHECKCAST) {
+                null // Don't treat NEW and ANEWARRAY as constants
+            } else Type.getType(Class::class.java)
+        }
+
+        else -> {
+            val index = CONSTANTS_ALL.indexOf(insn.opcode)
+            return if (index < 0) null else Type.getType(CONSTANTS_TYPES.get(index))
+        }
+    }
+}
Index: src/main/kotlin/platform/mixin/handlers/mixinextras/ModifyExpressionValueHandler.kt
===================================================================
--- src/main/kotlin/platform/mixin/handlers/mixinextras/ModifyExpressionValueHandler.kt	(revision 449a4b79ff378b209b2be82c6dc8e53d8a7f2c8a)
+++ src/main/kotlin/platform/mixin/handlers/mixinextras/ModifyExpressionValueHandler.kt	(revision 449a4b79ff378b209b2be82c6dc8e53d8a7f2c8a)
@@ -0,0 +1,51 @@
+/*
+ * Minecraft Development for IntelliJ
+ *
+ * https://mcdev.io/
+ *
+ * Copyright (C) 2023 minecraft-dev
+ *
+ * This program is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU Lesser General Public License as published
+ * by the Free Software Foundation, version 3.0 only.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU Lesser General Public License
+ * along with this program.  If not, see <https://www.gnu.org/licenses/>.
+ */
+
+package com.demonwav.mcdev.platform.mixin.handlers.mixinextras
+
+import com.demonwav.mcdev.platform.mixin.inspection.injector.ParameterGroup
+import com.demonwav.mcdev.util.Parameter
+import com.intellij.psi.PsiAnnotation
+import com.intellij.psi.PsiType
+import org.objectweb.asm.Type
+import org.objectweb.asm.tree.AbstractInsnNode
+import org.objectweb.asm.tree.ClassNode
+import org.objectweb.asm.tree.MethodNode
+
+class ModifyExpressionValueHandler : MixinExtrasInjectorAnnotationHandler() {
+    override val supportedInstructionTypes = listOf(
+        InstructionType.METHOD_CALL, InstructionType.FIELD_GET, InstructionType.INSTANTIATION, InstructionType.CONSTANT
+    )
+
+    override fun extraTargetRestrictions(insn: AbstractInsnNode): Boolean {
+        val returnType = getInsnReturnType(insn) ?: return false
+        return returnType != Type.VOID_TYPE
+    }
+
+    override fun expectedMethodSignature(
+        annotation: PsiAnnotation,
+        targetClass: ClassNode,
+        targetMethod: MethodNode,
+        insn: AbstractInsnNode
+    ): Pair<ParameterGroup, PsiType>? {
+        val psiType = getPsiReturnType(insn, annotation) ?: return null
+        return ParameterGroup(listOf(Parameter("original", psiType))) to psiType
+    }
+}
Index: src/main/kotlin/platform/mixin/handlers/mixinextras/ModifyReceiverHandler.kt
===================================================================
--- src/main/kotlin/platform/mixin/handlers/mixinextras/ModifyReceiverHandler.kt	(revision 449a4b79ff378b209b2be82c6dc8e53d8a7f2c8a)
+++ src/main/kotlin/platform/mixin/handlers/mixinextras/ModifyReceiverHandler.kt	(revision 449a4b79ff378b209b2be82c6dc8e53d8a7f2c8a)
@@ -0,0 +1,52 @@
+/*
+ * Minecraft Development for IntelliJ
+ *
+ * https://mcdev.io/
+ *
+ * Copyright (C) 2023 minecraft-dev
+ *
+ * This program is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU Lesser General Public License as published
+ * by the Free Software Foundation, version 3.0 only.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU Lesser General Public License
+ * along with this program.  If not, see <https://www.gnu.org/licenses/>.
+ */
+
+package com.demonwav.mcdev.platform.mixin.handlers.mixinextras
+
+import com.demonwav.mcdev.platform.mixin.inspection.injector.ParameterGroup
+import com.intellij.psi.PsiAnnotation
+import com.intellij.psi.PsiType
+import org.objectweb.asm.Opcodes
+import org.objectweb.asm.tree.AbstractInsnNode
+import org.objectweb.asm.tree.ClassNode
+import org.objectweb.asm.tree.MethodNode
+
+class ModifyReceiverHandler : MixinExtrasInjectorAnnotationHandler() {
+    override val supportedInstructionTypes = listOf(
+        InstructionType.METHOD_CALL, InstructionType.FIELD_GET, InstructionType.FIELD_SET
+    )
+
+    override fun extraTargetRestrictions(insn: AbstractInsnNode) = when (insn.opcode) {
+        Opcodes.INVOKEVIRTUAL, Opcodes.INVOKESPECIAL, Opcodes.INVOKEINTERFACE,
+        Opcodes.GETFIELD, Opcodes.PUTFIELD -> true
+
+        else -> false
+    }
+
+    override fun expectedMethodSignature(
+        annotation: PsiAnnotation,
+        targetClass: ClassNode,
+        targetMethod: MethodNode,
+        insn: AbstractInsnNode
+    ): Pair<ParameterGroup, PsiType>? {
+        val params = getPsiParameters(insn, targetClass, annotation) ?: return null
+        return ParameterGroup(params) to params[0].type
+    }
+}
Index: src/main/kotlin/platform/mixin/handlers/mixinextras/ModifyReturnValueHandler.kt
===================================================================
--- src/main/kotlin/platform/mixin/handlers/mixinextras/ModifyReturnValueHandler.kt	(revision 449a4b79ff378b209b2be82c6dc8e53d8a7f2c8a)
+++ src/main/kotlin/platform/mixin/handlers/mixinextras/ModifyReturnValueHandler.kt	(revision 449a4b79ff378b209b2be82c6dc8e53d8a7f2c8a)
@@ -0,0 +1,44 @@
+/*
+ * Minecraft Development for IntelliJ
+ *
+ * https://mcdev.io/
+ *
+ * Copyright (C) 2023 minecraft-dev
+ *
+ * This program is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU Lesser General Public License as published
+ * by the Free Software Foundation, version 3.0 only.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU Lesser General Public License
+ * along with this program.  If not, see <https://www.gnu.org/licenses/>.
+ */
+
+package com.demonwav.mcdev.platform.mixin.handlers.mixinextras
+
+import com.demonwav.mcdev.platform.mixin.inspection.injector.ParameterGroup
+import com.demonwav.mcdev.platform.mixin.util.getGenericReturnType
+import com.demonwav.mcdev.util.Parameter
+import com.intellij.psi.PsiAnnotation
+import com.intellij.psi.PsiType
+import org.objectweb.asm.tree.AbstractInsnNode
+import org.objectweb.asm.tree.ClassNode
+import org.objectweb.asm.tree.MethodNode
+
+class ModifyReturnValueHandler : MixinExtrasInjectorAnnotationHandler() {
+    override val supportedInstructionTypes = listOf(InstructionType.RETURN)
+
+    override fun expectedMethodSignature(
+        annotation: PsiAnnotation,
+        targetClass: ClassNode,
+        targetMethod: MethodNode,
+        insn: AbstractInsnNode
+    ): Pair<ParameterGroup, PsiType>? {
+        val returnType = targetMethod.getGenericReturnType(targetClass, annotation.project)
+        return ParameterGroup(listOf(Parameter("original", returnType))) to returnType
+    }
+}
Index: src/main/kotlin/platform/mixin/handlers/mixinextras/WrapOperationHandler.kt
===================================================================
--- src/main/kotlin/platform/mixin/handlers/mixinextras/WrapOperationHandler.kt	(revision 449a4b79ff378b209b2be82c6dc8e53d8a7f2c8a)
+++ src/main/kotlin/platform/mixin/handlers/mixinextras/WrapOperationHandler.kt	(revision 449a4b79ff378b209b2be82c6dc8e53d8a7f2c8a)
@@ -0,0 +1,71 @@
+/*
+ * Minecraft Development for IntelliJ
+ *
+ * https://mcdev.io/
+ *
+ * Copyright (C) 2023 minecraft-dev
+ *
+ * This program is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU Lesser General Public License as published
+ * by the Free Software Foundation, version 3.0 only.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU Lesser General Public License
+ * along with this program.  If not, see <https://www.gnu.org/licenses/>.
+ */
+
+package com.demonwav.mcdev.platform.mixin.handlers.mixinextras
+
+import com.demonwav.mcdev.platform.mixin.inspection.injector.ParameterGroup
+import com.demonwav.mcdev.platform.mixin.util.MixinConstants
+import com.demonwav.mcdev.util.Parameter
+import com.intellij.psi.JavaPsiFacade
+import com.intellij.psi.PsiAnnotation
+import com.intellij.psi.PsiElement
+import com.intellij.psi.PsiPrimitiveType
+import com.intellij.psi.PsiType
+import com.intellij.psi.search.GlobalSearchScope
+import org.objectweb.asm.tree.AbstractInsnNode
+import org.objectweb.asm.tree.ClassNode
+import org.objectweb.asm.tree.MethodNode
+
+class WrapOperationHandler : MixinExtrasInjectorAnnotationHandler() {
+    override val supportedInstructionTypes = listOf(
+        InstructionType.METHOD_CALL, InstructionType.FIELD_GET, InstructionType.FIELD_SET, InstructionType.INSTANCEOF
+    )
+
+    override fun getAtKey(annotation: PsiAnnotation): String {
+        return if (annotation.hasAttribute("constant")) "constant" else "at"
+    }
+
+    override fun expectedMethodSignature(
+        annotation: PsiAnnotation,
+        targetClass: ClassNode,
+        targetMethod: MethodNode,
+        insn: AbstractInsnNode
+    ): Pair<ParameterGroup, PsiType>? {
+        val params = getPsiParameters(insn, targetClass, annotation) ?: return null
+        val returnType = getPsiReturnType(insn, annotation) ?: return null
+        val operationType = getOperationType(annotation, returnType) ?: return null
+        return ParameterGroup(
+            params + Parameter("original", operationType)
+        ) to returnType
+    }
+
+    private fun getOperationType(context: PsiElement, type: PsiType): PsiType? {
+        val project = context.project
+        val boxedType = if (type is PsiPrimitiveType) {
+            type.getBoxedType(context) ?: return null
+        } else {
+            type
+        }
+        val psiClass = JavaPsiFacade.getInstance(project)
+            .findClass(MixinConstants.MixinExtras.OPERATION, GlobalSearchScope.allScope(project))
+            ?: return null
+        return JavaPsiFacade.getElementFactory(project).createType(psiClass, boxedType)
+    }
+}
Index: src/main/kotlin/platform/mixin/handlers/mixinextras/WrapWithConditionHandler.kt
===================================================================
--- src/main/kotlin/platform/mixin/handlers/mixinextras/WrapWithConditionHandler.kt	(revision 449a4b79ff378b209b2be82c6dc8e53d8a7f2c8a)
+++ src/main/kotlin/platform/mixin/handlers/mixinextras/WrapWithConditionHandler.kt	(revision 449a4b79ff378b209b2be82c6dc8e53d8a7f2c8a)
@@ -0,0 +1,49 @@
+/*
+ * Minecraft Development for IntelliJ
+ *
+ * https://mcdev.io/
+ *
+ * Copyright (C) 2023 minecraft-dev
+ *
+ * This program is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU Lesser General Public License as published
+ * by the Free Software Foundation, version 3.0 only.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU Lesser General Public License
+ * along with this program.  If not, see <https://www.gnu.org/licenses/>.
+ */
+
+package com.demonwav.mcdev.platform.mixin.handlers.mixinextras
+
+import com.demonwav.mcdev.platform.mixin.inspection.injector.ParameterGroup
+import com.intellij.psi.PsiAnnotation
+import com.intellij.psi.PsiType
+import org.objectweb.asm.Type
+import org.objectweb.asm.tree.AbstractInsnNode
+import org.objectweb.asm.tree.ClassNode
+import org.objectweb.asm.tree.MethodNode
+
+class WrapWithConditionHandler : MixinExtrasInjectorAnnotationHandler() {
+    override val oldSuperBehavior = true
+
+    override val supportedInstructionTypes = listOf(
+        InstructionType.METHOD_CALL, InstructionType.FIELD_SET
+    )
+
+    override fun extraTargetRestrictions(insn: AbstractInsnNode) = getInsnReturnType(insn) == Type.VOID_TYPE
+
+    override fun expectedMethodSignature(
+        annotation: PsiAnnotation,
+        targetClass: ClassNode,
+        targetMethod: MethodNode,
+        insn: AbstractInsnNode
+    ): Pair<ParameterGroup, PsiType>? {
+        val params = getPsiParameters(insn, targetClass, annotation) ?: return null
+        return ParameterGroup(params) to PsiType.BOOLEAN
+    }
+}
Index: src/main/kotlin/platform/mixin/util/MixinConstants.kt
===================================================================
--- src/main/kotlin/platform/mixin/util/MixinConstants.kt	(revision 9ae67da17c7109ef7ce9fd12290798cab82e1d02)
+++ src/main/kotlin/platform/mixin/util/MixinConstants.kt	(revision 449a4b79ff378b209b2be82c6dc8e53d8a7f2c8a)
@@ -75,4 +75,8 @@
         const val REDIRECT = "org.spongepowered.asm.mixin.injection.Redirect"
         const val SURROGATE = "org.spongepowered.asm.mixin.injection.Surrogate"
     }
+
+    object MixinExtras {
+        const val OPERATION = "com.llamalad7.mixinextras.injector.wrapoperation.Operation"
-}
+    }
+}
Index: src/main/resources/META-INF/plugin.xml
===================================================================
--- src/main/resources/META-INF/plugin.xml	(revision 9ae67da17c7109ef7ce9fd12290798cab82e1d02)
+++ src/main/resources/META-INF/plugin.xml	(revision 449a4b79ff378b209b2be82c6dc8e53d8a7f2c8a)
@@ -80,6 +80,7 @@
         <extensionPoint name="injectionPoint" beanClass="com.demonwav.mcdev.platform.mixin.handlers.injectionPoint.InjectionPointInfo" dynamic="true">
             <with attribute="implementation" implements="com.demonwav.mcdev.platform.mixin.handlers.injectionPoint.InjectionPoint"/>
         </extensionPoint>
+        <extensionPoint name="injectionPointAnnotation" beanClass="com.demonwav.mcdev.platform.mixin.handlers.injectionPoint.InjectionPointAnnotation" />
     </extensionPoints>
 
     <extensions defaultExtensionNs="com.demonwav.minecraft-dev">
@@ -135,6 +136,11 @@
         <mixinAnnotationHandler annotation="org.spongepowered.asm.mixin.injection.ModifyConstant" implementation="com.demonwav.mcdev.platform.mixin.handlers.ModifyConstantHandler" />
         <mixinAnnotationHandler annotation="org.spongepowered.asm.mixin.injection.ModifyVariable" implementation="com.demonwav.mcdev.platform.mixin.handlers.ModifyVariableHandler" />
         <mixinAnnotationHandler annotation="org.spongepowered.asm.mixin.injection.Redirect" implementation="com.demonwav.mcdev.platform.mixin.handlers.RedirectInjectorHandler" />
+        <mixinAnnotationHandler annotation="com.llamalad7.mixinextras.injector.ModifyExpressionValue" implementation="com.demonwav.mcdev.platform.mixin.handlers.mixinextras.ModifyExpressionValueHandler" />
+        <mixinAnnotationHandler annotation="com.llamalad7.mixinextras.injector.ModifyReceiver" implementation="com.demonwav.mcdev.platform.mixin.handlers.mixinextras.ModifyReceiverHandler" />
+        <mixinAnnotationHandler annotation="com.llamalad7.mixinextras.injector.ModifyReturnValue" implementation="com.demonwav.mcdev.platform.mixin.handlers.mixinextras.ModifyReturnValueHandler" />
+        <mixinAnnotationHandler annotation="com.llamalad7.mixinextras.injector.wrapoperation.WrapOperation" implementation="com.demonwav.mcdev.platform.mixin.handlers.mixinextras.WrapOperationHandler" />
+        <mixinAnnotationHandler annotation="com.llamalad7.mixinextras.injector.WrapWithCondition" implementation="com.demonwav.mcdev.platform.mixin.handlers.mixinextras.WrapWithConditionHandler" />
 
         <injectionPoint atCode="INVOKE" implementation="com.demonwav.mcdev.platform.mixin.handlers.injectionPoint.MethodInjectionPoint" />
         <injectionPoint atCode="INVOKE_ASSIGN" implementation="com.demonwav.mcdev.platform.mixin.handlers.injectionPoint.MethodInjectionPoint" />
@@ -147,6 +153,8 @@
         <injectionPoint atCode="LOAD" implementation="com.demonwav.mcdev.platform.mixin.handlers.injectionPoint.LoadInjectionPoint" />
         <injectionPoint atCode="STORE" implementation="com.demonwav.mcdev.platform.mixin.handlers.injectionPoint.StoreInjectionPoint" />
         <injectionPoint atCode="CONSTANT" implementation="com.demonwav.mcdev.platform.mixin.handlers.injectionPoint.ConstantInjectionPoint" />
+
+        <injectionPointAnnotation annotation="org.spongepowered.asm.mixin.injection.Constant" atCode="CONSTANT" />
     </extensions>
 
     <extensions defaultExtensionNs="com.intellij">