KSP for Code-Generation

ksp-for-code-generation
Photo by Birmingham Museums Trust on Unsplash

As a developer, what we’re doing every day is writing code. Not just writing unconsciously, we all wish we can write flawless code that has no bugs. To do so, there are lots of tools that help us to increase the quality like unit tests and static analysis. But still, no one can guarantee that their code has no bugs at all.

Human error seems to be inevitable from time to time and that’s why we might want to leverage on code-generation tool to help us write less code. And that will not only reduce the time we have to spend but also increase the code quality. If you were familiar with Java, you probably heard annotation processing before, and if you’re a Kotlin developer, now you have the KSP.

KSP

KSP stands for Kotlin Symbol Processing. As the name suggests, it’s not just a code generation tool but a compiler plugin API that helps you to get complete source code information during compile time. You can do whatever like generate more codes, analyze the structure or even throw an additional error to break the build, and code generation is definitely a very strong use case and we’ll mainly focus on this topic today.

Code-Generation

Generally, we’ll create a separate processor module with some processors that will be triggered when compiled. Each processor can read the main application source code information to generate code and the generated code will compile together with your source code to form the final output.

Using predefined annotations is a common trick. The processor module and your application module both depend on the same annotation module, and the processor module can then filter the information from the application module by these annotations easily, but it’s totally up to your preference.

I’ll mix the detailed step-by-step introduction with a simple use case to showcase what KSP can do for us in the next section.

Example

Let’s say we like to build some factory pattern like the below:

public enum class AnimalType {
CAT,
DOG,
}

public fun AnimalFactory(key: AnimalType): Animal = when (key) {
AnimalType.CAT -> Cat()
AnimalType.DOG -> Dog()
}
interface Animal
class Dog : Animal
class Cat : Animal

When we call AnimalFacoty(AnimalType.CAT), we’d like to get a Cat instance – not to worry about how to create it. This pattern is beautiful but not scalable as you can see whenever we want to add a new type, we need to modify the same function again and again. What can we do to make it better?

The most essential part is to define those animal classes and interfaces and the rest of them seem to just follow some rule behind. Let’s aim to build AnimalType and AnimalFactory automatically via KSP so we don’t have to write boilerplate code anymore.

Update project-level config

First, update your root build.gradle to enable the KSP feature like the below:

plugins {
id 'org.jetbrains.kotlin.jvm' version '1.7.10' apply false
id 'com.google.devtools.ksp' version '1.7.10-1.0.6' apply false
}
buildscript {
dependencies {
classpath 'org.jetbrains.kotlin:kotlin-gradle-plugin:1.7.10'
}
}

KSP version can be found here: https://github.com/google/ksp/releases

Annotation Module

We’ll define annotations first in a new module for both application and processor module to pass meta-data in between.

@Target(AnnotationTarget.CLASS)
@Retention(AnnotationRetention.SOURCE)
annotation class AutoElement
@Target(AnnotationTarget.CLASS)
@Retention(AnnotationRetention.SOURCE)
annotation class AutoFactory

We assume AutoElement will add to the concrete class and AutoFactory is used to declare the return type. Later on, we’ll rely on this contract to filter the needed information to build the graph.

Here you can see annotation class itself also has two eligible annotations. @Target is used to limit where it can be used. Given we only want it to be label on class or interface, we choose CLASS. @Retention is used to declare the lifecycle of the annotation, and for KSP use-case, we don’t need it anymore after compile so we set it to SOURCE.

Processor Module

This is where all the magic happened! We’ll create the processor soon to generate code automatically. Before that, we need to declare the KSP dependency like the below build.gradle:

plugins {
id 'org.jetbrains.kotlin.jvm'
}

dependencies {
implementation project(":annotation")
implementation 'com.google.devtools.ksp:symbol-processing-api:1.7.10-1.0.6'
}

A processor module might have several processors, and here is the base SymbolProcessor interface that all processors should implement.

interface SymbolProcessor {
fun process(resolver: Resolver): List //main logic
fun finish() {}
fun onError() {}
}

The process function is the entry point to parse the code and Resolver is an object that helps us to get code symbols. We can build a helper function like the below:

