Support remapping other mixin injection annotations (fixes #6, fixes #7)

1 files changed, 13 insertions, 3 deletions

diff --git a/src/main/kotlin/com/replaymod/gradle/remap/PsiMapper.kt b/src/main/kotlin/com/replaymod/gradle/remap/PsiMapper.kt
index 35d4fd5..9ddb78a 100644
--- a/src/main/kotlin/com/replaymod/gradle/remap/PsiMapper.kt
+++ b/src/main/kotlin/com/replaymod/gradle/remap/PsiMapper.kt
@@ -284,10 +284,16 @@ internal class PsiMapper(private val map: MappingSet, private val file: PsiFile)
         })
     }
 
-    private fun remapInjectsAndRedirects(mapping: ClassMapping<*, *>) {
+    private fun remapMixinInjections(mapping: ClassMapping<*, *>) {
         file.accept(object : JavaRecursiveElementVisitor() {
             override fun visitMethod(method: PsiMethod) {
-                val annotation = method.getAnnotation(CLASS_INJECT) ?: method.getAnnotation(CLASS_REDIRECT) ?: return
+                val annotation = method.getAnnotation(CLASS_INJECT)
+                        ?: method.getAnnotation(CLASS_MODIFY_ARG)
+                        ?: method.getAnnotation(CLASS_MODIFY_ARGS)
+                        ?: method.getAnnotation(CLASS_MODIFY_CONSTANT)
+                        ?: method.getAnnotation(CLASS_MODIFY_VARIABLE)
+                        ?: method.getAnnotation(CLASS_REDIRECT)
+                        ?: return
 
                 for (attribute in annotation.parameterList.attributes) {
                     if ("method" != attribute.name) continue
@@ -438,7 +444,7 @@ internal class PsiMapper(private val map: MappingSet, private val file: PsiFile)
                     remapAccessors(mapping)
                 }
                 if (!mapping.methodMappings.isEmpty()) {
-                    remapInjectsAndRedirects(mapping)
+                    remapMixinInjections(mapping)
                 }
             }
         })
@@ -509,6 +515,10 @@ internal class PsiMapper(private val map: MappingSet, private val file: PsiFile)
         private const val CLASS_INVOKER = "org.spongepowered.asm.mixin.gen.Invoker"
         private const val CLASS_AT = "org.spongepowered.asm.mixin.injection.At"
         private const val CLASS_INJECT = "org.spongepowered.asm.mixin.injection.Inject"
+        private const val CLASS_MODIFY_ARG = "org.spongepowered.asm.mixin.injection.ModifyArg"
+        private const val CLASS_MODIFY_ARGS = "org.spongepowered.asm.mixin.injection.ModifyArgs"
+        private const val CLASS_MODIFY_CONSTANT = "org.spongepowered.asm.mixin.injection.ModifyConstant"
+        private const val CLASS_MODIFY_VARIABLE = "org.spongepowered.asm.mixin.injection.ModifyVariable"
         private const val CLASS_REDIRECT = "org.spongepowered.asm.mixin.injection.Redirect"
 
         private fun isSwitchCase(e: PsiElement): Boolean {