Solution: KSP execution measurement wrapper
Here is an example solution code:
package academy.kt
import com.google.devtools.ksp.closestClassDeclaration
import com.google.devtools.ksp.getDeclaredFunctions
import com.google.devtools.ksp.isConstructor
import com.google.devtools.ksp.isPublic
import com.google.devtools.ksp.processing.*
import com.google.devtools.ksp.symbol.*
import com.squareup.kotlinpoet.*
import com.squareup.kotlinpoet.ksp.toAnnotationSpec
import com.squareup.kotlinpoet.ksp.toKModifier
import com.squareup.kotlinpoet.ksp.toTypeName
import java.io.OutputStreamWriter
import java.nio.charset.StandardCharsets
class MeasuredWrapperGenerator(
private val codeGenerator: CodeGenerator,
private val logger: KSPLogger,
) : SymbolProcessor {
private val annotationName = Measured::class.qualifiedName!!
override fun process(resolver: Resolver): List<KSAnnotated> {
resolver
.getSymbolsWithAnnotation(
annotationName
)
.filterIsInstance<KSFunctionDeclaration>()
.groupBy { it.closestClassDeclaration() }
.forEach { (classDeclaration, _) ->
if (classDeclaration != null) {
generateMeasuredClass(classDeclaration)
}
}
return emptyList()
}
private fun generateMeasuredClass(
classElement: KSClassDeclaration
) {
val className = classElement.simpleName.getShortName()
val measuredName = "Measured$className"
val measuredPackage = classElement.packageName.asString()
val publicMethods = classElement
.getDeclaredFunctions()
.filter { !it.isConstructor() && it.isPublic() }
.toList()
val fileSpec = FileSpec.builder(
measuredPackage,
"$measuredName.kt"
)
.addType(
TypeSpec.classBuilder(measuredName)
.primaryConstructor(
FunSpec.constructorBuilder()
.addParameter(
"wrapper",
classElement.asType(emptyList())
.toTypeName()
)
.build()
)
.addFunction(
FunSpec.constructorBuilder()
.addParameters(
classElement.primaryConstructor!!.parameters
.map { buildInterfaceMethodParameter(it) }
)
.callThisConstructor(
"$className(${
classElement.primaryConstructor!!.parameters.joinToString {
it.name?.getShortName().orEmpty()
}
})"
)
.build()
)
.addProperty(
PropertySpec.builder(
"wrapper",
classElement.asType(emptyList())
.toTypeName()
).initializer("wrapper")
.build()
)
.addFunctions(
publicMethods
.map { buildMethod(className, it) }
.toList()
)
.build()
)
.build()
val dependencies = Dependencies(
aggregating = false,
classElement.containingFile!!
)
val file = codeGenerator.createNewFile(
dependencies, measuredPackage, measuredName
)
OutputStreamWriter(file, StandardCharsets.UTF_8)
.use(fileSpec::writeTo)
}
private fun buildMethod(
className: String,
method: KSFunctionDeclaration
): FunSpec {
val methodName = method.simpleName.getShortName()
return FunSpec.builder(methodName)
.addModifiers(
method.modifiers
.mapNotNull { it.toKModifier() }.toList()
)
.addParameters(
method.parameters
.map { buildInterfaceMethodParameter(it) }
)
.returns(method.returnType!!.toTypeName())
.addAnnotations(method
.annotations
.filter { !isMeasured(it) }
.map { it.toAnnotationSpec() }
.toList())
.addCode(
if (method.annotations.none { isMeasured(it) })
"return wrapper.$methodName(${
method.parameters.joinToString {
it.name?.getShortName().orEmpty()
}
})"
else
"""
val before = System.currentTimeMillis()
val value = wrapper.$methodName(${
method.parameters.joinToString {
it.name?.getShortName().orEmpty()
}
})
val after = System.currentTimeMillis()
println("$methodName from $className took ${'$'}{after-before} ms")
return value
""".trimIndent()
)
.build()
}
private fun isMeasured(annotation: KSAnnotation) = annotation
.annotationType
.resolve()
.declaration
.qualifiedName
?.asString() == annotationName
private fun buildInterfaceMethodParameter(
variableElement: KSValueParameter,
): ParameterSpec = ParameterSpec
.builder(
variableElement.name!!.getShortName(),
variableElement.type.toTypeName(),
)
.addAnnotations(
variableElement.annotations
.map { it.toAnnotationSpec() }.toList()
)
.build()
}
