PreLaunch: Rework @Local

See the comment on the Local class for more info. This also adds a list
of Locals to other injectors so their use is a bit more normalized.

src/main/kotlin/com/chattriggers/ctjs/launch/Mappings.kt

@@ -55,8 +55,15 @@ object Mappings {
 
                 methods.getOrPut(method.unmappedName, ::mutableListOf).add(MappedMethod(
                     name = Mapping.fromMapped(method),
-                    parameterTypes = unmappedType.argumentTypes.zip(mappedType.argumentTypes).map {
-                        Mapping(it.first.descriptor, it.second.descriptor)
+                    parameters = method.parameters.map {
+                        MappedParameter(
+                            Mapping(it.unmappedName, it.mappedName),
+                            Mapping(
+                                unmappedType.getArgumentByLvtIndex(it.localVariableIndex).internalName,
+                                mappedType.getArgumentByLvtIndex(it.localVariableIndex).internalName,
+                            ),
+                            it.localVariableIndex,
+                        )
                     },
                     returnType = Mapping(unmappedType.returnType.descriptor, mappedType.returnType.descriptor)
                 ))
@@ -95,9 +102,15 @@ object Mappings {
 
     data class MappedField(val name: Mapping, val type: Mapping)
 
+    class MappedParameter(
+        val name: Mapping,
+        val type: Mapping,
+        val lvtIndex: Int,
+    )
+
     class MappedMethod(
         val name: Mapping,
-        val parameterTypes: List<Mapping>,
+        val parameters: List<MappedParameter>,
         val returnType: Mapping,
     )
 
@@ -134,4 +147,17 @@ object Mappings {
             drop(1).dropLast(1)
         } else this).replace('.', '/')
     }
+
+    fun Type.getArgumentByLvtIndex(index: Int): Type {
+        require(sort == Type.METHOD)
+
+        var currentIndex = 1
+        for (argumentType in argumentTypes) {
+            if (index == currentIndex)
+                return argumentType
+            currentIndex += argumentType.size
+        }
+
+        error("Invalid LVT index $index for method descriptor $this")
+    }
 }

src/main/kotlin/com/chattriggers/ctjs/launch/annotations.kt

@@ -85,15 +85,16 @@ data class Slice(
     val to: At?,
 )
 
-// This intentionally does not have a 'name' field since MC methods will not have
-// an LVT, and thus the injection would always fail. The local must instead be
-// referenced by index
+// There are three ways to capture a local:
+//   - By name, only if the local is a parameter: Local(parameterName = "foo")
+//   - By absolute local index: Local(index = 4)
+//   - By ordinal and type: Local(type = "F", ordinal = 1)
 data class Local(
-    val type: String,
     val print: Boolean?,
-    val ordinal: Int?,
+    val parameterName: String?,
     val index: Int?,
-    val argsOnly: Boolean?,
+    val ordinal: Int?,
+    val type: String?,
     val mutable: Boolean?,
 )
 
