Чиним сериализацию объектов в Kotlin раз и навсегда

Недавно я наткнулся на статью о проблеме c Java-сериализацией объектов в Kotlin. Автор предложил решать её добавлением метода readResolve к каждому объекту, который наследуется от java.io.Serializable.

Этот способ выглядит абсолютно правильным, однако его поддержка может оказаться слишком проблематичной. С учетом того, что в нашем проекте эта проблема возникала только при использовании объектов внутри Bundle, мы решили использовать проверку через is для каждой ветки when-выражений в случае sealed классов.

Тем не менее, размышляя об этом, я никак не мог понять, почему Kotlin не генерирует readResolve в компиляторе, поддерживая singleton-свойства объектов. Мне казалось, что это работа для инструментов, а не для человека. Но раз Kotlin не добавляет эту функцию сам, мы можем ему помочь! Этим мы сейчас и займёмся.

Взгляд ближе

Для начала внимательно посмотрим на метод, который нам нужно сгенерировать:

object Example : java.io.Serializable {
   // TODO: should be generated
   fun readResolve(): Any? = Example
}

Плагин должен добавить метод readResolve для каждого объекта, который наследуется от java.io.Serializable. Данная функция не имеет параметров и возвращает текущее значение объекта, замаскированное под типом Any?.

Этот метод должен существовать только в получившихся .class-файлах и желательно быть незаметным в IDE. Это значительно облегчает нам задачу, позволяя реализовать генерацию только на бэкенде компилятора.

Настраиваем среду

Начнём создание плагина с настройки сборки. Также мы убедимся в том, что плагин успешно подключается к компилятору через отдельный интеграционный модуль.

Плагин зависит от артефакта компилятора, который нужен только во время сборки плагина; в рантайме Kotlin содержит все необходимые классы по умолчанию.

К счастью, JetBrains публикует специальную версию компилятора для плагинов под идентификатором “kotlin-compiler-embeddable”:

// kotlin-plugin/build.gradle
apply plugin: "org.jetbrains.kotlin.jvm"
dependencies {
   implementation "org.jetbrains.kotlin:kotlin-stdlib"
   compileOnly "org.jetbrains.kotlin:kotlin-compiler-embeddable"
}

Входной точкой в плагин служит ComponentRegistrar, который вызывается перед компиляцией и позволяет зарегистрировать все расширения внутри компилятора:

class ObjectSerializationComponentRegistrar: ComponentRegistrar {
   override fun registerProjectComponents(
       project: MockProject, 
       configuration: CompilerConfiguration
   ) {
       println("Works")
   }
}

Kotlin использует ServiceLoader, чтобы подключить наш ComponentRegistrar. По этой причине плагин должен содержать файл с полным именем класса в папке META-INF/services. Альтернативой является использование AutoService от Google, который создаёт такие файлы за вас.

# resources/META-INF/services/org.jetbrains.kotlin.compiler.plugin.ComponentRegistrar
me.shika.ObjectSerializationComponentRegistrar

Создав минимальный плагин, переходим к интеграционному модулю:

// integration-test/build.gradle
apply plugin: "org.jetbrains.kotlin.jvm"
dependencies {
   kotlinCompilerPluginClasspath project(':kotlin-plugin')
}

Kotlin имеет отдельную конфигурацию, которая отвечает за подключение плагина и всех его зависимостей. Если мы попробуем скомпилировать какой-либо класс в текущем модуле, мы должны увидеть строчку “Works” в консоли.

Теперь, когда минимальный плагин настроен, мы можем смотреть в сторону кодогенерации. На текущий момент Kotlin поддерживает три разные платформы, из которых мы заинтересованы только в JVM (потому что java.io.Serializable существует только там). Для нее мы будем использовать ExpressionCodegenExtension.

Компилятор применяет это расширение на каждый класс на этапе генерации байт-кода. Здесь мы можем манипулировать вызовом функций, обращением к полям, а также добавлять синтетические части к классам. Последнее — как раз то, что нам нужно, чтобы добавить readResolve:

class ObjectSerializationJvmGeneration : ExpressionCodegenExtension {
   override fun generateClassSyntheticParts(codegen: ImplementationBodyCodegen) {
       println("Found ${codegen.descriptor}")
       // todo: generate
   }
}

На этом этапе мы просто выведем текстовую репрезентацию класса, для которого была вызвана генерация.

Большинство возможных расширений заданы как подкласс ProjectExtensionDescriptor<T>. Они имеют функцию registerExtension для добавления кастомной функциональности. С целью генерации байт-кода мы будем использовать только ExpressionCodegenExtension, но компилятор даёт нам намного больше возможностей для расширения.

Последний этап — подключение расширения в ComponentRegistrar:

override fun registerProjectComponents(
   project: MockProject, 
   configuration: CompilerConfiguration
) {
   ExpressionCodegenExtension.registerExtension(
       project,
       ObjectSerializationJvmGeneration()
   )
}

Теперь мы можем вызвать компиляцию модуля integration-test и увидеть, что выводится в консоль.

Генерируем байт-код

Компилятор предоставляет информацию о классе в виде дескриптора, который содержит данные о его функциях, полях и родителях. Этого достаточно, чтобы понять, нужно ли чинить сериализацию для обрабатываемого класса.

fun ClassDescriptor.needsSerializableFix() =
   DescriptorUtils.isObject(this)
       && isSerializable()
       && !hasReadMethod()

Проверка выше состоит из трёх шагов:

1. Имеем ли мы дело с object-классом?
2. Наследуется ли класс от java.io.Serializable?
3. Есть ли у класса созданный ранее метод readResolve?

