@CaptureSource and sourceOf (#4)

mfwgenerics · Jan 1, 2023 · 974cc9f · 974cc9f
1 parent 4452b54
commit 974cc9f
Show file tree

Hide file tree

Showing 8 changed files with 274 additions and 46 deletions.
diff --git a/README.md b/README.md
@@ -1,5 +1,5 @@
 # Kapshot
-Kapshot is a simple Kotlin compiler plugin for capturing source code text from closures.
+Kapshot is a simple Kotlin compiler plugin for capturing source code text from closure blocks and declarations.
 
 ## Usage
 
@@ -13,6 +13,8 @@ plugins {
 }
 ```
 
+### Capturing Blocks
+
 Now your Kotlin code can use `CapturedBlock<T>` as a source enriched replacement for `() -> T`.
 You can call `source()` on any instance of
 `CapturedBlock` to access the source text for that block.
@@ -41,18 +43,7 @@ fun equation(block: CapturedBlock<Int>): String {
 check(equation { 2 + 2 } == "2 + 2 = 4")
 ```
 
-## Purpose
-
-The purpose of this plugin is to support experimental literate
-programming and documentation generation techniques in Kotlin.
-
-An example of this is the code used to generate this README.md.
-Capturing source from blocks allows sample code to be run and
-tested during generation.
-
-View the source here: [readme/src/main/kotlin/Main.kt](readme/src/main/kotlin/Main.kt)
-
-## Parameterized Blocks
+### Parameterized Blocks
 
 The default `CapturedBlock` interface doesn't accept any
 arguments to `invoke` and is only generic on the return type. This
@@ -61,7 +52,6 @@ enclosing scope. To write source capturing versions of builder blocks
 or common higher-order functions like `map` and `filter` you will
 need to define your own capture interface that extends `Capturable`.
 
-
 ```kotlin
 /* must be a fun interface to support SAM conversion from blocks */
 fun interface CustomCapturable<T, R> : Capturable<CustomCapturable<T, R>> {
@@ -88,4 +78,46 @@ check(
 )
 ```
 
-If it is present, the block's argument list is considered part of its source text.
+If it is present, the block's argument list is considered part of its source text.
+
+### Declarations
+
+You can capture declaration sources using the `@CaptureSource`
+annotation. The source of annotated declarations can then be retrieved using
+`sourceOf<T>` for class declarations or `sourceOf(::method)` for method
+declarations. The source capture starts at the end of the `@CaptureSource`
+annotation.
+
+```kotlin
+@CaptureSource
+class MyClass {
+    @CaptureSource
+    fun twelve() = 12
+}
+
+check(
+    sourceOf<MyClass>() ==
+    """
+    class MyClass {
+        @CaptureSource
+        fun twelve() = 12
+    }
+    """.trimIndent()
+)
+
+check(
+    sourceOf(MyClass::twelve) ==
+    "fun twelve() = 12"
+)
+```
+
+## Purpose
+
+The purpose of this plugin is to support experimental literate
+programming and documentation generation techniques in Kotlin.
+
+An example of this is the code used to generate this README.md.
+Capturing source from blocks allows sample code to be run and
+tested during generation.
+
+View the source here: [readme/src/main/kotlin/Main.kt](readme/src/main/kotlin/Main.kt)
diff --git a/kapshot-plugin-kotlin/src/main/kotlin/io/koalaql/kapshot/plugin/CaptureTransformer.kt b/kapshot-plugin-kotlin/src/main/kotlin/io/koalaql/kapshot/plugin/CaptureTransformer.kt
@@ -4,25 +4,40 @@ import org.jetbrains.kotlin.backend.common.IrElementTransformerVoidWithContext
 import org.jetbrains.kotlin.backend.common.extensions.IrPluginContext
 import org.jetbrains.kotlin.backend.common.lower.DeclarationIrBuilder
 import org.jetbrains.kotlin.backend.common.sourceElement
+import org.jetbrains.kotlin.backend.jvm.ir.getStringConstArgument
+import org.jetbrains.kotlin.ir.IrElement
+import org.jetbrains.kotlin.ir.IrStatement
+import org.jetbrains.kotlin.ir.builders.IrSingleStatementBuilder
 import org.jetbrains.kotlin.ir.builders.irCall
 import org.jetbrains.kotlin.ir.builders.irString
 import org.jetbrains.kotlin.ir.declarations.IrClass
+import org.jetbrains.kotlin.ir.declarations.IrDeclarationBase
+import org.jetbrains.kotlin.ir.declarations.IrFunction
 import org.jetbrains.kotlin.ir.declarations.path
 import org.jetbrains.kotlin.ir.expressions.IrExpression
 import org.jetbrains.kotlin.ir.expressions.IrTypeOperator
 import org.jetbrains.kotlin.ir.expressions.IrTypeOperatorCall
+import org.jetbrains.kotlin.ir.expressions.impl.IrConstImpl
 import org.jetbrains.kotlin.ir.symbols.IrSimpleFunctionSymbol
 import org.jetbrains.kotlin.ir.types.IrSimpleType
 import org.jetbrains.kotlin.ir.types.IrType
+import org.jetbrains.kotlin.ir.util.dump
 import org.jetbrains.kotlin.ir.util.dumpKotlinLike
 import org.jetbrains.kotlin.ir.util.kotlinFqName
+import org.jetbrains.kotlin.ir.visitors.IrElementTransformer
 import java.io.File
 
 class CaptureTransformer(
     private val context: IrPluginContext,
     private val addSourceToBlock: IrSimpleFunctionSymbol,
-    private val capturableFqn: String = "io.koalaql.kapshot.Capturable"
+    private val capturableFqn: String = "io.koalaql.kapshot.Capturable",
+    private val captureSourceFqn: String = "io.koalaql.kapshot.CaptureSource",
 ): IrElementTransformerVoidWithContext() {
+    fun currentFileText(): String {
+        /* https://youtrack.jetbrains.com/issue/KT-41888 */
+        return File(currentFile.path).readText().replace("\r\n", "\n")
+    }
+
     private fun transformSam(expression: IrTypeOperatorCall): IrExpression {
         val symbol = currentScope!!.scope.scopeOwnerSymbol
 
@@ -34,8 +49,7 @@ class CaptureTransformer(
                 expression.typeOperand
             )
 
-            /* https://youtrack.jetbrains.com/issue/KT-41888 */
-            val fileText = File(currentFile.path).readText().replace("\r\n", "\n")
+            val fileText = currentFileText()
 
             var startOffset = sourceElement.startOffset
             var endOffset = sourceElement.endOffset
@@ -85,4 +99,26 @@ class CaptureTransformer(
 
         return super.visitTypeOperator(expression)
     }
+
+    override fun visitDeclaration(declaration: IrDeclarationBase): IrStatement {
+        val captureSource = declaration.annotations.singleOrNull {
+            typeIsFqn(it.type, captureSourceFqn)
+        }
+
+        if (captureSource != null) {
+            val fileText = currentFileText()
+
+            /* we start from end of captureSource rather than declaration.startOffset to exclude the capture annotation */
+            val startOffset = captureSource.endOffset
+            val endOffset = declaration.endOffset
+
+            val trimmed = fileText.substring(startOffset, endOffset).trimIndent().trim()
+
+            captureSource.putValueArgument(0,
+                IrConstImpl.string(captureSource.startOffset, captureSource.endOffset, context.irBuiltIns.stringType, trimmed)
+            )
+        }
+
+        return super.visitDeclaration(declaration)
+    }
 }
diff --git a/kapshot-plugin-kotlin/src/main/kotlin/io/koalaql/kapshot/plugin/GenerationExtension.kt b/kapshot-plugin-kotlin/src/main/kotlin/io/koalaql/kapshot/plugin/GenerationExtension.kt
@@ -4,7 +4,10 @@ import org.jetbrains.kotlin.backend.common.extensions.IrGenerationExtension
 import org.jetbrains.kotlin.backend.common.extensions.IrPluginContext
 import org.jetbrains.kotlin.cli.common.messages.CompilerMessageSeverity
 import org.jetbrains.kotlin.cli.common.messages.MessageCollector
+import org.jetbrains.kotlin.ir.IrElement
 import org.jetbrains.kotlin.ir.declarations.IrModuleFragment
+import org.jetbrains.kotlin.ir.util.dump
+import org.jetbrains.kotlin.ir.visitors.IrElementVisitor
 import org.jetbrains.kotlin.name.CallableId
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.name.Name

diff --git a/kapshot-runtime/build.gradle.kts b/kapshot-runtime/build.gradle.kts
@@ -6,4 +6,8 @@ plugins {
     id("publish")
 
     kotlin("jvm") version "1.8.0"
+}
+
+dependencies {
+    api(kotlin("reflect"))
 }
diff --git a/kapshot-runtime/src/main/kotlin/io/koalaql/kapshot/CaptureSource.kt b/kapshot-runtime/src/main/kotlin/io/koalaql/kapshot/CaptureSource.kt
@@ -0,0 +1,9 @@
+package io.koalaql.kapshot
+
+@Target(
+    AnnotationTarget.CLASS,
+    AnnotationTarget.FUNCTION
+)
+annotation class CaptureSource(
+    val text: String = ""
+)
diff --git a/kapshot-runtime/src/main/kotlin/io/koalaql/kapshot/sourceOf.kt b/kapshot-runtime/src/main/kotlin/io/koalaql/kapshot/sourceOf.kt
@@ -0,0 +1,29 @@
+package io.koalaql.kapshot
+
+import kotlin.reflect.KClass
+import kotlin.reflect.KFunction
+import kotlin.reflect.KType
+import kotlin.reflect.typeOf
+
+private fun sourceOf(annotations: List<Annotation>, name: () -> String): String {
+    val anno = annotations
+        .singleOrNull { it.annotationClass == CaptureSource::class }
+        ?: error("sourceOf call on ${name()} requires the CaptureSource annotation")
+
+    anno as CaptureSource
+
+    check(anno.text.isNotBlank()) { "missing source text" }
+
+    return anno.text
+}
+
+fun sourceOf(type: KType): String {
+    val classifier = checkNotNull(type.classifier as? KClass<*>) { "$type is not a supported target of source capture"}
+
+    return sourceOf(classifier.annotations) { "$type" }
+}
+
+inline fun <reified T : Any> sourceOf() = sourceOf(typeOf<T>())
+
+fun sourceOf(method: KFunction<*>): String =
+    sourceOf(method.annotations) { method.name }
diff --git a/readme/src/main/kotlin/Main.kt b/readme/src/main/kotlin/Main.kt
@@ -1,5 +1,7 @@
 import io.koalaql.kapshot.Capturable
+import io.koalaql.kapshot.CaptureSource
 import io.koalaql.kapshot.CapturedBlock
+import io.koalaql.kapshot.sourceOf
 import kotlin.io.path.Path
 import kotlin.io.path.writeText
 
@@ -8,21 +10,27 @@ fun execSource(block: CapturedBlock<*>): String {
     return block.source()
 }
 
+@CaptureSource
+/* must be a fun interface to support SAM conversion from blocks */
 fun interface CustomCapturable<T, R> : Capturable<CustomCapturable<T, R>> {
+    /* invoke is not special. this could be any single abstract method */
     operator fun invoke(arg: T): R
 
+    /* withSource is called by the plugin to add source information */
     override fun withSource(source: String): CustomCapturable<T, R> =
         object : CustomCapturable<T, R> by this { override fun source(): String = source }
 }
 
 fun generateMarkdown(): String {
     val capturableClass = Capturable::class
     val blockClass = CapturedBlock::class
+    val captureAnno = CaptureSource::class
+
     val importStatement = "import ${blockClass.qualifiedName}"
 
     return """
 # Kapshot
-Kapshot is a simple Kotlin compiler plugin for capturing source code text from closures.
+Kapshot is a simple Kotlin compiler plugin for capturing source code text from closure blocks and declarations.
 
 ## Usage
 
@@ -36,6 +44,8 @@ plugins {
 }
 ```
 
+### Capturing Blocks
+
 Now your Kotlin code can use `${blockClass.simpleName}<T>` as a source enriched replacement for `() -> T`.
 You can call `${CapturedBlock<*>::source.name}()` on any instance of
 `${blockClass.simpleName}` to access the source text for that block.
@@ -72,18 +82,7 @@ ${
 }
 ```
 
-## Purpose
-
-The purpose of this plugin is to support experimental literate
-programming and documentation generation techniques in Kotlin.
-
-An example of this is the code used to generate this README.md.
-Capturing source from blocks allows sample code to be run and
-tested during generation.
-
-View the source here: [readme/src/main/kotlin/Main.kt](readme/src/main/kotlin/Main.kt)
-
-## Parameterized Blocks
+### Parameterized Blocks
 
 The default `${blockClass.simpleName}` interface doesn't accept any
 arguments to `invoke` and is only generic on the return type. This
@@ -92,21 +91,8 @@ enclosing scope. To write source capturing versions of builder blocks
 or common higher-order functions like `map` and `filter` you will
 need to define your own capture interface that extends `${capturableClass.simpleName}`.
 
-${""/*
-The code below can't be captured from a block due to limitations on where
-interfaces can be declared. Perhaps a future version of this plugin could
-add a sourceOf<T>() operator for class/interface definitions
-*/}
 ```kotlin
-/* must be a fun interface to support SAM conversion from blocks */
-fun interface CustomCapturable<T, R> : Capturable<CustomCapturable<T, R>> {
-    /* invoke is not special. this could be any single abstract method */
-    operator fun invoke(arg: T): R
-
-    /* withSource is called by the plugin to add source information */
-    override fun withSource(source: String): CustomCapturable<T, R> =
-        object : CustomCapturable<T, R> by this { override fun source(): String = source }
-}
+${sourceOf<CustomCapturable<*, *>>()}
 ```
 
 Once you have declared your own `${capturableClass.simpleName}` you can use it
@@ -127,7 +113,54 @@ ${
 }
 ```
 
-If it is present, the block's argument list is considered part of its source text. 
+If it is present, the block's argument list is considered part of its source text.
+ 
+### Declarations
+
+You can capture declaration sources using the `@${captureAnno.simpleName}`
+annotation. The source of annotated declarations can then be retrieved using
+`sourceOf<T>` for class declarations or `sourceOf(::method)` for method
+declarations. The source capture starts at the end of the `@${captureAnno.simpleName}`
+annotation.
+
+```kotlin
+${
+    execSource {
+        @CaptureSource
+        class MyClass {
+            @CaptureSource
+            fun twelve() = 12
+        }
+        
+        check(
+            sourceOf<MyClass>() ==
+            """
+            class MyClass {
+                @CaptureSource
+                fun twelve() = 12
+            }
+            """.trimIndent()
+        )
+        
+        check(
+            sourceOf(MyClass::twelve) ==
+            "fun twelve() = 12"
+        )
+    }
+}
+```
+
+## Purpose
+
+The purpose of this plugin is to support experimental literate
+programming and documentation generation techniques in Kotlin.
+
+An example of this is the code used to generate this README.md.
+Capturing source from blocks allows sample code to be run and
+tested during generation.
+
+View the source here: [readme/src/main/kotlin/Main.kt](readme/src/main/kotlin/Main.kt)
+
     """.trim()
 }