Add new @Pattern feature to centralize version-aware code

That is, most of the business code should not be aware that it is being compiled
to multiple versions even when it heavily interacts with MC, preprocessor
statements should be an escape hatch, not the norm.
Similarly, code should not be forced to do `MCVer.getWindow(mc)` instead of the
much more intuitive `mc.getWindow()`, and this new preprocessor (technically remap)
feature makes this possible by defining "search and replace"-like patterns (but
smarter in that they are type-aware) in one or more central places which then
are applied all over the code base.

In a way, this is another step in the automatic back-porting process where
preprocessor statements are used when we cannot yet do something automatically.
Previously we "merely" automatically converted between different mapping, this
new feature now also allows us to automatically perform simple refactoring
tasks like changing field access to a getter+setter (e.g. `mc.getWindow()`), or
changing how a method is called (e.g. `BufferBuilder.begin`), or changing a
method call chain (e.g. `dispatcher.camera.getYaw()`), or most other
search-and-replace-like changes and any combination of those.
The only major limitation is that the replacement itself is not smart, so
arguments must be kept in same order (or be temporarily assigned to local
variables which then can be used in any order).

4 files changed, 257 insertions, 5 deletions

diff --git a/src/main/kotlin/com/replaymod/gradle/remap/PsiMapper.kt b/src/main/kotlin/com/replaymod/gradle/remap/PsiMapper.kt
index b131478..c51b801 100644
--- a/src/main/kotlin/com/replaymod/gradle/remap/PsiMapper.kt
+++ b/src/main/kotlin/com/replaymod/gradle/remap/PsiMapper.kt
@@ -21,7 +21,11 @@ import org.jetbrains.kotlin.synthetic.SyntheticJavaPropertyDescriptor
 import org.jetbrains.kotlin.synthetic.SyntheticJavaPropertyDescriptor.Companion.propertyNameByGetMethodName
 import java.util.*
 
