Room & Kotlin Symbol Processing

Photo by Marc Reichelt on Unsplash

Room is Jetpack’s database abstraction library that provides ability to write compile time verified SQL queries without any of the boilerplate. It achieves this by processing annotations in the code and generating Java sources.

Annotation processors are very powerful but they incur a build time penalty. This is usually acceptable for code written in Java but for Kotlin, this penalty is quite significant because Kotlin does not have a built-in annotation processing pipeline. Instead, it supports annotation processors by generating stub Java code from the Kotlin code, and then piping it into the Java compiler for processing.

Since not everything in Kotlin sources are representable in Java, some information gets lost in this translation. Similarly, Kotlin is a multi platform language, but KAPT only works if you are targeting Java bytecode.

Meet: Kotlin Symbol Processing

As annotation processors are widely used on Android, KAPT became a bottleneck for build performance. To address this issue, Google Kotlin Compilers Team started working on an alternative that provides first class annotation processing support for Kotlin. When the project was born, we got very excited for Room as it will allow Room to have a much better support for Kotlin. As of Room 2.4, it has experimental support for KSP and we observed up to 2x compilation speed increase, especially for clean compilation.

This post is not about teaching annotation processing, Room, or KSP. Rather, it is about the challenges we’ve faced and trade-offs we’ve made while adding KSP support to Room. You don’t need to know Room or KSP to understand it but familiarity with annotation processing is necessary.

Note that, we started using KSP long before it became stable. As such, some decisions we made may or may not be applicable today.

The goal of this post is to give a head start to annotation processor authors on what to look out for when adding KSP support to their projects.

A Brief on How Room Works

Room’s annotation processing involves two steps. There are “Processor” classes that traverse the user code, validate and extract necessarily information into “value objects”. These value objects are passed into “Writer” classes that turn them into code. As many other annotation processors, Room heavily relied on Auto-Common and frequently referenced classes from the javax.lang.model package (Java annotation processing API packages).

To support KSP, we had three options:

a) Duplicate each “Processor” class for JavaAP and KSP where they’ll have the same value objects as outputs that we can feed into writers.

b) Create an abstraction over KSP / JavaAP such that processors can have one implementation that uses the abstraction.

c) Replace JavaAP with KSP and require developers to use KSP to process Java code as well.

Option C was not really feasible as that would create a significant disruption to Java users. With the adoption numbers of Room, such disruptive changes are not possible. Between “A” and “B”, we’ve decided to go with “B” because processors have significant business logic and tearing it apart was not going to be easy.

Meet: X-Processing

Creating a general purpose abstraction over JavaAP and KSP is not very easy. Kotlin and Java can interoperate, but they have different models. For instance, there are special class types like value classes in Kotlin or static methods in Java. Moreover, Java has fields and methods in classes whereas Kotlin has properties and functions.

Instead of trying to shoot for the perfect abstraction, we decided to implement “what Room needs”. This literally meant, find every single file in Room where javax.lang.model is imported and move it into an abstraction in X-Processing. So TypeElement became XTypeElement , ExecutableElement became XExecutableElement etc.

Unfortunately, javax.lang.modelAPIs were used all over Room. Creating these X classes at one change would create unrecoverable mental load on reviewers. As a result, we needed a way to implement this iteratively.

On the other hand, we needed to prove that this is feasible. So we first prototyped it and once we were confident that it is a reasonable choice, we then re-implemented all X classes one by one with their own tests.

A good example of what I meant by implementing “what Room needs” can be seen in this change about fields in a class. When Room processes a class for its fields, it is always interested in all of its fields, including fields in the super classes. So when we created the corresponding X-Processing API, we only added the ability to get all fields.

If we were designing a general purpose library, this would never pass API review. But because our target was just Room, and it already had a helper method with the same functionality for TypeElement, copying it as is de-risked the project.

Once we had the basic X-Processing APIs with their own test, the next step was moving Room to use that abstraction. It is also where “implementing what Room needs” paid off nicely. Room already had extension functions/properties on javax.lang.model APIs for common functionality (e.g. getting methods of a TypeElement). We first updated these extensions to looks like the X-Processing APIs and then migrated Room to X-Processing in 1 CL.

API Usability Improvements

Keeping the JavaAP like API didn’t mean we couldn’t improve things. After migrating Room to X-Processing, we followed up with a bunch of API improvements.

For instance, Room had many calls to MoreElement/MoreTypes to convert between javax.lang.model types (e.g. MoreElements.asType). A call to that would usually look like:

We moved all of these calls into Kotlin contracts such that it would be sufficient to write:

Another good example is finding methods in a TypeElement . Usually in JavaAP, you need to call ElementFilter class to get methods in a TypeElement . Instead, we made it a property in XTypeElement .

One last example, maybe one of my favorites, is assignability. In JavaAP, if you want to check if a given TypeMirror is assignable from another TypeMirror, you need to call Types.isAssignable .

This code is really hard to read as you can’t even guess whether it validates type1 can be assigned from type2 or the other way around. We already had an extension function that looked like:

And in X-Processing , we were able to turn it into a regular function on XType that simply looks like:

Implementing KSP Backend for X-Processing

Each of these X-Processing interfaces came with their own test suites. We didn’t write them to test AutoCommon or JavaAP. Instead, they were written so that once we have the KSP implementation of them, we could run the test suite to validate that it conforms to what Room expects.

