openrs2/deob/src/main/java/dev/openrs2/deob/transform/OpaquePredicateTransformer.kt

package dev.openrs2.deob.transform

import com.github.michaelbull.logging.InlineLogger
import dev.openrs2.asm.InsnMatcher
import dev.openrs2.asm.MemberRef
import dev.openrs2.asm.classpath.ClassPath
import dev.openrs2.asm.classpath.Library
import dev.openrs2.asm.hasCode
import dev.openrs2.asm.transform.Transformer
import org.objectweb.asm.Opcodes
import org.objectweb.asm.tree.*

class OpaquePredicateTransformer : Transformer() {
    private val flowObstructors = mutableSetOf<MemberRef>()
    private var opaquePredicates = 0
    private var stores = 0

    override fun preTransform(classPath: ClassPath) {
        flowObstructors.clear()
        opaquePredicates = 0
        stores = 0

        for (library in classPath.libraries) {
            for (clazz in library) {
                for (method in clazz.methods) {
                    if (method.hasCode()) {
                        findFlowObstructors(library, method)
                    }
                }
            }
        }

        logger.info { "Identified flow obstructors $flowObstructors" }
    }

    private fun findFlowObstructors(library: Library, method: MethodNode) {
        FLOW_OBSTRUCTOR_INITIALIZER_MATCHER.match(method).forEach {
            // add flow obstructor to set
            val putstatic = it.last() as FieldInsnNode
            flowObstructors.add(MemberRef(putstatic))

            // remove initializer
            it.forEach(method.instructions::remove)

            // remove field
            val owner = library[putstatic.owner]!!
            owner.fields.removeIf { it.name == putstatic.name && it.desc == putstatic.desc }
        }
    }

    private fun isFlowObstructor(insn: FieldInsnNode): Boolean {
        return flowObstructors.contains(MemberRef(insn))
    }

    private fun isOpaquePredicate(method: MethodNode, match: List<AbstractInsnNode>): Boolean {
        val load = match[0]

        // flow obstructor loaded directly?
        if (load is FieldInsnNode && load.opcode == Opcodes.GETSTATIC) {
            return isFlowObstructor(load)
        }

        // flow obstructor loaded via local variable
        val iload = load as VarInsnNode
        return STORE_MATCHER.match(method).anyMatch {
            val getstatic = it[0] as FieldInsnNode
            val istore = it[1] as VarInsnNode
            return@anyMatch isFlowObstructor(getstatic) && iload.`var` == istore.`var`
        }
    }

    private fun isRedundantStore(match: List<AbstractInsnNode>): Boolean {
        val getstatic = match[0] as FieldInsnNode
        return isFlowObstructor(getstatic)
    }

    override fun transformCode(classPath: ClassPath, library: Library, clazz: ClassNode, method: MethodNode): Boolean {
        // find and fix opaque predicates
        OPAQUE_PREDICATE_MATCHER.match(method).filter { isOpaquePredicate(method, it) }.forEach {
            val branch = it[1] as JumpInsnNode
            if (branch.opcode == Opcodes.IFEQ) {
                // branch is always taken
                method.instructions.remove(it[0])
                branch.opcode = Opcodes.GOTO
            } else { // IFNE
                // branch is never taken
                it.forEach(method.instructions::remove)
            }

            opaquePredicates++
        }

        // remove redundant stores
        STORE_MATCHER.match(method).filter(this::isRedundantStore).forEach {
            it.forEach(method.instructions::remove)
            stores++
        }

        return false
    }

    override fun postTransform(classPath: ClassPath) {
        logger.info { "Removed $opaquePredicates opaque predicates and $stores redundant stores" }
    }

    companion object {
        private val logger = InlineLogger()
        private val FLOW_OBSTRUCTOR_INITIALIZER_MATCHER =
            InsnMatcher.compile("(GETSTATIC | ILOAD) IFEQ (((GETSTATIC ISTORE)? IINC ILOAD) | ((GETSTATIC | ILOAD) IFEQ ICONST GOTO ICONST)) PUTSTATIC")
        private val OPAQUE_PREDICATE_MATCHER = InsnMatcher.compile("(GETSTATIC | ILOAD) (IFEQ | IFNE)")
        private val STORE_MATCHER = InsnMatcher.compile("GETSTATIC ISTORE")
    }
}