How can I access the body of methods during Kotlin annotation processing?

Question

Overview

I'm wondering if there's a way I can apply annotations to functions and access the body of those functions during annotation processing. If it's not possible to directly get the method body through inspection of Element objects within the annotation processor, are there any other alternatives to accessing the source code of the function that these annotations are applied to?

Details

As part of a project I'm working on, I'm trying to use kapt to inspect Kotlin functions annotated with a specific type of annotation and generate classes based on them. For example, given an annotated function like this:

@ElementaryNode
fun addTwoNumbers(x: Int, y: Int) = x + y

My annotation processor currently generates this:

class AddTwoNumbers : Node {
    val x: InputPort<Int> = TODO("implement node port property")

    val y: InputPort<Int> = TODO("implement node port property")

    val output: OutputPort<Int> = TODO("implement node port property")
}

However, I need to include the original function itself in this class, essentially just as if it were copy/pasted in as a private function:

class AddTwoNumbers : Node {
    val x: InputPort<Int> = TODO("implement node port property")

    val y: InputPort<Int> = TODO("implement node port property")

    val output: OutputPort<Int> = TODO("implement node port property")

    private fun body(x: Int, y: Int) = x + y
}

What I've tried

Based on this answer, I tried using com.sun.source.util.Trees to access the method body of the ExecutableElement corresponding to the annotated functions:

override fun inspectElement(element: Element) {
    if (element !is ExecutableElement) {
        processingEnv.messager.printMessage(
            Diagnostic.Kind.ERROR,
            "Cannot generate elementary node from non-executable element"
        )

        return
    }

    val docComment = processingEnv.elementUtils.getDocComment(element)
    val trees = Trees.instance(processingEnv)
    val body = trees.getTree(element).body
    processingEnv.messager.printMessage(Diagnostic.Kind.WARNING, "Processing ${element.simpleName}: $body")
}

However, kapt only generates stubs of method bodies, so all I got for each method body was stuff like this:

$ gradle clean build
...
> Task :kaptGenerateStubsKotlin
w: warning: Processing addTwoNumbers: {
      return 0;
  }
w: warning: Processing subtractTwoNumbers: {
      return 0.0;
  }
w: warning: Processing transform: {
      return null;
  }
w: warning: Processing minAndMax: {
      return null;
  }
w: warning: Processing dummy: {
  }

Update

Accessing Element.enclosingElement on the ExecutableElement representing each function gives me the qualified name of the package/module where the function is defined. For example, addTwoNumbers is declared as a top-level function in Main.kt, and during annotation processing I get this output: Processing addTwoNumbers: com.mycompany.testmaster.playground.MainKt.

Is there a way I can access the original source file (Main.kt) given this information?

Answer 1

It wasn't easy, but I eventually managed to figure out one (rather hacky) solution for this.

I found that during annotation processing, Kotlin was generating metadata files under a temporary build output directory. These metadata files contained serialized information that included the paths to the original source files containing the annotations I was processing:

Looking through the source code for the Kapt plugin, I found this file that allowed me to figure out how to deserialize the information in these files, allowing me to extract the locations of the original source code.

I created a Kotlin object SourceCodeLocator that put this all together so that I could I could pass it an Element representing a function, and it would return me a string representation of the source code containing it:

package com.mycompany.testmaster.nodegen.parsers

import com.mycompany.testmaster.nodegen.KAPT_KOTLIN_GENERATED_OPTION_NAME
import com.mycompany.testmaster.nodegen.KAPT_METADATA_EXTENSION
import java.io.ByteArrayInputStream
import java.io.File
import java.io.ObjectInputStream
import javax.annotation.processing.ProcessingEnvironment
import javax.lang.model.element.Element
import javax.lang.model.element.ElementKind
import javax.lang.model.element.ExecutableElement

internal object SourceCodeLocator {
    fun sourceOf(function: Element, environment: ProcessingEnvironment): String {
        if (function !is ExecutableElement)
            error("Cannot extract source code from non-executable element")

        return getSourceCodeContainingFunction(function, environment)
    }

    private fun getSourceCodeContainingFunction(function: Element, environment: ProcessingEnvironment): String {
        val metadataFile = getMetadataForFunction(function, environment)
        val path = deserializeMetadata(metadataFile.readBytes()).entries
            .single { it.key.contains(function.simpleName) }
            .value

        val sourceFile = File(path)
        assert(sourceFile.isFile) { "Source file does not exist at stated position within metadata" }

        return sourceFile.readText()
    }