As the initial X-Processing APIs were modeled after javax.lang.model, they didn’t always fit KSP so we also improved the API to have better Kotlin support when needed.

This also presented a new problem. Existing Room codebase was written to process Java source code. When the application is written in Kotlin, Room only understood what that Kotlin looked like in Java stubs. We’ve decided to keep it similar in the KSP implementations of X-Processing.

For instance, suspend functions in Kotlin have the following signature when compiled:

To keep the same behavior, XMethodElement implementation in KSP synthesizes a new argument for suspend methods, along with the new return type. (KspMethodElement.kt)

Note that this works well because Room generates Java code, even in KSP. When we add support for Kotlin code generation, this will probably change.

Another example of this is properties. A Kotlin property might also have a synthetic getter/setter (accessor) based on its signature. XTypeElement implementation synthesizes methods for them because Room expects to find them as methods (KspTypeElement.kt).

Note: We actually have plans to change the XTypeElement API to rather provide properties instead of fields as that is what Room really wants to know. As you’ve probably guessed by now, we decided not to do that “yet” to reduce the changes for Room. Hopefully, we’ll get to it one day and when we do that, JavaAP implementation of XTypeElement will change to bundle methods with fields as properties.

One last interesting problem while adding KSP implementation for X-Processing was the API coupling. These processor APIs frequently access each-other, such that you cannot implement XTypeElement in KSP without implementing XField / XMethod which themselves reference XType etc. While adding these KSP implementations, we wrote separate tests for them for the implemented parts. When the KSP implementation became more complete, we gradually enabled all X-Processing tests with the KSP backend.

Note that, we only run tests in the X-Processing project at this phase so even though we knew what we were testing was good, we didn’t know if all of Room’s tests would pass (call it unit vs integration testing :) ). We needed a way to run all of Room’s tests with the KSP backend. And this is when “X-Processing-Testing” artifact was born.

Meet: X-Processing-Testing

Writing an annotation processor is 20% processor code and 80% test code. You need to account for variety of possible developer errors and make sure you report proper error messages. To write these tests, Room already had a helper method that looked like this:

Under the hood, runTestused the Google Compile Testing library and allowed us to simply unit test Processors. It synthesized a Java annotation processor and invoked the given process method inside it.

Unfortunately, Google Compile Testing only supports Java source code. To test Kotlin, we needed another library and luckily there was Kotlin Compile Testing. It did allow us to write tests targeting Kotlin and we’ve contributed KSP support to the library.

Note: We later replaced Kotlin Compile Testing with an internal implementation to ease Kotlin/KSP updates in the AndroidX Repo. We also added better assertion APIs that would require breaking API changes for KCT.

As the final step to be able to run all tests with KSP, we’ve created the following testing API:

The main difference between this one and the original is that, it runs the test with both KSP and JavaAP (or KAPT, depending on sources). Because it is running the test multiple times, it cannot return a single result as the assertions might be different between KSP and JavaAP.

Instead, we came up with:

After each compilation, it invokes the result assertion (and if there are no failure assertions, check the compilation succeeded). We refactored each Room test to look like:

The rest was straightforward. Migrate each Room compilation test to the new API, discover new KSP / X-Processing bugs, report, implement workarounds; rinse and repeat. As KSP was under heavy development, we did hit fair number of bugs. Each time, we reported the bug, linked to it from Room source and moved on (or contribute a fix). After each KSP release, we searched the codebase for fixed issues and removed the workarounds / enabled tests.

Once we got to a good coverage in compilation tests, the next step was also running Room’s integration tests with KSP. These are actual Android test apps that also test the behavior at runtime. Luckily, Android has Gradle variants support so it was fairly easy to run our Kotlin integration tests with both KSP and KAPT.

What is Next?

Adding KSP support to Room was only the first step. Now, we need to update Room to take advantage of it. For instance, all type checks in Room ignored nullability because javax.lang.model's TypeMirror does not understand nullability. As such, Room could sometimes trigger NullPointerException at runtime when calling your Kotlin code. With KSP, these checks now work which created new KSP bugs in Room (e.g. b/193437407). We’ve added some workarounds but ideally we want to improve Room to properly handle these cases.

Similarly, even though we support KSP, Room still only generates Java code. This limitation prevents us from adding support for certain Kotlin features like value classes. Hopefully in the future, we’ll add support for generating Kotlin code as well to provide first class support for Kotlin in Room. And then, maybe more :).

Can I use X-Processing in my project?

Well, not really; at least not the way you can use any other Jetpack library. As mentioned before, we’ve only implemented what Room needs. Writing a true Jetpack library has a significant overhead, like documentation, API stability, Codelabs etc and we are not at a place to take that work. That being said, both Dagger and Airbnb (Paris, DeeplinkDispatch) started using X-Processing to support KSP (and contributed what they needed 🙏). Maybe one day we’ll unbundle it from Room. Technically, you can still use it as artifacts are in Google’s Maven repository, but there is no API guarantees such that you should definitely shade it.

tl;dr;

We’ve added KSP support to Room and even though it wasn’t easy, it was definitely worth it. If you maintain an annotation processor, please add support for KSP to provide a better Kotlin developer experience.

Special thanks to Zac Sweers and Eli Hart for reviewing early versions of this post and also being awesome KSP contributors.

Resources:

• Issue Tracker for Room’s KSP support (where you can find links to commits)
• X-Processing and X-Processing-Testing source code
• KSP source code