Kotlinコードを解析してコードをコードで書き換えよう

この記事ではKotlinのソースコードを読み込んで解析し、コードの一括置換を行う方法を説明します。

解説

build.gradle.kts

plugins {
    kotlin("jvm") version "2.1.20"
}

repositories {
    mavenCentral()
    maven(url = "https://www.jetbrains.com/intellij-repository/releases")
    maven(url = "https://packages.jetbrains.team/maven/p/ij/intellij-dependencies")
}

dependencies {
    testImplementation(kotlin("test"))
    val kotlinVersion = "2.1.20"

    implementation("org.jetbrains.kotlin:kotlin-stdlib:$kotlinVersion")
    implementation("org.jetbrains.kotlin:kotlin-compiler-embeddable:$kotlinVersion")

}

kotlin {
    jvmToolchain(21)
}

解析コード

まずはenvironmentを作成します。

val disposable = Disposer.newDisposable()
try {
    val environment = KotlinCoreEnvironment.createForProduction(
        disposable,
        CompilerConfiguration().apply {
            put(
                CommonConfigurationKeys.MESSAGE_COLLECTOR_KEY,
                PrintingMessageCollector(System.err, MessageRenderer.PLAIN_FULL_PATHS, false),
            )
        },
        EnvironmentConfigFiles.JVM_CONFIG_FILES,
    )
} finally {
    Disposer.dispose(disposable)
}

ファイルを読み込み、トップレベルの定義から必要なものをフィルターします。

val file = File("path/to/TargetFile.kt").readText()
val ktFile = KtPsiFactory(environment.project, false).createFile(file)
ktFile.declarations
    .filterIsInstance<KtClassOrObject>()
    .onEach { println(it.name) }
ktFile.declarations
    .filterIsInstance<KtObjectDeclaration>()
    .onEach { println(it.name) }
ktFile.declarations
    .filterIsInstance<KtFunction>()
    .onEach { println(it.name) }

実例

Sampleのval property: StringをMigrateにfun property(): String関数として移植します。

sealed interface Sample {
    val property: String

    class One : Sample {
        override val property: String = "One"
    }

    class Two : Sample {
        override val property: String get() = "Two"
    }

    class Three : Sample {
        override val property: String
            get() {
                return "Three"
            }
    }
}

sealed interface Migrate {
    class One : Migrate {
        
    }
    
    class Two : Migrate {
        
    }
    
    class Three : Migrate {
        
    }
}

add系の関数もありますが、これはIntelliJのPluginではないと動かないです。CLIではStringを直接組み立てるしかないので、その方法で行います。
offsetで位置を取得し、コードをinsertしていきます。コードは後でktlint等で整えてください。

val sample = ktFile.declarations
    .filterIsInstance<KtClassOrObject>()
    .first { it.name == "Sample" }

val samplesImplementations = sample.declarations
    .filterIsInstance<KtClassOrObject>()
    .filter { it.superTypeListEntries.any { it.text == "Sample" } }

val builder = StringBuilder(file.readText())
for (sampleImpl in samplesImplementations) {
    val property = sampleImpl.declarations.filterIsInstance<KtProperty>()
        .first { it.name == "property" }

    val text: String

    val getter = property.getter
    if (getter == null) {
        val initializer = property.initializer
        if (initializer != null) {
            // val property = ""
            text = buildString {
                appendLine("fun property(): String {")
                appendLine("return ${initializer.text}")
                appendLine("}")
            }
        } else {
            throw IllegalStateException()
        }
    } else {
        // val property = get()
        if (getter.hasBlockBody()) {
            // get() {}
            text = buildString {
                append("fun property(): String")
                appendLine(getter.bodyExpression!!.text)
            }
        } else {
            // get() = ""
            text = buildString {
                appendLine("fun property(): String {")
                appendLine("return ${getter.bodyExpression!!.text}")
                appendLine("}")
            }
        }
    }

    val migrateObject = KtPsiFactory(environment.project, false)
        .createFile(builder.toString())
        .declarations
        .filterIsInstance<KtClassOrObject>()
        .first { it.name == "Migrate" }
        .declarations
        .filterIsInstance<KtClassOrObject>()
        .first { it.name == sampleImpl.name!! }
    builder.insert(migrateObject.body!!.startOffset + 2, text)
}
file.writeText(builder.toString())