MigrateDataBindingToViewBindingPlugin.kt

import com.android.build.api.dsl.CommonExtension
import java.io.File
import javax.xml.parsers.DocumentBuilderFactory
import javax.xml.transform.TransformerFactory
import javax.xml.transform.dom.DOMSource
import javax.xml.transform.stream.StreamResult
import org.gradle.api.Plugin
import org.gradle.api.Project
import org.gradle.configurationcache.extensions.capitalized
import org.gradle.kotlin.dsl.getByName
import org.gradle.kotlin.dsl.provideDelegate
import org.w3c.dom.Attr
import org.w3c.dom.Document
import org.w3c.dom.Element
import org.w3c.dom.Node

/**
 * Usage:
 * ```
 * # if you applied this plugin to `:app` module
 * ./gradlew :app:migrateDataBindingToViewBinding
 * ./gradlew :app:migrateDataBindingToViewBinding -Ppattern="item_*.xml"
 * ```
 */
class MigrateDataBindingToViewBindingPlugin : Plugin<Project> {
    override fun apply(target: Project) = target.applyPlugin()
}

private fun Project.applyPlugin() {
    val android = extensions.getByName<CommonExtension<*, *, *, *, *>>("android")
    tasks.register("migrateDataBindingToViewBinding") {
        doLast {
            val pattern: String? by project
            layout.projectDirectory.dir("src/main/res/layout").asFileTree
                .matching { pattern?.let { include(it) } }
                .forEach { xmlFile ->
                    println("## Processing: $xmlFile")
                    parseAndRemoveLayout(xmlFile, android.namespace!!)
                }
        }
    }
}

private fun parseAndRemoveLayout(xmlFile: File, namespace: String) {
    val documentBuilderFactory = DocumentBuilderFactory.newInstance()
    val documentBuilder = documentBuilderFactory.newDocumentBuilder()
    val document = documentBuilder.parse(xmlFile)

    val layoutElement = document.documentElement
    if (layoutElement.nodeName != "layout") {
        return
    }

    val dataElement: Node? = layoutElement.getElementsByTagName("data").item(0)
    val variableElements = if (dataElement == null) emptyList() else with(dataElement.childNodes) {
        (0 until length).map { item(it) }.filterIsInstance<Element>().filter { it.nodeName == "variable" }
    }
    val viewElement = with(layoutElement.childNodes) {
        (0 until length).map { item(it) }.filterIsInstance<Element>().first { it.nodeName != "data" }
    }

    val newDocument = documentBuilder.newDocument()
    val newViewElement = newDocument.importNode(viewElement, true)
    newDocument.appendChild(newViewElement)
    with(layoutElement.attributes) {
        (0 until length).map { item(it) }.forEach { attr ->
            newViewElement.attributes.setNamedItemNS(newDocument.importNode(attr, true))
        }
    }

    val bindingAttrs = findAttrsWithAttributesStartingWith(newDocument, "@{")
    bindingAttrs.forEach { (element, attrs) ->
        println("Element: ${element}, Attribute: ${attrs.map { it.value }}")
        attrs.forEach { attr ->
            element.attributes.removeNamedItem(attr.name)
        }
    }

    writeXml(newDocument, xmlFile)

    createViewBindingExtensionKtFile(
        xmlFile,
        variableElements,
        bindingAttrs,
        namespace,
    )
}

private fun findAttrsWithAttributesStartingWith(document: Document, attributePrefix: String): Map<Element, List<Attr>> {
    val matchingAttrs = mutableMapOf<Element, List<Attr>>()
    findAttrsWithAttributesStartingWith(document.documentElement, attributePrefix, matchingAttrs)
    return matchingAttrs
}

private fun findAttrsWithAttributesStartingWith(
    element: Element,
    attributePrefix: String,
    matchingAttrs: MutableMap<Element, List<Attr>>
) {
    val attributes = element.attributes
    for (i in 0 until attributes.length) {
        val attr = attributes.item(i) as Attr
        if (attr.nodeValue.startsWith(attributePrefix)) {
            matchingAttrs[element] = matchingAttrs.getOrDefault(element, emptyList()) + attr
        }
    }

    // Recursive call for all child elements
    val childNodes = element.childNodes
    for (i in 0 until childNodes.length) {
        val node = childNodes.item(i) as? Element ?: continue
        findAttrsWithAttributesStartingWith(node, attributePrefix, matchingAttrs)
    }
}