@@ -150,7 +151,7 @@ class ModifyArg(
     val slice: Slice?,
     val at: At,
     val captureAllParams: Boolean?,
-    val locals: List<Local>? = null,
+    val locals: List<Local>?,
     val index: Int?,
     val remap: Boolean?,
     val require: Int?,
@@ -172,6 +173,7 @@ class ModifyArgs(
     val method: String,
     val slice: Slice?,
     val at: At,
+    val locals: List<Local>?,
     val remap: Boolean?,
     val require: Int?,
     val expect: Int?,
@@ -192,6 +194,7 @@ class ModifyExpressionValue(
     val method: String,
     val at: At,
     val slice: List<Slice>?,
+    val locals: List<Local>?,
     val remap: Boolean?,
     val require: Int?,
     val expect: Int?,
@@ -211,6 +214,7 @@ class ModifyReceiver(
     val method: String,
     val at: At,
     val slice: List<Slice>?,
+    val locals: List<Local>?,
     val remap: Boolean?,
     val require: Int?,
     val expect: Int?,

src/main/kotlin/com/chattriggers/ctjs/launch/generation/InjectGenerator.kt

@@ -2,11 +2,11 @@ package com.chattriggers.ctjs.launch.generation
 
 import com.chattriggers.ctjs.launch.Inject
 import com.chattriggers.ctjs.launch.Mappings
+import gg.essential.elementa.state.map
 import org.objectweb.asm.Type
 import org.objectweb.asm.tree.MethodNode
 import org.spongepowered.asm.mixin.injection.callback.CallbackInfo
 import org.spongepowered.asm.mixin.injection.callback.CallbackInfoReturnable
-import org.spongepowered.asm.mixin.injection.callback.LocalCapture
 import org.spongepowered.asm.mixin.injection.Inject as SPInject
 
 class InjectGenerator(
@@ -19,8 +19,10 @@ class InjectGenerator(
     override fun getInjectionSignature(): InjectionSignature {
         val (_, name, possibleParameters) = Utils.splitMethodDescriptor(inject.method)
         val (mappedMethod, method) = ctx.findMethod(name, possibleParameters)
-        val parameters = mappedMethod.parameterTypes.mapTo(mutableListOf()) {
-            Parameter(Type.getType(it.value), null)
+        val parameters = mutableListOf<Parameter>()
+
+        inject.locals?.forEach {
+            parameters.add(Parameter(Utils.getLocalType(it, mappedMethod), it))
         }
 
         if (mappedMethod.returnType.value == "V") {
@@ -29,11 +31,6 @@ class InjectGenerator(
             parameters.add(Parameter(Type.getType(CallbackInfoReturnable::class.java), null))
         }
 
-        inject.locals?.forEach {
-            val type = Type.getType(Mappings.getMappedClassName(it.type) ?: it.type)
-            parameters.add(Parameter(type, it))
-        }
-
         return InjectionSignature(
             name,
             parameters,

src/main/kotlin/com/chattriggers/ctjs/launch/generation/InjectorGenerator.kt

@@ -60,7 +60,7 @@ abstract class InjectorGenerator(protected val ctx: GenerationContext, val id: I
                     local.print?.let { visit("print", it) }
                     local.ordinal?.let { visit("ordinal", it) }
                     local.index?.let { visit("index", it) }
-                    local.argsOnly?.let { visit("argsOnly", it) }
+                    // local.argsOnly?.let { visit("argsOnly", it) }
                 }
             }
 

src/main/kotlin/com/chattriggers/ctjs/launch/generation/ModifyArgGenerator.kt

@@ -34,30 +34,30 @@ class ModifyArgGenerator(
             targetDescriptor.parameters,
         )
 
-        if (modifyArg.index != null && modifyArg.index !in targetMappedMethod.parameterTypes.indices)
+        if (modifyArg.index != null && modifyArg.index !in targetMappedMethod.parameters.indices)
             error("ModifyArg received an out-of-bounds index ${modifyArg.index}")
 
         val parameters = when {
-            modifyArg.captureAllParams == true ->
-                targetMappedMethod.parameterTypes.mapTo(mutableListOf()) { Parameter(Type.getType(it.value), null) }
+            modifyArg.captureAllParams == true -> targetMappedMethod.parameters.mapTo(mutableListOf()) {
+                Parameter(Type.getType(it.type.value), null)
+            }
             modifyArg.index == null -> error("ModifyArg must capture all parameters or specify an index")
             else -> {
                 mutableListOf(
-                    Parameter(Type.getType(targetMappedMethod.parameterTypes[modifyArg.index].value), null)
+                    Parameter(Type.getType(targetMappedMethod.parameters[modifyArg.index].type.value), null)
                 )
             }
         }
 
         val returnType = if (modifyArg.index != null) {
-            Type.getType(targetMappedMethod.parameterTypes[modifyArg.index].value)
+            Type.getType(targetMappedMethod.parameters[modifyArg.index].type.value)
         } else {
             check(parameters.size == 1) // Should be true given the above checks
             parameters[0].type
         }
 
         modifyArg.locals?.forEach {
-            val type = Type.getType(Mappings.getMappedClassName(it.type) ?: it.type)
-            parameters.add(Parameter(type, it))
+            parameters.add(Parameter(Utils.getLocalType(it, mappedMethod), it))
         }
 
         return InjectionSignature(

src/main/kotlin/com/chattriggers/ctjs/launch/generation/ModifyArgsGenerator.kt

@@ -21,8 +21,8 @@ class ModifyArgsGenerator(
 
         val parameters = mutableListOf<Parameter>()
         parameters.add(Parameter(Type.getType(Args::class.java), null))
-        mappedMethod.parameterTypes.forEach {
-            parameters.add(Parameter(Type.getType(it.value), null))
+        modifyArgs.locals?.forEach {
+            parameters.add(Parameter(Utils.getLocalType(it, mappedMethod), null))
         }
 
         return InjectionSignature(

src/main/kotlin/com/chattriggers/ctjs/launch/generation/ModifyExpressionValueGenerator.kt

@@ -25,9 +25,13 @@ class ModifyExpressionValueGenerator(
             is Utils.ConstantAtTarget -> atTarget.type
         }
 
+        val parameters = listOf(Parameter(exprType, null)) + modifyExpressionValue.locals?.map {
+            Parameter(Utils.getLocalType(it, mappedMethod), null)
+        }.orEmpty()
+
         return InjectionSignature(
             name,
-            listOf(Parameter(exprType, null)),
+            parameters,
             exprType,
             method.isStatic,
             mappedMethod,

src/main/kotlin/com/chattriggers/ctjs/launch/generation/ModifyReceiverGenerator.kt

@@ -30,9 +30,13 @@ class ModifyReceiverGenerator(
             else -> error("Unsupported At.value for ModifyReceiver: ${atTarget.name}")
         }
 
+        val params = listOf(Parameter(owner, null)) +
+            extraParams.map { Parameter(it, null) } +
+            modifyReceiver.locals?.map { Parameter(Utils.getLocalType(it, mappedMethod), null) }.orEmpty()
+
         return InjectionSignature(
             name,
-            listOf(Parameter(owner, null)) + extraParams.map { Parameter(it, null) },
+            params,
             owner,
             method.isStatic,
             mappedMethod,

src/main/kotlin/com/chattriggers/ctjs/launch/generation/RedirectGenerator.kt

@@ -21,11 +21,9 @@ class RedirectGenerator(
         val parameters = redirect.args.orEmpty().mapTo(mutableListOf()) {
             Parameter(Type.getType(Mappings.getMappedClassName(it) ?: it), null)
         }
-        mappedMethod.parameterTypes.mapTo(parameters) { Parameter(Type.getType(it.value), null) }
 
         redirect.locals?.forEach {
-            val type = Type.getType(Mappings.getMappedClassName(it.type) ?: it.type)
-            parameters.add(Parameter(type, it))
+            parameters.add(Parameter(Utils.getLocalType(it, mappedMethod), it))
         }
 
         return InjectionSignature(

src/main/kotlin/com/chattriggers/ctjs/launch/generation/Utils.kt

@@ -1,9 +1,6 @@
 package com.chattriggers.ctjs.launch.generation
 
-import com.chattriggers.ctjs.launch.At
-import com.chattriggers.ctjs.launch.Desc
-import com.chattriggers.ctjs.launch.Mappings
-import com.chattriggers.ctjs.launch.Slice
+import com.chattriggers.ctjs.launch.*
 import net.fabricmc.accesswidener.AccessWidenerReader
 import net.fabricmc.loader.impl.FabricLoaderImpl
 import org.objectweb.asm.Opcodes
@@ -50,7 +47,7 @@ internal object Utils {
                 val mappedTarget = buildString {
                     append("L${mappedClass.name.value};")
                     append(mappedMethod.name.value)
-                    append(mappedMethod.parameterTypes.joinToString("", prefix = "(", postfix = ")") { it.value })
+                    append(mappedMethod.parameters.joinToString("", prefix = "(", postfix = ")") { it.type.value })
                     append(mappedMethod.returnType.value)
                 }
 
@@ -172,7 +169,7 @@ internal object Utils {
 
         val desc = buildString {
             append('(')
-            append(mappedMethod.parameterTypes.joinToString("") { it.value })
+            append(mappedMethod.parameters.joinToString("") { it.type.value })
             append(')')
             append(mappedMethod.returnType.value)
         }
@@ -209,15 +206,15 @@ internal object Utils {
             ?: error("Cannot find name $unmappedName in class ${mappedClass.name.original}")
 
         for (method in mappedMethods) {
-            if (method.parameterTypes.size != parameters.size)
+            if (method.parameters.size != parameters.size)
                 continue
 
-            if (method.parameterTypes.zip(parameters).any { it.first.original != it.second })
+            if (method.parameters.zip(parameters).any { it.first.type.original != it.second })
                 continue
 
             val desc = Type.getMethodDescriptor(
                 Type.getType(method.returnType.value),
-                *method.parameterTypes.map { Type.getType(it.value) }.toTypedArray(),
+                *method.parameters.map { Type.getType(it.type.value) }.toTypedArray(),
             )
 
             val result = classInfo.findMethodInHierarchy(
@@ -246,7 +243,7 @@ internal object Utils {
         for (method in mappedMethods) {
             val desc = Type.getMethodDescriptor(
                 Type.getType(method.returnType.value),
-                *method.parameterTypes.map { Type.getType(it.value) }.toTypedArray(),
+                *method.parameters.map { Type.getType(it.type.value) }.toTypedArray(),
             )
 
             val result = classInfo.findMethodInHierarchy(
@@ -269,6 +266,36 @@ internal object Utils {
         error("Unable to find method $unmappedName in class ${mappedClass.name.original}")
     }
 
+    fun getLocalType(local: Local, method: Mappings.MappedMethod): Type {
+        return when {
+            local.print == true -> {
+                // The type doesn't matter, it won't actually be applied
+                "I"
+            }
+            local.parameterName != null -> {
+                require(local.type == null && local.index == null && local.ordinal == null) {
+                    "Local that specifies parameterName cannot specify type, index, or ordinal"
+                }
+                method.parameters.find { p -> p.name.original == local.parameterName }?.type?.value
+                    ?: error("Could not find parameter \"${local.parameterName}\" in method ${method.name.original}")
+            }
+            local.type != null -> {
+                if (local.index != null) {
+                    require(local.ordinal == null) {
+                        "Local that specifies a type and index cannot specify an ordinal"
+                    }
+                } else {
+                    require(local.ordinal != null) {
+                        "Local that specifies a type must also specify an index or ordinal"
+                    }
+                }
+                Mappings.getMappedClassName(local.type) ?: local.type
+            }
+            else -> error("Local must specify one of the following options: \"print\"; \"parameterName\"; " +
+                "\"type\" and either \"ordinal\" or \"index\"")
+        }.let(Type::getType)
+    }
+
     sealed class AtTarget(val name: String)
 
     data class InvokeAtTarget(val owner: Type, val parameters: List<Type>, val returnType: Type) : AtTarget("INVOKE")