2 files changed, 38 insertions, 9 deletions

diff --git a/src/main/kotlin/com/replaymod/gradle/remap/PsiPattern.kt b/src/main/kotlin/com/replaymod/gradle/remap/PsiPattern.kt
index fcba558..c1cdf0f 100644
--- a/src/main/kotlin/com/replaymod/gradle/remap/PsiPattern.kt
+++ b/src/main/kotlin/com/replaymod/gradle/remap/PsiPattern.kt
@@ -7,7 +7,7 @@ import org.jetbrains.kotlin.psi.psiUtil.endOffset
 import org.jetbrains.kotlin.psi.psiUtil.startOffset
 
 internal class PsiPattern(
-        private val parameters: List<String>,
+        private val parameters: Set<PsiParameter>,
         private val varArgs: Boolean,
         private val pattern: PsiStatement,
         private val replacement: List<String>
@@ -43,6 +43,8 @@ internal class PsiPattern(
 
     inner class Matcher(private val root: PsiElement, private val arguments: MutableList<TextRange> = mutableListOf()) {
 
+        private val localVariables = mutableMapOf<PsiElement, PsiElement>()
+
         fun toChanges(): List<Pair<TextRange, String>> {
             val sortedArgs = arguments.toList().sortedBy { it.startOffset }
             val changes = mutableListOf<Pair<TextRange, String>>()
@@ -97,25 +99,51 @@ internal class PsiPattern(
                     && pattern.classReference?.resolve() == expr.classReference?.resolve()
                     && match(pattern.qualifier, expr.qualifier)
                     && match(pattern.argumentList, expr.argumentList)
+            is PsiLambdaExpression -> expr is PsiLambdaExpression
+                    && match(pattern.parameterList, expr.parameterList)
+                    && match(pattern.body, expr.body)
+            is PsiParameterList -> expr is PsiParameterList
+                    && pattern.parametersCount == expr.parametersCount
+                    && pattern.parameters.zip(expr.parameters).all { (p, e) -> match(p, e) }
             is PsiLiteralExpression -> expr is PsiLiteralExpression
                     && pattern.text == expr.text
             else -> false
         }
 
+        private fun match(pattern: PsiParameter, expr: PsiParameter): Boolean {
+            if (pattern.isVarArgs != expr.isVarArgs) {
+                return false
+            }
+
+            localVariables[pattern] = expr
+            return true
+        }
+
         private fun match(pattern: PsiReferenceExpression, expr: PsiExpression): Boolean {
-            return if (pattern.firstChild is PsiReferenceParameterList && pattern.referenceName in parameters) {
+            val resolvedPattern = pattern.resolve()
+            if (resolvedPattern in parameters) {
                 val patternType = pattern.type ?: return false
                 val exprType = expr.type ?: return false
-                if (patternType.isAssignableFrom(exprType)) {
+                return if (patternType.isAssignableFrom(exprType)) {
                     arguments.add(expr.textRange)
                     true
                 } else {
                     false
                 }
             }
-            else expr is PsiReferenceExpression
-                    && pattern.referenceName == expr.referenceName
-                    && match(pattern.qualifierExpression, expr.qualifierExpression)
+
+            // If the pattern is not a free variable, the expression must match it structurally
+            if (expr !is PsiReferenceExpression) {
+                return false
+            }
+
+            // If the pattern refers to a specific local variable, so must the expression
+            val localVariable = localVariables[resolvedPattern]
+            if (localVariable != null) {
+                return expr.resolve() == localVariable
+            }
+
+            return pattern.referenceName == expr.referenceName && match(pattern.qualifierExpression, expr.qualifierExpression)
         }
 
         private fun match(pattern: PsiExpressionList, expr: PsiExpressionList): Boolean {
@@ -167,6 +195,6 @@ internal class PsiPattern(
         }
 
         private fun isVarArgsParameter(expr: PsiReferenceExpression): Boolean =
-            varArgs && expr.firstChild is PsiReferenceParameterList && expr.referenceName == parameters.last()
+            varArgs && expr.firstChild is PsiReferenceParameterList && expr.resolve() == parameters.last()
     }
 }
\ No newline at end of file
diff --git a/src/main/kotlin/com/replaymod/gradle/remap/PsiPatterns.kt b/src/main/kotlin/com/replaymod/gradle/remap/PsiPatterns.kt
index 14eee07..fe6040c 100644
--- a/src/main/kotlin/com/replaymod/gradle/remap/PsiPatterns.kt
+++ b/src/main/kotlin/com/replaymod/gradle/remap/PsiPatterns.kt
@@ -23,7 +23,8 @@ internal class PsiPatterns(private val annotationFQN: String) {
         val body = method.body!!
         val methodLine = offsetToLineNumber(file.text, body.startOffset)
 
-        val parameters = method.parameterList.parameters.map { it.name }
+        val parameters = method.parameterList.parameters.toSet()
+        val parameterNames = parameters.map { it.name }.toSet()
         val varArgs = method.parameterList.parameters.lastOrNull()?.isVarArgs ?: false
 
         val project = file.project
@@ -55,7 +56,7 @@ internal class PsiPatterns(private val annotationFQN: String) {
             val arguments = mutableListOf<PsiExpression>()
             replacementExpression.accept(object : JavaRecursiveElementVisitor() {
                 override fun visitReferenceExpression(expr: PsiReferenceExpression) {
-                    if (expr.firstChild is PsiReferenceParameterList && expr.referenceName in parameters) {
+                    if (expr.firstChild is PsiReferenceParameterList && expr.referenceName in parameterNames) {
                         arguments.add(expr)
                     } else {
                         super.visitReferenceExpression(expr)