// extract variable tag from data tag, create ViewBinding class extension with extracted variables.
private fun createViewBindingExtensionKtFile(
    layoutXmlFile: File,
    extractedVariables: List<Element>,
    bindingAttrs: Map<Element, List<Attr>>,
    androidNamespace: String
) {
    val layoutName = layoutXmlFile.nameWithoutExtension.split('_').map { it.capitalized() }.joinToString("")
    val viewBindingClassName = "${layoutName}Binding"
    val viewBindingExtensionFileName = "${layoutName}BindingExt"

    val viewBindingExtensionKtFile = File(
        layoutXmlFile.parentFile.parentFile.parentFile,
        "kotlin/${androidNamespace.replace('.', '/')}/databinding/${viewBindingExtensionFileName}.kt",
    )
    viewBindingExtensionKtFile.delete()
    viewBindingExtensionKtFile.parentFile.mkdirs()

    val viewBindingExtensionKtFileContent = buildString {
        appendLine("/*")
        appendLine(" * WARNING: This file is auto-generated by MigrateDataBindingToViewBindingPlugin.")
        appendLine(" * TODOs:")
        appendLine(" * - [ ] You must edit bind() function by your hand.")
        appendLine(" * - [ ] You must reformat the layout XML code by your hand.")
        appendLine(" * - [ ] Check the view codes that uses binding and fix it if there are errors.")
        appendLine(" * - [ ] Remove this comment if you finished editing.")
        appendLine(" */")
        appendLine("package ${androidNamespace}.databinding")
        appendLine()

        appendLine("import android.util.ArrayMap")
        extractedVariables.forEach { element ->
            val type = element.attributes.getNamedItem("type").nodeValue
            appendLine("import $type")
        }

        appendLine()
        appendLine("""@Suppress("UNCHECKED_CAST")""")
        appendLine("private val $viewBindingClassName.variables: MutableMap<String, Any?>")
        appendLine("    get() {")
        appendLine("        val existing = root.tag as? MutableMap<String, Any?>")
        appendLine("        return if (existing == null) {")
        appendLine("            val new = ArrayMap<String, Any?>()")
        appendLine("            root.tag = new")
        appendLine("            new")
        appendLine("        } else {")
        appendLine("            existing")
        appendLine("        }")
        appendLine("    }")

        extractedVariables.forEach { element ->
            appendLine()
            val name = element.attributes.getNamedItem("name").nodeValue
            val type = element.attributes.getNamedItem("type").nodeValue
            val simpleType = type.split('.').last()
            appendLine("""internal var ${viewBindingClassName}.$name: $simpleType?""")
            appendLine("""    get() = variables["$name"] as? $simpleType""")
            appendLine("""    set(value) {""")
            appendLine("""        variables["$name"] = value""")
            appendLine("""        executeBinding()""")
            appendLine("""    }""")
        }
        appendLine()
        appendLine("""private val $viewBindingClassName.bindRunnable: Runnable""")
        appendLine("""    get() = variables.getOrPut("bindRunnable") { Runnable { bindVariables() } } as Runnable""")
        appendLine()
        appendLine("""internal fun $viewBindingClassName.executeBinding() {""")
        appendLine("""    val runnable = bindRunnable""")
        appendLine("""    root.removeCallbacks(runnable)""")
        appendLine("""    root.post(runnable)""")
        appendLine("""}""")
        appendLine()
        appendLine("private fun $viewBindingClassName.bindVariables() {")
        bindingAttrs.forEach { (element, attrs) ->
            val id = element.attributes.getNamedItem("android:id")?.nodeValue?.removePrefix("@+id/")
            attrs.forEach { attr ->
                appendLine("""    TODO("${element.tagName}(id:$id) ${attr.name}=${attr.value}")""")
            }
        }
        appendLine("}")
    }

    viewBindingExtensionKtFile.writeText(viewBindingExtensionKtFileContent)
    println("Extension created to ${viewBindingExtensionKtFile.absolutePath}")
}

private fun writeXml(document: Document, outputFile: File) {
    try {
        document.xmlStandalone = true
        val transformerFactory = TransformerFactory.newInstance()
        val transformer = transformerFactory.newTransformer()
        val source = DOMSource(document)
        val result = StreamResult(outputFile)
        transformer.transform(source, result)
    } catch (e: Exception) {
        e.printStackTrace()
    }
}