Первый шаг компилятор делает за нас. В DescriptorUtils уже содержится нужная нам функция:

fun ClassDescriptor.isSerializable(): Boolean =
   getSuperInterfaces().any {
       it.fqNameSafe == SERIALIZABLE_FQ_NAME
       || it.isSerializable()
   } || getSuperClassNotAny()?.isSerializable() == true
val SERIALIZABLE_FQ_NAME = FqName("java.io.Serializable")

На втором этапе проверки нам придётся пройти по всему дереву родителей и найти интерфейс Serializable.

Последний шаг — найти readResolve среди функций класса:

fun ClassDescriptor.hasReadMethod() =
    unsubstitutedMemberScope
        .getContributedFunctions(
            SERIALIZABLE_READ, 
            NoLookupLocation.FROM_BACKEND
        )
        .any { 
            it.name == SERIALIZABLE_READ 
            && it.valueParameters.isEmpty()
        }
val SERIALIZABLE_READ = Name.identifier("readResolve")

У дескриптора есть доступ к каждой функции, находящейся в скоупе класса. Мы находим вариант с нужным нам именем и нулевым количеством параметров.

Теперь, когда мы знаем, какие классы нам нужно модифицировать, мы можем приступить к генерации самого метода. Компилятор Kotlin использует ASM для манипуляций с байт-кодом и передаёт уже инициализированный инстанс ClassBuilder в наше расширение:

private fun ImplementationBodyCodegen.addReadResolveFunction(
   block: InstructionAdapter.() -> Unit
) {
   val visitor = v.newMethod(
       NO_ORIGIN,
       ACC_PUBLIC or ACC_SYNTHETIC,
       SERIALIZABLE_READ.identifier,
       "()Ljava/lang/Object;",
       null,
       EMPTY_STRING_ARRAY
   )

   visitor.visitCode()
   val iv = InstructionAdapter(visitor)
   iv.apply(block)
   FunctionCodegen.endVisit(iv, "JVM serialization bindings")
}

Мы создаём новый метод с модификаторами public и synthetic, так что он не будет виден в IDE. Строка ()Ljava/lang/Object; передаёт параметры и возвращаемый тип. Помимо этого, мы генерируем тело функции, которое передаётся через лямбда-параметр.

Самый простой способ узнать байт-код инструкции для метода — посмотреть на объект Example из примера выше:

GETSTATIC Example.INSTANCE : LExample;
ARETURN

InstructionAdapter, который используется для генерации тела функции, имеет синтаксис, очень близкий к инструкциям байт-кода, которые он создаёт. Используя приведённый выше сниппет, мы наконец можем закончить создание метода:

if (codegen.descriptor.needsSerializableFix()) {
   val selfType = codegen.typeMapper.mapType(codegen.descriptor)

   codegen.addReadResolveFunction {
       getstatic(codegen.className, "INSTANCE", selfType.descriptor)
       areturn(selfType)
   }
}

Тестируем

Команда компилятора Kotlin тестирует плагины на многих уровнях, включая использование интеграционных и юнит-тестов. Некоторые тесты (например, с валидацией байт-кода) немного сложны в настройке, так что их мы касаться не будем.

Я предлагаю остановиться на более высокоуровневых тестах: мы протестируем получившиеся классы на валидность, а потом проведём интеграционный тест в уже существующем у нас модуле.
Для тестирования вывода компилятора я использую kotlin-compile-testing. Эта прекрасная библиотека позволяет получить доступ к сгенерированным файлам через Java-рефлексию. На вход она принимает как директории файлов (например, через test/resources/), так и простые сниппеты.

private val SERIALIZABLE_OBJECT = """
   import java.io.Serializable

   object Serial : Serializable
""".source()

@Test
fun `adds readResolve to obj extending Serializable`() {
   compiler.sources = listOf(SERIALIZABLE_OBJECT)
   val result = compiler.compile()

   val klass = result.classLoader.loadClass("Serial")
   assertTrue(klass.methods.any { it.addedReadResolve()})
}
private fun Method.addedReadResolve() =
   name == "readResolve"
       && parameterCount == 0
       && returnType == Object::class.java
       && isSynthetic

Приведённый тест компилирует класс из строки и проверяет наличие readResolve с помощью рефлексии.

С интеграционными тестами всё намного проще. Мы уже создали модуль с подключённым плагином. Единственное, что осталось сделать, — добавить ваш любимый тестовый фреймворк и проверить инстанс объекта после сериализации:

private object TestObject : Serializable

@Test
fun `object instance is the same after deserialization`() {
   assertEquals(TestObject, serializeDeserialize(TestObject))
}

private fun serializeDeserialize(instance: Serializable): Serializable {
   val out = ByteArrayOutputStream()
   ObjectOutputStream(out).use {
       it.writeObject(instance)
   }
   return ObjectInputStream(
       ByteArrayInputStream(out.toByteArray())
   ).use {
       it.readObject() as TestObject
   }
}

Заключение

Расширения для компилятора Kotlin — удобный инструмент для генерации кода и метапрограммирования. Я открыл для себя огромное количество возможностей в этой платформе и, несмотря на высокий порог вхождения, предлагаю вам попробовать самим.

Конечно же, разработка и поддержка такого плагина имеет подводные камни, которых я не коснулся в этой статье: например, постоянно ломающийся API или отсутствие какой-либо документации. Надеюсь, что ситуация изменится в сторону официальной поддержки плагинов после выхода Kotlin версии 1.4.

Репозиторий с этим плагином доступен на GitHub. Также артефакт доступен через Maven (если вы захотите попробовать использовать его в своих проектах).