    private fun getMetadataForFunction(element: Element, environment: ProcessingEnvironment): File {
        val enclosingClass = element.enclosingElement
        assert(enclosingClass.kind == ElementKind.CLASS)

        val stubDirectory = locateStubDirectory(environment)
        val metadataPath = enclosingClass.toString().replace(".", "/")
        val metadataFile = File("$stubDirectory/$metadataPath.$KAPT_METADATA_EXTENSION")

        if (!metadataFile.isFile) error("Cannot locate kapt metadata for function")
        return metadataFile
    }

    private fun deserializeMetadata(data: ByteArray): Map<String, String> {
        val metadata = mutableMapOf<String, String>()

        val ois = ObjectInputStream(ByteArrayInputStream(data))
        ois.readInt() // Discard version information

        val lineInfoCount = ois.readInt()
        repeat(lineInfoCount) {
            val fqName = ois.readUTF()
            val path = ois.readUTF()
            val isRelative = ois.readBoolean()
            ois.readInt() // Discard position information

            assert(!isRelative)

            metadata[fqName] = path
        }

        return metadata
    }

    private fun locateStubDirectory(environment: ProcessingEnvironment): File {
        val kaptKotlinGeneratedDir = environment.options[KAPT_KOTLIN_GENERATED_OPTION_NAME]
        val buildDirectory = File(kaptKotlinGeneratedDir).ancestors.firstOrNull { it.name == "build" }
        val stubDirectory = buildDirectory?.let { File("${buildDirectory.path}/tmp/kapt3/stubs/main") }

        if (stubDirectory == null || !stubDirectory.isDirectory)
            error("Could not locate kapt stub directory")

        return stubDirectory
    }

    // TODO: convert into generator for Kotlin 1.3
    private val File.ancestors: Iterable<File>
        get() {
            val ancestors = mutableListOf<File>()
            var currentAncestor: File? = this

            while (currentAncestor != null) {
                ancestors.add(currentAncestor)
                currentAncestor = currentAncestor.parentFile
            }

            return ancestors
        }
}

Caveats

This solution seems to work for me but I can't guarantee it will work in the general case. In particular, I'm configured Kapt in my project via the Kapt Gradle plugin (version 1.3.0-rc-198 currently), which determines the directories where all generated files (including the metadata files) are stored. I then make the assumption that the metadata files are stored at /tmp/kapt3/stubs/main under the project build output folder.

I've created a feature request in JetBrain's issue tracker to make this process easier and more reliable so these sorts of hacks aren't necessary.

Example

In my case, I've been able to use this to transform source code like this:

minAndMax.kt

package com.mycompany.testmaster.playground.nodes

import com.mycompany.testmaster.nodegen.annotations.ElementaryNode

@ElementaryNode
private fun <T: Comparable<T>> minAndMax(values: Iterable<T>) =
    Output(values.min(), values.max())
private data class Output<T : Comparable<T>>(val min: T?, val max: T?)

And generate source code like this, containing a modified version of the original source code:

MinAndMax.gen.kt

// This code was generated by the <Company> Test Master node generation tool at 2018-10-29T08:31:35.847.
//
// Do not modify this file. Any changes may be overwritten at a later time.
package com.mycompany.testmaster.playground.nodes.gen

import com.mycompany.testmaster.domain.ElementaryNode
import com.mycompany.testmaster.domain.InputPort
import com.mycompany.testmaster.domain.OutputPort
import com.mycompany.testmaster.domain.Port
import kotlin.collections.Set
import kotlinx.coroutines.async
import kotlinx.coroutines.coroutineScope

class MinAndMax<T : Comparable<in T>> : ElementaryNode() {
    private val _values: Port<Iterable<out T>> = Port<Iterable<out T>>()

    val values: InputPort<Iterable<out T>> = _values

    private val _min: Port<T?> = Port<T?>()

    val min: OutputPort<T?> = _min

    private val _max: Port<T?> = Port<T?>()

    val max: OutputPort<T?> = _max

    override val ports: Set<Port<*>> = setOf(_values, _min, _max)

    override suspend fun executeOnce() {
        coroutineScope {
            val values = async { _values.receive() }
            val output = _nodeBody(values.await())
            _min.forward(output.min)
            _max.forward(output.max)
        }
    }
}



private  fun <T: Comparable<T>> _nodeBody(values: Iterable<T>) =
    Output(values.min(), values.max())
private data class Output<T : Comparable<T>>(val min: T?, val max: T?)