MixinExtras: Add handler signature support for expressions.

build.gradle.kts

@@ -101,7 +101,7 @@ dependencies {
     implementation(files(Jvm.current().toolsJar))
 
     // TODO: temporary waiting for MixinExtras expression library
-    implementation("com.github.LlamaLad7.MixinExtras:mixinextras-common:fe1e2b9")
+    implementation("com.github.LlamaLad7.MixinExtras:mixinextras-common:debfdc8")
     implementation("org.spongepowered:mixin:0.8.4")
     implementation("org.ow2.asm:asm-util:9.3")
 

src/main/kotlin/platform/mixin/handlers/injectionPoint/InjectionPoint.kt

@@ -410,6 +410,7 @@ abstract class CollectVisitor<T : PsiElement>(protected val mode: Mode) {
         insn: AbstractInsnNode,
         element: T,
         qualifier: String? = null,
+        decorations: Map<String, Any?> = emptyMap(),
     ) {
         // apply shift.
         // being able to break out of the shift loops is important to prevent IDE freezes in case of large shift bys.
@@ -430,7 +431,14 @@ abstract class CollectVisitor<T : PsiElement>(protected val mode: Mode) {
             }
         }
 
-        val result = Result(nextIndex++, insn, shiftedInsn ?: return, element, qualifier)
+        val result = Result(
+            nextIndex++,
+            insn,
+            shiftedInsn ?: return,
+            element,
+            qualifier,
+            if (insn === shiftedInsn) decorations else emptyMap()
+        )
         var isFiltered = false
         for ((name, filter) in resultFilters) {
             if (!filter(result, method)) {
@@ -466,6 +474,7 @@ abstract class CollectVisitor<T : PsiElement>(protected val mode: Mode) {
         val insn: AbstractInsnNode,
         val target: T,
         val qualifier: String? = null,
+        val decorations: Map<String, Any?>
     )
 
     enum class Mode { MATCH_ALL, MATCH_FIRST, COMPLETION }

src/main/kotlin/platform/mixin/handlers/mixinextras/ExpressionInjectionPoint.kt

@@ -290,37 +290,37 @@ class ExpressionInjectionPoint : InjectionPoint<PsiElement>() {
                 return
             }
 
-            val result = IdentityHashMap<AbstractInsnNode, PsiElement>()
+            val result = IdentityHashMap<AbstractInsnNode, Pair<PsiElement, Map<String, Any?>>>()
 
             for ((expr, psiExpr) in expressions) {
-                insns.iterator().forEachRemaining { insn ->
-                    val genericDecorations = IdentityHashMap<AbstractInsnNode, MutableMap<String, Any?>>()
-                    val injectorSpecificDecorations = IdentityHashMap<AbstractInsnNode, MutableMap<String, Any?>>()
-                    val captured = mutableListOf<Pair<AbstractInsnNode, Int>>()
-
+                val decorations = IdentityHashMap<AbstractInsnNode, MutableMap<String, Any?>>()
+                val captured = mutableListOf<Pair<AbstractInsnNode, Int>>()
+                for (insn in insns) {
                     val sink = object : Expression.OutputSink {
                         override fun capture(node: FlowValue, expr: Expression?) {
                             captured += node.insn to (expr?.src?.startIndex ?: 0)
+                            decorations.getOrPut(insn, ::mutableMapOf).putAll(node.decorations)
                         }
 
                         override fun decorate(insn: AbstractInsnNode, key: String, value: Any?) {
-                            genericDecorations.computeIfAbsent(insn) { mutableMapOf() }[key] = value
+                            decorations.getOrPut(insn, ::mutableMapOf)[key] = value
                         }
 
                         override fun decorateInjectorSpecific(insn: AbstractInsnNode, key: String, value: Any?) {
-                            injectorSpecificDecorations.computeIfAbsent(insn) { mutableMapOf() }[key] = value
+                            // Our maps are per-injector anyway, so this is just a normal decoration.
+                            decorations.getOrPut(insn, ::mutableMapOf)[key] = value
                         }
                     }
 
-                    val flow = flows[insn] ?: return@forEachRemaining
+                    val flow = flows[insn] ?: continue
                     try {
                         if (expr.matches(flow, ExpressionContext(pool, sink, targetClass, methodNode))) {
                             for ((capturedInsn, startOffset) in captured) {
                                 val capturedExpr = psiExpr.findElementAt(startOffset)
                                     ?.parentOfType<MECapturingExpression>(withSelf = true)
                                     ?.expression
                                     ?: psiExpr
-                                result.putIfAbsent(capturedInsn, capturedExpr)
+                                result.putIfAbsent(capturedInsn, capturedExpr to decorations[capturedInsn].orEmpty())
                             }
                         }
                     } catch (e: ProcessCanceledException) {
@@ -335,11 +335,9 @@ class ExpressionInjectionPoint : InjectionPoint<PsiElement>() {
                 return
             }
 
-            insns.iterator().forEachRemaining { insn ->
-                val element = result[insn]
-                if (element != null) {
-                    addResult(insn, element)
-                }
+            for (insn in insns) {
+                val (element, decorations) = result[insn] ?: continue
+                addResult(insn, element, decorations = decorations)
             }
         }
     }

src/main/kotlin/platform/mixin/handlers/mixinextras/MixinExtrasInjectorAnnotationHandler.kt

@@ -35,8 +35,10 @@ 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.PsiElement
 import com.intellij.psi.PsiType
 import com.intellij.psi.PsiTypes
+import com.llamalad7.mixinextras.utils.Decorations
 import org.objectweb.asm.Opcodes
 import org.objectweb.asm.Type
 import org.objectweb.asm.tree.AbstractInsnNode
@@ -52,30 +54,38 @@ abstract class MixinExtrasInjectorAnnotationHandler : InjectorAnnotationHandler(
 
     enum class InstructionType {
         METHOD_CALL {
-            override fun matches(insn: AbstractInsnNode) = insn is MethodInsnNode && insn.name != "<init>"
+            override fun matches(target: TargetInsn) = target.insn is MethodInsnNode && target.insn.name != "<init>"
         },
         FIELD_GET {
-            override fun matches(insn: AbstractInsnNode) =
-                insn.opcode == Opcodes.GETFIELD || insn.opcode == Opcodes.GETSTATIC
+            override fun matches(target: TargetInsn) =
+                target.insn.opcode == Opcodes.GETFIELD || target.insn.opcode == Opcodes.GETSTATIC
         },
         FIELD_SET {
-            override fun matches(insn: AbstractInsnNode) =
-                insn.opcode == Opcodes.PUTFIELD || insn.opcode == Opcodes.PUTSTATIC
+            override fun matches(target: TargetInsn) =
+                target.insn.opcode == Opcodes.PUTFIELD || target.insn.opcode == Opcodes.PUTSTATIC
         },
         INSTANTIATION {
-            override fun matches(insn: AbstractInsnNode) = insn.opcode == Opcodes.NEW
+            override fun matches(target: TargetInsn) = target.insn.opcode == Opcodes.NEW
         },
         INSTANCEOF {
-            override fun matches(insn: AbstractInsnNode) = insn.opcode == Opcodes.INSTANCEOF
+            override fun matches(target: TargetInsn) = target.insn.opcode == Opcodes.INSTANCEOF
         },
         CONSTANT {
-            override fun matches(insn: AbstractInsnNode) = isConstant(insn)
+            override fun matches(target: TargetInsn) = isConstant(target.insn)
         },
         RETURN {
-            override fun matches(insn: AbstractInsnNode) = insn.opcode in Opcodes.IRETURN..Opcodes.ARETURN
+            override fun matches(target: TargetInsn) = target.insn.opcode in Opcodes.IRETURN..Opcodes.ARETURN
+        },
+        SIMPLE_OPERATION {
+            override fun matches(target: TargetInsn) =
+                target.hasDecoration(Decorations.SIMPLE_OPERATION_ARGS) &&
+                    target.hasDecoration(Decorations.SIMPLE_OPERATION_RETURN_TYPE)
+        },
+        SIMPLE_EXPRESSION {
+            override fun matches(target: TargetInsn) = target.hasDecoration(Decorations.SIMPLE_EXPRESSION_TYPE)
         };
 
-        abstract fun matches(insn: AbstractInsnNode): Boolean
+        abstract fun matches(target: TargetInsn): Boolean
     }
 
     abstract val supportedInstructionTypes: Collection<InstructionType>
@@ -86,7 +96,7 @@ abstract class MixinExtrasInjectorAnnotationHandler : InjectorAnnotationHandler(
         annotation: PsiAnnotation,
         targetClass: ClassNode,
         targetMethod: MethodNode,
-        insn: AbstractInsnNode
+        target: TargetInsn,
     ): Pair<ParameterGroup, PsiType>?
 
     override val allowCoerce = true
@@ -98,9 +108,11 @@ abstract class MixinExtrasInjectorAnnotationHandler : InjectorAnnotationHandler(
     ): List<MethodSignature>? {
         val insns = resolveInstructions(annotation, targetClass, targetMethod)
             .ifEmpty { return emptyList() }
-            .map { it.insn }
+            .map { TargetInsn(it.insn, it.decorations) }
         if (insns.any { insn -> supportedInstructionTypes.none { it.matches(insn) } }) return emptyList()
-        val signatures = insns.map { expectedMethodSignature(annotation, targetClass, targetMethod, it) }
+        val signatures = insns.map { insn ->
+            expectedMethodSignature(annotation, targetClass, targetMethod, insn)
+        }
         val firstMatch = signatures[0] ?: return emptyList()
         if (signatures.drop(1).any { it != firstMatch }) return emptyList()
         return listOf(
@@ -287,7 +299,12 @@ abstract class MixinExtrasInjectorAnnotationHandler : InjectorAnnotationHandler(
             }
 
             else -> null
-        } ?: getInsnArgTypes(insn, targetClass)?.map { Parameter(null, it.toPsiType(elementFactory)) }
+        } ?: getInsnArgTypes(insn, targetClass)?.toParameters(annotation)
+    }
+
+    protected fun List<Type>.toParameters(context: PsiElement): List<Parameter> {
+        val elementFactory = JavaPsiFacade.getElementFactory(context.project)
+        return map { Parameter(null, it.toPsiType(elementFactory)) }
     }
 }
 

src/main/kotlin/platform/mixin/handlers/mixinextras/ModifyExpressionValueHandler.kt

@@ -21,17 +21,21 @@
 package com.demonwav.mcdev.platform.mixin.handlers.mixinextras
 
 import com.demonwav.mcdev.platform.mixin.inspection.injector.ParameterGroup
+import com.demonwav.mcdev.platform.mixin.util.toPsiType
 import com.demonwav.mcdev.util.Parameter
+import com.intellij.psi.JavaPsiFacade
 import com.intellij.psi.PsiAnnotation
 import com.intellij.psi.PsiType
+import com.llamalad7.mixinextras.utils.Decorations
 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
+        InstructionType.METHOD_CALL, InstructionType.FIELD_GET, InstructionType.INSTANTIATION, InstructionType.CONSTANT,
+        InstructionType.SIMPLE_EXPRESSION
     )
 
     override fun extraTargetRestrictions(insn: AbstractInsnNode): Boolean {
@@ -43,9 +47,23 @@ class ModifyExpressionValueHandler : MixinExtrasInjectorAnnotationHandler() {
         annotation: PsiAnnotation,
         targetClass: ClassNode,
         targetMethod: MethodNode,
-        insn: AbstractInsnNode
+        target: TargetInsn
     ): Pair<ParameterGroup, PsiType>? {
-        val psiType = getPsiReturnType(insn, annotation) ?: return null
+        val psiType = getReturnType(target, annotation) ?: return null
         return ParameterGroup(listOf(Parameter("original", psiType))) to psiType
     }
+
+    private fun getReturnType(
+        target: TargetInsn,
+        annotation: PsiAnnotation
+    ): PsiType? {
+        val psiReturnType = getPsiReturnType(target.insn, annotation)
+        val rawReturnType = getInsnReturnType(target.insn)
+        val exprType = target.getDecoration<Type>(Decorations.SIMPLE_EXPRESSION_TYPE)
+        if (exprType != null && rawReturnType != exprType) {
+            // The expression knows more than the standard logic does.
+            return exprType.toPsiType(JavaPsiFacade.getElementFactory(annotation.project))
+        }
+        return psiReturnType
+    }
 }

src/main/kotlin/platform/mixin/handlers/mixinextras/ModifyReceiverHandler.kt

@@ -44,9 +44,9 @@ class ModifyReceiverHandler : MixinExtrasInjectorAnnotationHandler() {
         annotation: PsiAnnotation,
         targetClass: ClassNode,
         targetMethod: MethodNode,
-        insn: AbstractInsnNode
+        target: TargetInsn
     ): Pair<ParameterGroup, PsiType>? {
-        val params = getPsiParameters(insn, targetClass, annotation) ?: return null
+        val params = getPsiParameters(target.insn, targetClass, annotation) ?: return null
         return ParameterGroup(params) to params[0].type
     }
 }

src/main/kotlin/platform/mixin/handlers/mixinextras/ModifyReturnValueHandler.kt

@@ -25,7 +25,6 @@ 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
 
@@ -36,8 +35,8 @@ class ModifyReturnValueHandler : MixinExtrasInjectorAnnotationHandler() {
         annotation: PsiAnnotation,
         targetClass: ClassNode,
         targetMethod: MethodNode,
-        insn: AbstractInsnNode
-    ): Pair<ParameterGroup, PsiType>? {
+        target: TargetInsn
+    ): Pair<ParameterGroup, PsiType> {
         val returnType = targetMethod.getGenericReturnType(targetClass, annotation.project)
         return ParameterGroup(listOf(Parameter("original", returnType))) to returnType
     }

src/main/kotlin/platform/mixin/handlers/mixinextras/TargetInsn.kt

@@ -0,0 +1,10 @@
+package com.demonwav.mcdev.platform.mixin.handlers.mixinextras
+
+import org.objectweb.asm.tree.AbstractInsnNode
+
+class TargetInsn(val insn: AbstractInsnNode, private val decorations: Map<String, Any?>) {
+    fun hasDecoration(key: String) = key in decorations
+
+    @Suppress("UNCHECKED_CAST")
+    fun <T> getDecoration(key: String): T? = decorations[key] as T
+}

src/main/kotlin/platform/mixin/handlers/mixinextras/WrapOperationHandler.kt

@@ -22,20 +22,22 @@ 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.MixinExtras.OPERATION
+import com.demonwav.mcdev.platform.mixin.util.toPsiType
 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 org.objectweb.asm.tree.AbstractInsnNode
+import com.llamalad7.mixinextras.utils.Decorations
+import org.objectweb.asm.Type
 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,
-        InstructionType.INSTANTIATION
+        InstructionType.INSTANTIATION, InstructionType.SIMPLE_OPERATION
     )
 
     override fun getAtKey(annotation: PsiAnnotation): String {
@@ -46,16 +48,35 @@ class WrapOperationHandler : MixinExtrasInjectorAnnotationHandler() {
         annotation: PsiAnnotation,
         targetClass: ClassNode,
         targetMethod: MethodNode,
-        insn: AbstractInsnNode
+        target: TargetInsn
     ): Pair<ParameterGroup, PsiType>? {
-        val params = getPsiParameters(insn, targetClass, annotation) ?: return null
-        val returnType = getPsiReturnType(insn, annotation) ?: return null
+        val params = getParameterTypes(target, targetClass, annotation) ?: return null
+        val returnType = getReturnType(target, annotation) ?: return null
         val operationType = getOperationType(annotation, returnType) ?: return null
         return ParameterGroup(
             params + Parameter("original", operationType)
         ) to returnType
     }
 
+    private fun getParameterTypes(
+        target: TargetInsn,
+        targetClass: ClassNode,
+        annotation: PsiAnnotation
+    ): List<Parameter>? {
+        getPsiParameters(target.insn, targetClass, annotation)?.let { return it }
+        val args = target.getDecoration<Array<Type>>(Decorations.SIMPLE_OPERATION_ARGS) ?: return null
+        return args.toList().toParameters(annotation)
+    }
+
+    private fun getReturnType(
+        target: TargetInsn,
+        annotation: PsiAnnotation
+    ): PsiType? {
+        getPsiReturnType(target.insn, annotation)?.let { return it }
+        val type = target.getDecoration<Type>(Decorations.SIMPLE_OPERATION_RETURN_TYPE) ?: return null
+        return type.toPsiType(JavaPsiFacade.getElementFactory(annotation.project))
+    }
+
     private fun getOperationType(context: PsiElement, type: PsiType): PsiType? {
         val project = context.project
         val boxedType = if (type is PsiPrimitiveType) {

src/main/kotlin/platform/mixin/handlers/mixinextras/WrapWithConditionHandler.kt

@@ -42,9 +42,9 @@ class WrapWithConditionHandler : MixinExtrasInjectorAnnotationHandler() {
         annotation: PsiAnnotation,
         targetClass: ClassNode,
         targetMethod: MethodNode,
-        insn: AbstractInsnNode
+        target: TargetInsn
     ): Pair<ParameterGroup, PsiType>? {
-        val params = getPsiParameters(insn, targetClass, annotation) ?: return null
+        val params = getPsiParameters(target.insn, targetClass, annotation) ?: return null
         return ParameterGroup(params) to PsiTypes.booleanType()
     }
 }