Writing a Kotlin Compiler Plugin

Compiler plugins provide a gateway to extend the capabilities of the Kotlin compiler, allowing developers to inject custom logic into the compilation pipeline.

Without compiler plugins, Android Jetpack Compose would have never been possible. Compiler plugins enable the use of Aspect Oriented Programming to apply cross-cutting logic for use cases when Kotlin's language is not modular enough, or create your own implementation of Kotlin’s Coroutine.

In this session, we will explain the fundamentals of Kotlin's Internal Representation (IR) and explore the details of compiler phases. We will then showcase some meta-programming samples, illustrating how to inspect and modify the output of the compiler.

Transcript

1. Writing a Kotlin Compiler Plugin
2. Xavier F. GOUCHET Staff Engineer @ Joko Android since 2009 Kotlin since 2016
3. AI is about to change shopping forever We’re hiring! welcometothejungle.com/en/companies/joko
4. Introduction 01. Programming your programming…
5. Metaprogramming? • Code that can generate or modify code • Can happen at runtime or compile time • Exists in Ruby, Python, Java, Typescript, Javascript… and Kotlin
6. Java -> Java APT KSP Annotation Based Generate code Kotlin -> * KCP Access everything Generate/Modify code Compatible with KMP Kotlin -> Kotlin Annotation Based Generate code Meta Programming Kotlin -> Java Annotation Based KAPT Generate code Compatible with KMP
7. Why use a KCP? Logging Performances Security Transactions Serialization Validation Cache Notification Multithreading Memoization Testing UI/UX
8. Why use a KCP? Logging Performances Security Transactions Serialization Validation Cache Notification Multithreading Memoization Testing UI/UX
9. Why use a KCP? Logging Performances Security Transactions Serialization Validation Cache Notification Multithreading Memoization Testing Coroutines Kotlinx serialization Power Assert UI/UX Compose
10. Kotlin Compiler 02. Let’s dive into it…
11. A history… 2011 Project Inception 2016 Kotlin 1.0 2023 KMM Stable 2019 Kotlin becomes preferred Android Language 2012 2017 2024 2021 Open Sourced Official on Android KMP announced K2 Compiler Stable Jetpack Compose
12. JS *.js WASM *.wasm Kotlin Compilation Kotlin *.kt Kotlin Compiler JVM *.class Native *.so
13. Compiler Frontend Kotlin *.kt AST FIR AST = Abstract Syntax Tree FIR = Frontend Internal Representation override fun toString(): String { return "…" } function name toString params return kotlin.String
14. Compiler Backend FIR IR (+Lowering) JS *.js WASM *.wasm JVM *.class Native *.so IR = Internal Representation
15. Writing a KCP 03. Show me the code!
16. Setup the project • github.com/Kotlin/compiler-plugin-template • plugin-annotations • compiler-plugin • gradle-plugin
17. class SampleCompilerPluginRegistrar : CompilerPluginRegistrar() { override val supportsK2: Boolean = true override fun ExtensionStorage.registerExtensions( configuration: CompilerConfiguration) { FirExtensionRegistrarAdapter.registerExtension( SampleFirExtensionRegistrar() ) IrGenerationExtension.registerExtension( SampleIrGenerationExtension() ) } } Plugin entry point
18. class SampleCompilerPluginRegistrar : CompilerPluginRegistrar() { override val supportsK2: Boolean = true override fun ExtensionStorage.registerExtensions( configuration: CompilerConfiguration) { FirExtensionRegistrarAdapter.registerExtension( SampleFirExtensionRegistrar() ) IrGenerationExtension.registerExtension( SampleIrGenerationExtension() ) } } Plugin entry point
19. class SampleCompilerPluginRegistrar : CompilerPluginRegistrar() { override val supportsK2: Boolean = true override fun ExtensionStorage.registerExtensions( configuration: CompilerConfiguration) { FirExtensionRegistrarAdapter.registerExtension( SampleFirExtensionRegistrar() ) IrGenerationExtension.registerExtension( SampleIrGenerationExtension() ) } } Plugin entry point
20. class SampleCompilerPluginRegistrar : CompilerPluginRegistrar() { override val supportsK2: Boolean = true override fun ExtensionStorage.registerExtensions( configuration: CompilerConfiguration) { FirExtensionRegistrarAdapter.registerExtension( SampleFirExtensionRegistrar() ) IrGenerationExtension.registerExtension( SampleIrGenerationExtension() ) } } Plugin entry point
21. class SampleCompilerPluginRegistrar : CompilerPluginRegistrar() { override val supportsK2: Boolean = true override fun ExtensionStorage.registerExtensions( configuration: CompilerConfiguration) { FirExtensionRegistrarAdapter.registerExtension( SampleFirExtensionRegistrar() ) IrGenerationExtension.registerExtension( SampleIrGenerationExtension() ) } } Plugin entry point
22. Now what ? • Create Diagnostics • Modify existing code • Generate new code
23. Create diagnostics • Inspect the FIR • Raise warnings / errors
24. fun doSomething(user : User) { requireNotNull(user.phoneNumber) check(user.age > 18) } becomes ❗ requireNotNull() can make your application crash and must be avoided ❗ check() can make your application crash and must be avoided Goal
25. class CustomChecksExtension(session: FirSession) : FirAdditionalCheckersExtension(session) { override val declarationCheckers: DeclarationCheckers = … override val expressionCheckers: ExpressionCheckers = … override val typeCheckers: TypeCheckers = … } Compile Time Checks (FIR side)
26. class CustomChecksExtension(session: FirSession) : FirAdditionalCheckersExtension(session) { override val declarationCheckers: DeclarationCheckers = … override val expressionCheckers: ExpressionCheckers = … override val typeCheckers: TypeCheckers = … } Compile Time Checks (FIR side)
27. class CustomChecksExtension(session: FirSession) : FirAdditionalCheckersExtension(session) { override val declarationCheckers: DeclarationCheckers = … override val expressionCheckers: ExpressionCheckers = object : ExpressionCheckers() { override val functionCallCheckers: Set<FirFunctionCallChecker> = setOf(UnsafeCallChecker(session, compileCheckAsErrors)) } override val typeCheckers: TypeCheckers = … } Compile Time Checks (FIR side)
28. class UnsafeCallChecker( val session: FirSession ) : FirFunctionCallChecker(MppCheckerKind.Common) { val invalidCallDiagnostic by error0<KtExpression>() context(context: CheckerContext, reporter: DiagnosticReporter) override fun check(expression: FirFunctionCall) { if (isInvalid(expression)) { reporter.reportOn(expression.source, invalidCallDiagnostic) } } } Compile Time Checks (FIR side)
29. class UnsafeCallChecker( val session: FirSession ) : FirFunctionCallChecker(MppCheckerKind.Common) { val invalidCallDiagnostic by error0<KtExpression>() context(context: CheckerContext, reporter: DiagnosticReporter) override fun check(expression: FirFunctionCall) { if (isInvalid(expression)) { reporter.reportOn(expression.source, invalidCallDiagnostic) } } } Compile Time Checks (FIR side)
30. class UnsafeCallChecker( val session: FirSession ) : FirFunctionCallChecker(MppCheckerKind.Common) { val invalidCallDiagnostic by error0<KtExpression>() context(context: CheckerContext, reporter: DiagnosticReporter) override fun check(expression: FirFunctionCall) { if (isInvalid(expression)) { reporter.reportOn(expression.source, invalidCallDiagnostic) } } } Compile Time Checks (FIR side)
31. class UnsafeCallChecker( val session: FirSession ) : FirFunctionCallChecker(MppCheckerKind.Common) { val invalidCallDiagnostic by error0<KtExpression>() context(context: CheckerContext, reporter: DiagnosticReporter) override fun check(expression: FirFunctionCall) { if (isInvalid(expression)) { reporter.reportOn(expression.source, invalidCallDiagnostic) } } } Compile Time Checks (FIR side)
32. class UnsafeCallChecker( val session: FirSession ) : FirFunctionCallChecker(MppCheckerKind.Common) { val invalidCallDiagnostic by error0<KtExpression>() context(context: CheckerContext, reporter: DiagnosticReporter) override fun check(expression: FirFunctionCall) { if (isInvalid(expression)) { reporter.reportOn(expression.source, invalidCallDiagnostic) } } } Compile Time Checks (FIR side)
33. Result
34. Going further Let user disable diagnostic with an annotation
35. Modifying code • Manipulate code in the IR extension
36. @LogMethod fun doSomething(i : Int, s: String) { // … } becomes fun doSomething(i : Int, s: String) { println("⇢ doSomething(i=$i, s=$s”) // … } Goal
37. class LogMethodCallVisitor( val pluginContext: IrPluginContext, val messageCollector: MessageCollector ) : IrVisitorVoid() { // … } Modify code (IR side)
38. override fun visitFunction(declaration: IrFunction) { super.visitFunction(declaration) val isAnnotated = declaration.annotations.any { it.isAnnotationWithEqualFqName(ANNOTATION_FQNAME) } if (isAnnotated) { declaration.body = declaration.body?.let { wrapFunction(declaration, it) } } } Modify code (IR side)
39. override fun visitFunction(declaration: IrFunction) { super.visitFunction(declaration) val isAnnotated = declaration.annotations.any { it.isAnnotationWithEqualFqName(ANNOTATION_FQNAME) } if (isAnnotated) { declaration.body = declaration.body?.let { wrapFunction(declaration, it) } } } Modify code (IR side)
40. override fun visitFunction(declaration: IrFunction) { super.visitFunction(declaration) val isAnnotated = declaration.annotations.any { it.isAnnotationWithEqualFqName(ANNOTATION_FQNAME) } if (isAnnotated) { declaration.body = declaration.body?.let { wrapFunction(declaration, it) } } } Modify code (IR side)
41. override fun visitFunction(declaration: IrFunction) { super.visitFunction(declaration) val isAnnotated = declaration.annotations.any { it.isAnnotationWithEqualFqName(ANNOTATION_FQNAME) } if (isAnnotated) { declaration.body = declaration.body?.let { wrapFunction(declaration, it) } } } Modify code (IR side)
42. fun wrapFunction(function: IrFunction, body: IrBody): IrBlockBody { return DeclarationIrBuilder(pluginContext, function.symbol) .irBlockBody { +logFunStart(function) for (statement in body.statements) +statement } } Modify code (IR side)
43. fun wrapFunction(function: IrFunction, body: IrBody): IrBlockBody { return DeclarationIrBuilder(pluginContext, function.symbol) .irBlockBody { +logFunStart(function) for (statement in body.statements) +statement } } Modify code (IR side)
44. fun IrBuilderWithScope.logFunStart(function: IrFunction): IrCall { val concat = irConcat() concat.addArgument(irString("⇢ ${function.name}(")) for ((index, valueParameter) in function.parameters.withIndex()) { if (index > 0) concat.addArgument(irString(", ")) concat.addArgument(irString("${valueParameter.name}=")) concat.addArgument(irGet(valueParameter)) } concat.addArgument(irString(")")) return irCall(funPrintln).also { call -> call.arguments[0] = concat } } Modify code (IR side)
45. fun IrBuilderWithScope.logFunStart(function: IrFunction): IrCall { val concat = irConcat() concat.addArgument(irString("⇢ ${function.name}(")) for ((index, valueParameter) in function.parameters.withIndex()) { if (index > 0) concat.addArgument(irString(", ")) concat.addArgument(irString("${valueParameter.name}=")) concat.addArgument(irGet(valueParameter)) } concat.addArgument(irString(")")) return irCall(funPrintln).also { call -> call.arguments[0] = concat } } Modify code (IR side)
46. fun IrBuilderWithScope.logFunStart(function: IrFunction): IrCall { val concat = irConcat() concat.addArgument(irString("⇢ ${function.name}(")) for ((index, valueParameter) in function.parameters.withIndex()) { if (index > 0) concat.addArgument(irString(", ")) concat.addArgument(irString("${valueParameter.name}=")) concat.addArgument(irGet(valueParameter)) } concat.addArgument(irString(")")) return irCall(funPrintln).also { call -> call.arguments[0] = concat } } Modify code (IR side)
47. fun IrBuilderWithScope.logFunStart(function: IrFunction): IrCall { val concat = irConcat() concat.addArgument(irString("⇢ ${function.name}(")) for ((index, valueParameter) in function.parameters.withIndex()) { if (index > 0) concat.addArgument(irString(", ")) concat.addArgument(irString("${valueParameter.name}=")) concat.addArgument(irGet(valueParameter)) } concat.addArgument(irString(")")) return irCall(funPrintln).also { call -> call.arguments[0] = concat } } Modify code (IR side)
48. fun IrBuilderWithScope.logFunStart(function: IrFunction): IrCall { val concat = irConcat() concat.addArgument(irString("⇢ ${function.name}(")) for ((index, valueParameter) in function.parameters.withIndex()) { if (index > 0) concat.addArgument(irString(", ")) concat.addArgument(irString("${valueParameter.name}=")) concat.addArgument(irGet(valueParameter)) } concat.addArgument(irString(")")) return irCall(funPrintln).also { call -> call.arguments[0] = concat } } Modify code (IR side)
49. fun IrBuilderWithScope.logFunStart(function: IrFunction): IrCall { val concat = irConcat() concat.addArgument(irString("⇢ ${function.name}(")) for ((index, valueParameter) in function.parameters.withIndex()) { if (index > 0) concat.addArgument(irString(", ")) concat.addArgument(irString("${valueParameter.name}=")) concat.addArgument(irGet(valueParameter)) } concat.addArgument(irString(")")) return irCall(funPrintln).also { call -> call.arguments[0] = concat } } Modify code (IR side)
50. Result
51. Result
52. Generating new code • Declare the new Class and functions in FIR extension • Inject the function body in IR extension
53. Warning Generated code only exist after compilation, not in the IDE Use KSP instead if you can
54. Testing… 04. Because you should always test…
55. kotlin-compiler-devkit IDE plugin Testing Automated framework
56. // RUN_PIPELINE_TILL: FRONTEND import com.droidcon.kcp.AllowUnsafeCalls fun testCheck() { val s = (3 * 7 * 9) <!CHECK_CALL!>check(s % 11 == 0)<!> } @AllowUnsafeCalls fun testAllowed() { val s = (208 + 195) check(s % 13 == 0) } Test Diagnostics: unsafeCall.kt
57. // RUN_PIPELINE_TILL: FRONTEND import com.droidcon.kcp.AllowUnsafeCalls fun testCheck() { val s = (3 * 7 * 9) <!CHECK_CALL!>check(s % 11 == 0)<!> } @AllowUnsafeCalls fun testAllowed() { val s = (208 + 195) check(s % 13 == 0) } Test Diagnostics: unsafeCall.kt
58. // RUN_PIPELINE_TILL: FRONTEND import com.droidcon.kcp.AllowUnsafeCalls fun testCheck() { val s = (3 * 7 * 9) <!CHECK_CALL!>check(s % 11 == 0)<!> } @AllowUnsafeCalls fun testAllowed() { val s = (208 + 195) check(s % 13 == 0) } Test Diagnostics: unsafeCall.kt
59. // RUN_PIPELINE_TILL: FRONTEND import com.droidcon.kcp.AllowUnsafeCalls fun testCheck() { val s = (3 * 7 * 9) <!CHECK_CALL!>check(s % 11 == 0)<!> } @AllowUnsafeCalls fun testAllowed() { val s = (208 + 195) check(s % 13 == 0) } Test Diagnostics: unsafeCall.kt
60. // RUN_PIPELINE_TILL: FRONTEND import com.droidcon.kcp.AllowUnsafeCalls fun testCheck() { val s = (3 * 7 * 9) <!CHECK_CALL!>check(s % 11 == 0)<!> } @AllowUnsafeCalls fun testAllowed() { val s = (208 + 195) check(s % 13 == 0) } Test Diagnostics: unsafeCall.kt
61. Testing code generation Each file.kt in the test data has a matching file.fir.txt and file.fir.ir.txt
62. // WITH_STDLIB import com.droidcon.kcp.LogMethod fun box(): String { val result = logMePlease(42, "Lorem ipsum dolor sit amet…") return "OK" } @LogMethod fun logMePlease(i: Int, s: String): String? { return if (i >= 0 && i < s.length) "[" + s.get(i) + "]" else null } Test Code Gen: logMethod.kt
63. FILE: logMethod.kt public final fun box(): R|kotlin/String| { lval result: R|kotlin/String?| = R|/logMePlease|(Int(42), String(Lorem ipsum dolor sit amet…)) ^box String(OK) } @R|com/droidcon/kcp/LogMethod|() public final fun logMePlease(i: R|kotlin/Int|, s: R|kotlin/String|): R|kotlin/String?| { ^logMePlease when () { CMP(>=, R|<local>/i|.R|kotlin/Int.compareTo|(Int(0))) && CMP(<, R|<local>/i|.R|kotlin/Int.compareTo|(R|<local>/s|.R|kotlin/String.length|)) -> { String([).R|kotlin/String.plus|(R|<local>/s|.R|kotlin/String.get|(R|<local>/i| )).R|kotlin/String.plus|(String(])) Test Code Gen: logMethod.fir.txt
64. FILE fqName:<root> fileName:/logMethod.kt FUN name:box visibility:public modality:FINAL returnType:kotlin.String BLOCK_BODY VAR name:result type:kotlin.String? [val] CALL 'public final fun logMePlease (i: kotlin.Int, s: kotlin.String): kotlin.String? declared in <root>' type=kotlin.String? origin=null ARG i: CONST Int type=kotlin.Int value=42 ARG s: CONST String type=kotlin.String value="Lorem ipsum dolor sit amet…" RETURN type=kotlin.Nothing from='public final fun box (): kotlin.String declared in <root>' CONST String type=kotlin.String value="OK" FUN name:logMePlease visibility:public modality:FINAL returnType:kotlin.String? VALUE_PARAMETER kind:Regular name:i index:0 type:kotlin.Int VALUE_PARAMETER kind:Regular name:s index:1 type:kotlin.String annotations: LogMethod BLOCK_BODY CALL 'public final fun println (message: kotlin.Any?): kotlin.Unit declared in kotlin.io' type=kotlin.Unit origin=null Test Code Gen: logMethod.fir.ir.txt
65. Going further 05. The fun doesn’t stop there…
66. Expose them in Gradle plugin Compiler Options Adjust the plugin’s behavior Additional inputs - Files - Primitives (Strings, Numbers, …) - …
67. FIR Tree (IDE plugin)
68. Injekt Plugins fun-cache Allows caching function results Dependency injection at compile time Poko Generates equals() and toString() on annotated class github.com/linhnvtit/fun-cache github.com/drewhamilton/Poko github.com/IVIanuu/injekt Mokkery Generates Mocks at compile time github.com/lupuuss/Mokkery
69. Thanks! Any question?