kotlin

Adding KSP to your toolbelt

Vaios Tsitsonis

• 7 min read
Adding KSP to your toolbelt
Kotlin Symbol Processing can be a valuable tool in your toolbelt

Let's prepare the ground

Before we start examining how to use KSP a.k.a Kotlin Symbol Processing, let's say a few words about what it is. For the familiars with the annotation processing, we could say that is an annotation processing tool for Kotlin. This means that it can be used for code written in Kotlin JVM, Kotlin JS, and Kotlin Native.

So why KSP?

Everyone knows that there is interoperability between Java and Kotlin. So why do we need a new tool for annotation processing as we can already implement the same functionality with the existing Java APIs? The pretty obvious advantage is that it can unlock the full power of Kotlin. For a Java annotation processor is impossible to handle all the Kotlin's features. With KSP we can handle things like sealed class, data class, etc.

Apart from the obvious advantage, there are more. Kotlin exceeds the JVM world (with Kotlin JS and Kotlin Native). In contrast to Java annotation processors, KSP can be used to write tools that are not used only on JVM projects.

Last but not least is the decreased build time. During Java annotation processing there is a phase in which Kotlin classes are converted into Java. With KSP we avoid doing these unnecessary conversions. As a result, we have smaller build times (according to the documentation, switching to KSP would immediately reduce the time spent in the compiler by 25%).

Our use case

So let's give an example of how we can use KSP to build tools. In our codebase, we noticed that we do a lot of conversions from a string value to an enum. Basically, in every enum class we had to duplicate the following code:

enum class MyEnum(val value: String) {

    FIRST_ENTRY("first entry"),

    SECOND_ENTRY("second_entry");

    // This code was duplicated in almost every enum :S
    companion object {

        fun from(value: String?): MyEnum? = values().find {
            it.value.equals(value, ignoreCase = true)
        }
    }
}

Our plan is to remove all this boilerplate with the help of KSP.

Set up KSP

Before we start writing our KSP logic we need to add some dependencies to our project. In the root Gradle file, we need to add a dependency to the KSP Gradle plugin.

buildscript {

    repositories {
        mavenCentral()
    }

    dependencies {
        classpath(kotlin("gradle-plugin", version = "1.6.10"))
    }
}

Then we need to create two new Gradle modules. The first one will contain all the annotations and the second will contain the processor. This step is not necessary but is considered a best practice. What we achieve with that is to not include KSP logic in our code base. We can only include the annotations. The Gradle file for the processor module should look like that:

plugins {
  kotlin("jvm")
}

dependencies {
    implementation(kotlin("stdlib-jdk8"))
    implementation(project(":enum-convertible-annotation"))
    implementation("com.google.devtools.ksp:symbol-processing-api:1.6.10-1.0.6")
}

Define the annotation

After adding all the dependencies, we need to start writing some code. Firstly, we define an annotation that is going to indicate that our processor will generate a mapper class for the annotated enum. This annotation is needed only during compilation so we defined Retention as SOURCE. Additionally, we want to use this annotation only in enums so we defined the Target as CLASS (unfortunately, there is no dedicated target for enums 😞). The annotation will look like that: 

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

KSP

Visitor

To start providing KSP functionality we need to implement some interfaces. The first interface that we are going to implement is the KSVisitor. This implementation is needed in order to extract the desired data from the code. In KSP, every part of our code (fields, getters, setters, classes, enums, etc) is modeled as nodes. A KSVisitor accepts these nodes, analyzes them, and extracts information. Default implementations are provided and probably the simplest one is the KSVisitorVoid. We are going to extend this class to extract the needed information (package, enum name, and field name) from our enum. The implementation will look like that:

data class ProcessedEnumConvertible(
    val packageName: String,
    val className: String,
    val keyName: String,
    val classDeclaration: KSClassDeclaration
)

class EnumConvertibleVisitor(
    private val processedEnumConvertibles: MutableList<ProcessedEnumConvertible>
) : KSVisitorVoid() {

    override fun visitClassDeclaration(classDeclaration: KSClassDeclaration, data: Unit) {
        if (classDeclaration.classKind != ClassKind.ENUM_CLASS) {
            throw IllegalStateException(
                "${EnumConvertible::class.simpleName} can be used only in an enum class."
            )
        }

        val processedEnumConvertible = ProcessedEnumConvertible(
            packageName = classDeclaration.packageName.asString(),
            className = classDeclaration.simpleName.asString(),
            keyName =  classDeclaration.extractKeyName(),
            classDeclaration = classDeclaration
        )
        processedEnumConvertibles.add(processedEnumConvertible)
    }

    private fun KSClassDeclaration.extractKeyName(): String = 
        this.primaryConstructor!!.parameters.first().type.resolve().declaration.simpleName.asString()
}

Generator

This class is just a matter of personal taste in order to better organize your KSP code. It is responsible for generating files with the expected source code based on the extracted information. The file generation happens with the help of a CodeGenerator which is provided to us through the KSP environment (see SymbolProcessorProvider section). The source code is just a created String as we would write it by ourselves in our IDE. The most important thing here is the Dependencies param on the Code Generator.create New File function, where we can declare on which files the generated file depends. This means that the generated file will be regenerated only if any of its dependency files have been updated which saves us build time.