-internal class PsiMapper(private val map: MappingSet, private val file: PsiFile) {
+internal class PsiMapper(
+        private val map: MappingSet,
+        private val file: PsiFile,
+        private val patterns: PsiPatterns?
+) {
     private val mixinMappings = mutableMapOf<String, ClassMapping<*, *>>()
     private val errors = mutableListOf<Pair<Int, String>>()
     private val changes = TreeMap<TextRange, String>(Comparator.comparing<TextRange, Int> { it.startOffset })
@@ -31,8 +35,9 @@ internal class PsiMapper(private val map: MappingSet, private val file: PsiFile)
         errors.add(Pair(line, message))
     }
 
-    private fun replace(e: PsiElement, with: String) {
-        changes[e.textRange] = with
+    private fun replace(e: PsiElement, with: String) = replace(e.textRange, with)
+    private fun replace(textRange: TextRange, with: String) {
+        changes[textRange] = with
     }
 
     private fun replaceIdentifier(parent: PsiElement, with: String) {
@@ -49,10 +54,19 @@ internal class PsiMapper(private val map: MappingSet, private val file: PsiFile)
 
     private fun valid(e: PsiElement): Boolean {
         val range = e.textRange
+        // FIXME This implementation is technically wrong but some parts of the
+        //       remapper now rely on that, so fixing it is non-trivial.
+        //       For a proper implementation see the TextRange version below.
         val before = changes.ceilingKey(range)
         return before == null || !before.intersects(range)
     }
 
+    private fun valid(range: TextRange): Boolean {
+        val before = changes.floorKey(range) ?: TextRange.EMPTY_RANGE
+        val after = changes.ceilingKey(range) ?: TextRange.EMPTY_RANGE
+        return !before.intersectsStrict(range) && !after.intersectsStrict(range)
+    }
+
     private fun getResult(text: String): Pair<String, List<Pair<Int, String>>> {
         var result = text
         for ((key, value) in changes.descendingMap()) {
@@ -427,7 +441,32 @@ internal class PsiMapper(private val map: MappingSet, private val file: PsiFile)
         })
     }
 
+    private fun applyPatternMatch(matcher: PsiPattern.Matcher) {
+        val changes = matcher.toChanges()
+        if (changes.all { valid(it.first) }) {
+            changes.forEach { (range, text) -> replace(range, text)}
+        } else if (changes.any { it.first !in this.changes }) {
+            System.err.println("Conflicting pattern changes in $file")
+            System.err.println("Proposed changes:")
+            changes.forEach { println("${it.first}: \"${it.second}\" (${if (valid(it.first)) "accepted" else "rejected"})") }
+            System.err.println("Current changes:")
+            this.changes.forEach { println("${it.key}: \"${it.value}\"") }
+        }
+    }
+
     fun remapFile(bindingContext: BindingContext): Pair<String, List<Pair<Int, String>>> {
+        if (patterns != null) {
+            file.accept(object : JavaRecursiveElementVisitor() {
+                override fun visitCodeBlock(block: PsiCodeBlock) {
+                    patterns.find(block).forEach { applyPatternMatch(it) }
+                }
+
+                override fun visitExpression(expression: PsiExpression) {
+                    patterns.find(expression).forEach { applyPatternMatch(it) }
+                }
+            })
+        }
+
         file.accept(object : JavaRecursiveElementVisitor() {
             override fun visitClass(psiClass: PsiClass) {
                 val annotation = psiClass.getAnnotation(CLASS_MIXIN) ?: return
diff --git a/src/main/kotlin/com/replaymod/gradle/remap/PsiPattern.kt b/src/main/kotlin/com/replaymod/gradle/remap/PsiPattern.kt
new file mode 100644
index 0000000..95d6156
--- /dev/null
+++ b/src/main/kotlin/com/replaymod/gradle/remap/PsiPattern.kt
@@ -0,0 +1,107 @@
+package com.replaymod.gradle.remap
+
+import org.jetbrains.kotlin.backend.common.push
+import org.jetbrains.kotlin.com.intellij.openapi.util.TextRange
+import org.jetbrains.kotlin.com.intellij.psi.*
+import org.jetbrains.kotlin.psi.psiUtil.endOffset
+import org.jetbrains.kotlin.psi.psiUtil.startOffset
+
+internal class PsiPattern(
+        private val parameters: List<String>,
+        private val pattern: PsiStatement,
+        private val replacement: List<String>
+) {
+    private fun find(pattern: PsiElement, tree: PsiElement, result: MutableList<Matcher>) {
+        tree.accept(object : JavaRecursiveElementVisitor() {
+            override fun visitElement(element: PsiElement) {
+                val matcher = Matcher(element)
+                if (matcher.match(pattern)) {
+                    result.add(matcher)
+                } else {
+                    super.visitElement(element)
+                }
+            }
+        })
+    }
+
+    fun find(statements: Array<PsiStatement>, result: MutableList<Matcher>) {
+        for (statement in statements) {
+            when (pattern) {
+                is PsiReturnStatement -> find(pattern.returnValue!!, statement, result)
+                else -> find(pattern, statement, result)
+            }
+        }
+    }
+
+    fun find(expr: PsiExpression, result: MutableList<Matcher>) {
+        when (pattern) {
+            is PsiReturnStatement -> find(pattern.returnValue!!, expr, result)
+            else -> find(pattern, expr, result)
+        }
+    }
+
+    inner class Matcher(private val root: PsiElement, private val arguments: MutableList<PsiElement> = mutableListOf()) {
+
+        fun toChanges(): List<Pair<TextRange, String>> {
+            val sortedArgs = arguments.toList().sortedBy { it.startOffset }
+            val changes = mutableListOf<Pair<TextRange, String>>()
+
+            val replacementIter = replacement.iterator()
+            var start = root.startOffset
+            for (argPsi in sortedArgs) {
+                changes.push(Pair(TextRange(start, argPsi.startOffset), replacementIter.next()))
+                start = argPsi.endOffset
+            }
+            changes.push(Pair(TextRange(start, root.endOffset), replacementIter.next()))
+
+            return changes.filterNot { it.first.isEmpty && it.second.isEmpty() }
+        }
+
+        fun match(pattern: PsiElement): Boolean = match(pattern, root)
+
+        private fun match(pattern: PsiElement?, expr: PsiElement?): Boolean = when (pattern) {
+            null -> expr == null
+            is PsiAssignmentExpression -> expr is PsiAssignmentExpression
+                    && match(pattern.lExpression, expr.lExpression)
+                    && match(pattern.rExpression!!, expr.rExpression!!)
+            is PsiBlockStatement -> expr is PsiBlockStatement
+                    && pattern.codeBlock.statementCount == expr.codeBlock.statementCount
+                    && pattern.codeBlock.statements.asSequence().zip(expr.codeBlock.statements.asSequence())
+                    .all { (pattern, expr) -> match(pattern, expr) }
+            is PsiReferenceExpression -> expr is PsiExpression
+                    && match(pattern, expr)
+            is PsiMethodCallExpression -> expr is PsiMethodCallExpression
+                    && match(pattern.methodExpression, expr.methodExpression)
+                    && match(pattern.argumentList, expr.argumentList)
+            is PsiExpressionList -> expr is PsiExpressionList
+                    && pattern.expressionCount == expr.expressionCount
+                    && pattern.expressions.asSequence().zip(expr.expressions.asSequence())
+                    .all { (pattern, expr) -> match(pattern, expr) }
+            is PsiExpressionStatement -> expr is PsiExpressionStatement
+                    && match(pattern.expression, expr.expression)
+            is PsiTypeCastExpression -> expr is PsiTypeCastExpression
+                    && match(pattern.operand, expr.operand)
+            is PsiParenthesizedExpression -> expr is PsiParenthesizedExpression
+                    && match(pattern.expression, expr.expression)
+            is PsiNewExpression -> expr is PsiNewExpression
+                    && match(pattern.argumentList, expr.argumentList)
+            else -> false
+        }
+
+        private fun match(pattern: PsiReferenceExpression, expr: PsiExpression): Boolean {
+            return if (pattern.firstChild is PsiReferenceParameterList && pattern.referenceName in parameters) {
+                val patternType = pattern.type ?: return false
+                val exprType = expr.type ?: return false
+                if (patternType.isAssignableFrom(exprType)) {
+                    arguments.add(expr)
+                    true
+                } else {
+                    false
+                }
+            }
+            else expr is PsiReferenceExpression
+                    && pattern.referenceName == expr.referenceName
+                    && match(pattern.qualifierExpression, expr.qualifierExpression)
+        }
+    }
+}
\ 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
new file mode 100644
index 0000000..1000666
--- /dev/null
+++ b/src/main/kotlin/com/replaymod/gradle/remap/PsiPatterns.kt
@@ -0,0 +1,85 @@
+package com.replaymod.gradle.remap
+
+import com.replaymod.gradle.remap.PsiPattern.Matcher
+import org.jetbrains.kotlin.backend.common.push
+import org.jetbrains.kotlin.com.intellij.openapi.util.text.StringUtil.offsetToLineNumber
+import org.jetbrains.kotlin.com.intellij.psi.*
+import org.jetbrains.kotlin.psi.psiUtil.endOffset
+import org.jetbrains.kotlin.psi.psiUtil.startOffset
+
+internal class PsiPatterns(private val annotationFQN: String) {
+    private val patterns = mutableListOf<PsiPattern>()
+
+    fun read(file: PsiFile, replacementFile: String) {
+        file.accept(object : JavaRecursiveElementVisitor() {
+            override fun visitMethod(method: PsiMethod) {
+                method.getAnnotation(annotationFQN) ?: return
+                addPattern(file, method, replacementFile)
+            }
+        })
+    }
+
+    private fun addPattern(file: PsiFile, method: PsiMethod, replacementFile: String) {
+        val body = method.body!!
+        val methodLine = offsetToLineNumber(file.text, body.startOffset)
+
+        val parameters = method.parameterList.parameters.map { it.name }
+
+        val project = file.project
+        val psiFileFactory = PsiFileFactory.getInstance(project)
+        val replacementPsi = psiFileFactory.createFileFromText(file.language, replacementFile) as PsiJavaFile
+        val replacementClass = replacementPsi.classes.first()
+        val replacementMethod = replacementClass.findMethodsByName(method.name, false).let { candidates ->
+            if (candidates.size > 1) {
+                candidates.find { offsetToLineNumber(replacementFile, it.body!!.startOffset) == methodLine }
+            } else {
+                candidates.firstOrNull()
+            } ?: throw RuntimeException("Failed to find updated method \"${method.name}\" (line ${methodLine + 1})")
+        }
+
+        if (method.text == replacementMethod.text) return
+
+        val replacementExpression = when (val statement = replacementMethod.body!!.statements.last()) {
+            is PsiReturnStatement -> statement.returnValue!!
+            else -> statement
+        }
+
+        val replacement = mutableListOf<String>().also { replacement ->
+            val arguments = mutableListOf<PsiExpression>()
+            replacementExpression.accept(object : JavaRecursiveElementVisitor() {
+                override fun visitReferenceExpression(expr: PsiReferenceExpression) {
+                    if (expr.firstChild is PsiReferenceParameterList && expr.referenceName in parameters) {
+                        arguments.add(expr)
+                    } else {
+                        super.visitReferenceExpression(expr)
+                    }
+                }
+            })
+            val sortedArgs = arguments.toList().sortedBy { it.startOffset }
+            var start = replacementExpression.startOffset
+            for (argPsi in sortedArgs) {
+                replacement.push(replacementFile.slice(start until argPsi.startOffset))
+                start = argPsi.endOffset
+            }
+            replacement.push(replacementFile.slice(start until replacementExpression.endOffset))
+        }
+
+        patterns.add(PsiPattern(parameters, body.statements.last(), replacement))
+    }
+
+    fun find(block: PsiCodeBlock): MutableList<Matcher> {
+        val results = mutableListOf<Matcher>()
+        for (pattern in patterns) {
+            pattern.find(block.statements, results)
+        }
+        return results
+    }
+
+    fun find(expr: PsiExpression): MutableList<Matcher> {
+        val results = mutableListOf<Matcher>()
+        for (pattern in patterns) {
+            pattern.find(expr, results)
+        }
+        return results
+    }
+}
diff --git a/src/main/kotlin/com/replaymod/gradle/remap/Transformer.kt b/src/main/kotlin/com/replaymod/gradle/remap/Transformer.kt
index 85e5acd..dd4215f 100644
--- a/src/main/kotlin/com/replaymod/gradle/remap/Transformer.kt
+++ b/src/main/kotlin/com/replaymod/gradle/remap/Transformer.kt
@@ -41,9 +41,14 @@ import kotlin.system.exitProcess
 
 class Transformer(private val map: MappingSet) {
     var classpath: Array<String>? = null
+    var patternAnnotation: String? = null
 
     @Throws(IOException::class)
-    fun remap(sources: Map<String, String>): Map<String, Pair<String, List<Pair<Int, String>>>> {
+    fun remap(sources: Map<String, String>): Map<String, Pair<String, List<Pair<Int, String>>>> =
+            remap(sources, emptyMap())
+
+    @Throws(IOException::class)
+    fun remap(sources: Map<String, String>, processedSource: Map<String, String>): Map<String, Pair<String, List<Pair<Int, String>>>> {
         val tmpDir = Files.createTempDirectory("remap")
         val disposable = Disposer.newDisposable()
         try {
@@ -90,13 +95,29 @@ class Transformer(private val map: MappingSet) {
                     { scope: GlobalSearchScope -> environment.createPackagePartProvider(scope) }
             )
 
+            val patterns = patternAnnotation?.let { annotationFQN ->
+                val patterns = PsiPatterns(annotationFQN)
+                val annotationName = annotationFQN.substring(annotationFQN.lastIndexOf('.') + 1)
+                for ((unitName, source) in sources) {
+                    if (!source.contains(annotationName)) continue
+                    try {
+                        val patternFile = vfs.findFileByIoFile(tmpDir.resolve(unitName).toFile())!!
+                        val patternPsiFile = psiManager.findFile(patternFile)!!
+                        patterns.read(patternPsiFile, processedSource[unitName]!!)
+                    } catch (e: Exception) {
+                        throw RuntimeException("Failed to read patterns from file \"$unitName\".", e)
+                    }
+                }
+                patterns
+            }
+
             val results = HashMap<String, Pair<String, List<Pair<Int, String>>>>()
             for (name in sources.keys) {
                 val file = vfs.findFileByIoFile(tmpDir.resolve(name).toFile())!!
                 val psiFile = psiManager.findFile(file)!!
 
                 val mapped = try {
-                    PsiMapper(map, psiFile).remapFile(analysis.bindingContext)
+                    PsiMapper(map, psiFile, patterns).remapFile(analysis.bindingContext)
                 } catch (e: Exception) {
                     throw RuntimeException("Failed to map file \"$name\".", e)
                 }