fun Resolver.getSymbols(cls: KClass<*>) =
this.getSymbolsWithAnnotation(cls.qualifiedName.orEmpty())
.filterIsInstance()
.filter(KSNode::validate)

KSClassDeclaration is represented as a class or interface symbol in source code, and Resolver.getSymbolsWithAnnotation is a function that returns every symbol that has the target annotation.

Let’s simplify the problem and assume we only have one factory need to be generated. We can get all the symbols annotated with our annotation like below:

val factory = resolver.getSymbols(AutoFactory::class).firstOrNull()
val elements = resolver.getSymbols(AutoElement::class).toList()

If you’re interested in full implementation, feel free to check here!

KotlinPoet

Now it’s time to write the graph into a file. We can leverage a great tool developed by Square — KotlinPoet. KotlinPoet also has great integration with KSP. Just simply add the dependency first:

dependencies {
implementation 'com.squareup:kotlinpoet:1.11.0'
implementation 'com.squareup:kotlinpoet-ksp:1.12.0'
}

Then create a function to generate the file like below:

private fun genFile(key: ClassName, list: List): FileSpec {
val packageName = key.packageName
val funcName = key.simpleName + "Factory"
val enumName = key.simpleName + "Type"

return FileSpec.builder(packageName, funcName)
.addType(TypeSpec.enumBuilder(enumName)
.apply {
list.forEach {
addEnumConstant(it.simpleName.uppercase())
}
}
.build())
.addFunction(FunSpec.builder(funcName)
.addParameter("key", ClassName(packageName, enumName))
.returns(key)
.beginControlFlow("return when (key)")
.apply {
list.forEach {
addStatement("${enumName}.${it.simpleName.uppercase()} -> %T()", it)
}
}
.endControlFlow()
.build())
.build()
}

It’s a bit long but quite declarative, you can tell KotlinPoet provides a more systematic way to generate Kotlin code than StringBuilder or other manual works.

After we generate the FileSpec, we can finish the whole flow in process function like below:

override fun process(resolver: Resolver): List {
val factory = ...
if (factory != null) {
val elements = ...
genFile(
factory.toClassName(),
elements.map(KSClassDeclaration::toClassName)
).writeTo(codeGenerator, Dependencies(true))
}
return emptyList()
}

Seems great, but how does the compiler know where the processor is? The processor is created by a provider and here is the definition

fun interface SymbolProcessorProvider {
fun create(env: SymbolProcessorEnvironment): SymbolProcessor
}

So we can create our processor in a provider like the below:

fun create(environment: SymbolProcessorEnvironment) =
BuilderProcessor(
environment.codeGenerator,
environment.logger
)

But is that all? The compiler still has no idea where the provider is. We need to register our provider with the full name including the package name in a file named as resources/META-INF/services/com.google.devtools.ksp.processing.SymbolProcessorProvider for the compiler to link.

Application module

We’re close to what we want, now we can use the processor in the application module. The build.gradle file in the application module will look like the below:

plugins {
id 'com.google.devtools.ksp'
}

sourceSets {
main {
java {
srcDir "${buildDir.absolutePath}/generated/ksp/"
}
}
}
dependencies {
implementation project(":annotation")
ksp project(":processor")
}

First, apply the KSP plugin. And declare the annotation and the processor module as dependencies. Noted that for the processor module we need to use ksp keyword. And since KSP is quite new to IDE for now, if the auto-generated files are created but IDE doesn’t know how to find them manually add it as sourceSets section did.

@AutoFactory
interface Animal
@AutoElement
class Dog : Animal
@AutoElement
class Cat : Animal

Finally, add the annotation to your target class then rebuild the project and wait for KSP to generate the file for you and that’s it!!

Reference

https://kotlinlang.org/docs/ksp-overview.html

https://github.com/Jintin/KFactory


KSP for Code-Generation was originally published in Google Developers Experts on Medium, where people are continuing the conversation by highlighting and responding to this story.

Total
1
Shares
Leave a Reply

Your email address will not be published. Required fields are marked *

Previous Post
what’s-tiktok’s-attribution-manager-&-how-can-advertisers-use-it?

What’s TikTok’s Attribution Manager & How Can Advertisers Use It?

Next Post
awesome-tools-for-web-developer-–-part-2

Awesome tools for Web Developer – Part 2

Related Posts