class EnumConvertibleExtGenerator(
    private val codeGenerator: CodeGenerator
) {

    fun generate(processedEnumConvertible: ProcessedEnumConvertible) {
        // Generate file and declare its dependencies
        val mapperFile: OutputStream = codeGenerator.createNewFile(
            dependencies = Dependencies(
                aggregating = true,
                sources = arrayOf(classDeclaration.containingFile!!)
            ),
            packageName = packageName,
            fileName = "${className}Mapper",
            extensionName = "kt"
        )

        // Generate source code based on data and a template string
        val sourceCode = buildConvertibleMapperSource(
            packageName = processedEnumConvertible.packageName,
            className = processedEnumConvertible.className,
            keyName = processedEnumConvertible.keyName
        ).toByteArray()

        // Write source code to file
        mapperFile.write(sourceCode)
        mapperFile.close()
    }

    private fun buildConvertibleMapperSource(
        packageName: String,
        className: String,
        keyName: String
    ): String = """
            |package $packageName
            |
            |object ${className}Mapper {
            |    
            |    fun from(${keyName}: String?): ${className}? = values().find {
            |        it.value.equals(${keyName}, ignoreCase = true)
            |    }
            |}
        """.trimMargin()
}

SymbolProcessor

Another interface that we need to implement, is the SymbolProcessor. This class will be responsible for orchestrating the implemented KSVisitor and Generator to start generating the desired files. Specifically, we need to implement two functions. The first one is the process function. In this function, we process the source code to extract all the needed information. This is the place where we need to use the implemented KSVisitor. The other function is the finish. This will be called when the process is finished. So, it is the best place to generate our code based on the extracted data.

class EnumConvertibleProcessor(codeGenerator: CodeGenerator) : SymbolProcessor {

    private val enumConvertibleExtGenerator =
        EnumConvertibleExtGenerator(codeGenerator = codeGenerator)

    private val processedEnumConvertibles = mutableListOf<ProcessedEnumConvertible>()

    private val enumConvertibleVisitor = EnumConvertibleVisitor(
        processedEnumConvertibles = processedEnumConvertibles
    )

    override fun process(resolver: Resolver): List<KSAnnotated> {
        val symbols = resolver.getSymbolsWithAnnotation(
            annotationName = EnumConvertible::class.java.name
        )

        symbols.filterIsInstance<KSClassDeclaration>()
            .filter { kSClassDeclaration -> kSClassDeclaration.validate() }
            .forEach { kSClassDeclaration ->
                kSClassDeclaration.accept(enumConvertibleVisitor, Unit)
            }

        return emptyList()
    }

    override fun finish(): Unit = processedEnumConvertibles.forEach {
        enumConvertibleExtGenerator.generate(processedEnumConvertible = it)
    }
}

SymbolProcessorProvider

At this point, we have implemented our processor but there are some major things missing. Firstly, we have not yet provided the CodeGenerator. This is possible with the help of SymbolProcessorProvider which is responsible for creating our SymbolProcessor. The SymbolProcessorProvider has access to the environment in which our KSP is going to run. This environment can give us access to some tools (such as CodeGenerator and Logger) as well as some options that have passed through the terminal or the KSP Gradle plugin (which is useful in order to provide different configurations to our processor).

class EnumConvertibleProcessorProvider : SymbolProcessorProvider {

    override fun create(environment: SymbolProcessorEnvironment): SymbolProcessor =
        EnumConvertibleProcessor(codeGenerator = environment.codeGenerator)
}

This was the last class that we needed to implement. The last step that we need to execute is to define how this service is going to be loaded during compilation. To achieve that, we have to create the src/main/resources/META-INF/services/com.google.devtools.ksp.processing.SymbolProcessorProvider file (this should happen in the processor's Gradle module).

After creating the com.google.devtools.ksp.processing.SymbolProcessorProvider file, we need to write the fully-qualified name of the created SymbolProcessorProvider in it. The content of the file should look like that:

Configure it in app module

After finishing our KSP processor, we want to see it in action. To do so, we have to add a dependency to it from a consumer Gradle module.

plugins {
    id("com.google.devtools.ksp")
}

kotlin {
    sourceSets.main {
        kotlin.srcDir("build/generated/ksp/main/kotlin")
    }
    sourceSets.test {
        kotlin.srcDir("build/generated/ksp/test/kotlin")
    }
}

dependencies {
    implementation(kotlin("stdlib-jdk8"))
    implementation(project(":enum-convertible-annotation"))
    ksp(project(":enum-convertible-processor"))
}

Also, we need to mark the following directories as Generated Sources to make our IDE aware of the generated files (in another case, it would be difficult to add imports).

build/generated/ksp/main/kotlin/
build/generated/ksp/main/java/

Now, if we build our project we should see a mapper class for every enum that was annotated with the @EnumConvertible annotation. The mapper classes should look like that:

object MyEnumMapper {
    
    fun from(value: String?): MyEnum = values().find {
        it.value.equals(value, ignoreCase = true)
    }
}

Final thoughts

In this blog post, we showcased how to use KSP to built tools and reduce boilerplate code. Some last tips that we need to keep in mind before implementing real case processors are:

  • KSP does not modify existing code. It generates new based on the existing code. If you need to modify the existing code you should create a Kotlin Compiler Plugin.
  • Calling resolve() is a costly operation, so use it wisely. (This function was used in our KSVisitor implementation and give access to more information for the corresponding element).
  • In our snippets, we represented the templates of the generated code as strings for simplicity. You can use KotlinPoet to generate your source files. KotlinPoet uses the builder pattern to specify the logic in the generated files. By doing so, you can eliminate compilation errors that occur due to the fact that you write your code in a String form and enhance your productivity.

A library based on the described example can be found here.