POC - Compiler Plugin (IR) to wrap Theme Composable around Preview function body

android/foundation/designsystem/build.gradle.kts

@@ -1,6 +1,7 @@
 plugins {
     id("kstreamlined.android.library")
     id("kstreamlined.android.screenshot-test")
+    id("io.github.reactivecircus.cocoon")
     id("kstreamlined.compose")
 }
 
@@ -9,6 +10,11 @@ android {
     androidResources.enable = true
 }
 
+cocoon {
+    annotation.set("io.github.reactivecircus.kstreamlined.android.foundation.designsystem.preview.PreviewKStreamlined")
+    wrappingFunction.set("io.github.reactivecircus.kstreamlined.android.foundation.designsystem.preview.KSThemeWithSurface")
+}
+
 dependencies {
     implementation(libs.androidx.compose.materialIcons)
     implementation(libs.androidx.compose.material3)

android/foundation/designsystem/src/main/kotlin/io/github/reactivecircus/kstreamlined/android/foundation/designsystem/component/Button.kt

@@ -11,9 +11,9 @@ import androidx.compose.ui.semantics.Role
 import androidx.compose.ui.semantics.role
 import androidx.compose.ui.semantics.semantics
 import androidx.compose.ui.text.font.FontWeight
-import androidx.compose.ui.tooling.preview.PreviewLightDark
 import androidx.compose.ui.unit.dp
 import io.github.reactivecircus.kstreamlined.android.foundation.designsystem.foundation.KSTheme
+import io.github.reactivecircus.kstreamlined.android.foundation.designsystem.preview.PreviewKStreamlined
 
 @Composable
 public fun Button(
@@ -50,15 +50,11 @@ public fun Button(
 }
 
 @Composable
-@PreviewLightDark
+@PreviewKStreamlined
 private fun PreviewButton() {
-    KSTheme {
-        Surface {
-            Button(
-                text = "Button",
-                onClick = {},
-                modifier = Modifier.padding(8.dp),
-            )
-        }
-    }
+    Button(
+        text = "Button",
+        onClick = {},
+        modifier = Modifier.padding(8.dp),
+    )
 }

android/foundation/designsystem/src/main/kotlin/io/github/reactivecircus/kstreamlined/android/foundation/designsystem/preview/PreviewKStreamlined.kt

@@ -0,0 +1,22 @@
+package io.github.reactivecircus.kstreamlined.android.foundation.designsystem.preview
+
+import androidx.compose.runtime.Composable
+import androidx.compose.ui.tooling.preview.PreviewLightDark
+import io.github.reactivecircus.kstreamlined.android.foundation.designsystem.component.Surface
+import io.github.reactivecircus.kstreamlined.android.foundation.designsystem.foundation.KSTheme
+
+@Retention(AnnotationRetention.BINARY)
+@Target(AnnotationTarget.ANNOTATION_CLASS, AnnotationTarget.FUNCTION)
+@PreviewLightDark
+public annotation class PreviewKStreamlined
+
+@Composable
+public fun KSThemeWithSurface(
+    content: @Composable () -> Unit,
+) {
+    KSTheme {
+        Surface {
+            content()
+        }
+    }
+}

build-logic/cocoon/cocoon-compiler-plugin/build.gradle.kts

@@ -0,0 +1,52 @@
+import dev.detekt.gradle.Detekt
+import org.gradle.kotlin.dsl.withType
+import org.jetbrains.kotlin.gradle.dsl.JvmTarget
+import org.jetbrains.kotlin.gradle.tasks.KotlinCompile
+
+plugins {
+    alias(libs.plugins.kotlin.jvm)
+    alias(libs.plugins.detekt)
+}
+
+group = "io.github.reactivecircus.cocoon"
+version = "0.1.0"
+
+kotlin {
+    explicitApi()
+}
+
+tasks.withType<KotlinCompile>().configureEach {
+    compilerOptions {
+        jvmTarget.set(JvmTarget.JVM_11)
+        optIn.add("org.jetbrains.kotlin.compiler.plugin.ExperimentalCompilerApi")
+    }
+}
+
+tasks.withType<JavaCompile>().configureEach {
+    sourceCompatibility = JavaVersion.VERSION_11.toString()
+    targetCompatibility = JavaVersion.VERSION_11.toString()
+}
+
+detekt {
+    source.from(files("src/"))
+    config.from(files("$rootDir/../detekt.yml"))
+    buildUponDefaultConfig = true
+    parallel = true
+}
+
+tasks.withType<Detekt>().configureEach {
+    jvmTarget = JvmTarget.JVM_11.target
+    reports {
+        xml.required.set(false)
+        sarif.required.set(false)
+        md.required.set(false)
+    }
+}
+
+dependencies {
+    // enable Ktlint formatting
+    detektPlugins(libs.plugin.detektKtlintWrapper)
+
+    compileOnly(libs.kotlin.compiler)
+    compileOnly(libs.kotlin.stblib)
+}

build-logic/cocoon/cocoon-compiler-plugin/gradle.properties

@@ -0,0 +1 @@
+kotlin.stdlib.default.dependency=false

build-logic/cocoon/cocoon-compiler-plugin/src/main/kotlin/io/github/reactivecircus/cocoon/compiler/CocoonCommandLineProcessor.kt

@@ -0,0 +1,38 @@
+package io.github.reactivecircus.cocoon.compiler
+
+import org.jetbrains.kotlin.compiler.plugin.AbstractCliOption
+import org.jetbrains.kotlin.compiler.plugin.CliOption
+import org.jetbrains.kotlin.compiler.plugin.CommandLineProcessor
+import org.jetbrains.kotlin.config.CompilerConfiguration
+import org.jetbrains.kotlin.config.CompilerConfigurationKey
+
+public class CocoonCommandLineProcessor : CommandLineProcessor {
+    override val pluginId: String = "io.github.reactivecircus.cocoon.compiler"
+
+    @Suppress("MaxLineLength")
+    override val pluginOptions: Collection<AbstractCliOption> = listOf(
+        CliOption(
+            optionName = CompilerOptions.Annotation.toString(),
+            valueDescription = "Fully qualified annotation class name",
+            description = "The fully qualified name of the annotation to be used for marking functions for transformation.",
+        ),
+        CliOption(
+            optionName = CompilerOptions.WrappingFunction.toString(),
+            valueDescription = "Fully qualified name of the higher-order function to be used for wrapping the transformed function's body.",
+            description = "The fully qualified name of the function to be used for wrapping the transformed function.",
+        ),
+    )
+
+    override fun processOption(option: AbstractCliOption, value: String, configuration: CompilerConfiguration) {
+        when (option.optionName) {
+            CompilerOptions.Annotation.toString() -> configuration.put(CompilerOptions.Annotation, value)
+            CompilerOptions.WrappingFunction.toString() -> configuration.put(CompilerOptions.WrappingFunction, value)
+            else -> throw IllegalArgumentException("Unknown plugin option: ${option.optionName}")
+        }
+    }
+
+    internal object CompilerOptions {
+        val Annotation = CompilerConfigurationKey<String>("annotation")
+        val WrappingFunction = CompilerConfigurationKey<String>("wrappingFunction")
+    }
+}

build-logic/cocoon/cocoon-compiler-plugin/src/main/kotlin/io/github/reactivecircus/cocoon/compiler/CocoonCompilerPluginRegistrar.kt

@@ -0,0 +1,36 @@
+package io.github.reactivecircus.cocoon.compiler
+
+import org.jetbrains.kotlin.backend.common.extensions.IrGenerationExtension
+import org.jetbrains.kotlin.cli.common.messages.MessageCollector
+import org.jetbrains.kotlin.compiler.plugin.CompilerPluginRegistrar
+import org.jetbrains.kotlin.config.CommonConfigurationKeys
+import org.jetbrains.kotlin.config.CompilerConfiguration
+
+public class CocoonCompilerPluginRegistrar : CompilerPluginRegistrar() {
+    override val supportsK2: Boolean = true
+
+    override fun ExtensionStorage.registerExtensions(configuration: CompilerConfiguration) {
+        val annotationString = requireNotNull(
+            configuration.get(CocoonCommandLineProcessor.CompilerOptions.Annotation),
+        )
+        val annotationClassId = annotationString.toClassId()
+
+        val wrappingFunctionString = requireNotNull(
+            configuration.get(CocoonCommandLineProcessor.CompilerOptions.WrappingFunction),
+        )
+        val wrappingFunctionCallableId = wrappingFunctionString.toCallableId()
+
+        val messageCollector = configuration.get(
+            CommonConfigurationKeys.MESSAGE_COLLECTOR_KEY,
+            MessageCollector.NONE,
+        )
+
+        IrGenerationExtension.registerExtension(
+            extension = CocoonIrGenerationExtension(
+                annotationName = annotationClassId,
+                wrappingFunctionName = wrappingFunctionCallableId,
+                messageCollector = messageCollector,
+            ),
+        )
+    }
+}

build-logic/cocoon/cocoon-compiler-plugin/src/main/kotlin/io/github/reactivecircus/cocoon/compiler/CocoonFunctionTransformer.kt

@@ -0,0 +1,106 @@
+package io.github.reactivecircus.cocoon.compiler
+
+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.cli.common.messages.CompilerMessageSeverity
+import org.jetbrains.kotlin.cli.common.messages.MessageCollector
+import org.jetbrains.kotlin.descriptors.DescriptorVisibilities
+import org.jetbrains.kotlin.ir.IrStatement
+import org.jetbrains.kotlin.ir.UNDEFINED_OFFSET
+import org.jetbrains.kotlin.ir.builders.declarations.buildFun
+import org.jetbrains.kotlin.ir.builders.irBlock
+import org.jetbrains.kotlin.ir.builders.irCall
+import org.jetbrains.kotlin.ir.declarations.IrDeclarationOrigin
+import org.jetbrains.kotlin.ir.declarations.IrFunction
+import org.jetbrains.kotlin.ir.declarations.IrSimpleFunction
+import org.jetbrains.kotlin.ir.declarations.createBlockBody
+import org.jetbrains.kotlin.ir.expressions.IrExpression
+import org.jetbrains.kotlin.ir.expressions.IrStatementOrigin
+import org.jetbrains.kotlin.ir.expressions.impl.IrFunctionExpressionImpl
+import org.jetbrains.kotlin.ir.symbols.UnsafeDuringIrConstructionAPI
+import org.jetbrains.kotlin.ir.types.IrType
+import org.jetbrains.kotlin.ir.util.dump
+import org.jetbrains.kotlin.ir.util.hasAnnotation
+import org.jetbrains.kotlin.ir.util.statements
+import org.jetbrains.kotlin.name.CallableId
+import org.jetbrains.kotlin.name.ClassId
+import org.jetbrains.kotlin.name.SpecialNames
+
+internal class CocoonFunctionTransformer(
+    private val pluginContext: IrPluginContext,
+    private val messageCollector: MessageCollector,
+    private val annotation: ClassId,
+    private val wrappingFunction: CallableId,
+) : IrElementTransformerVoidWithContext() {
+    @OptIn(UnsafeDuringIrConstructionAPI::class)
+    override fun visitFunctionNew(declaration: IrFunction): IrStatement {
+        if (!declaration.hasAnnotation(annotation) || declaration.body == null) {
+            return super.visitFunctionNew(declaration)
+        }
+
+        // TODO check for $composer and report error
+
+        val originalBody = declaration.body!!
+
+        declaration.body = pluginContext.irFactory.createBlockBody(
+            startOffset = originalBody.startOffset,
+            endOffset = originalBody.endOffset,
+        ).apply {
+            val wrappingFunction = pluginContext.referenceFunctions(wrappingFunction).single()
+            val irBuilder = DeclarationIrBuilder(pluginContext, declaration.symbol)
+
+            // TODO move up and check early:
+            //  - must have at least 1 param
+            //  - last must be kotlin.Function0<kotlin.Unit>)
+            //  - move to FIR?
+            val wrappingFunctionParameters = wrappingFunction.owner.parameters
+
+            statements.add(
+                irBuilder.irBlock {
+                    +irCall(wrappingFunction).apply {
+                        val lambdaExpression = pluginContext.createLambdaIrFunctionExpression(
+                            lambdaReturnType = wrappingFunctionParameters.last().type,
+                        ) {
+                            parent = declaration
+                            body = pluginContext.irFactory.createBlockBody(
+                                startOffset,
+                                endOffset,
+                                originalBody.statements,
+                            )
+                        }
+                        arguments[wrappingFunctionParameters.size - 1] = lambdaExpression
+                    }
+                },
+            )
+        }
+
+        log("Transformed function IR: \n${declaration.dump()}")
+
+        return super.visitFunctionNew(declaration)
+    }
+
+    private fun IrPluginContext.createLambdaIrFunctionExpression(
+        lambdaReturnType: IrType,
+        block: IrSimpleFunction.() -> Unit = {},
+    ): IrExpression {
+        val lambda = irFactory.buildFun {
+            name = SpecialNames.ANONYMOUS
+            origin = IrDeclarationOrigin.LOCAL_FUNCTION_FOR_LAMBDA
+            visibility = DescriptorVisibilities.LOCAL
+            returnType = lambdaReturnType
+        }.apply(block)
+
+        return IrFunctionExpressionImpl(
+            startOffset = UNDEFINED_OFFSET,
+            endOffset = UNDEFINED_OFFSET,
+            type = lambda.returnType,
+            function = lambda,
+            origin = IrStatementOrigin.LAMBDA,
+        )
+    }
+
+    private fun log(message: String) {
+        messageCollector.report(CompilerMessageSeverity.LOGGING, "Cocoon Compiler Plugin (IR) - $message")
+    }
+}

build-logic/cocoon/cocoon-compiler-plugin/src/main/kotlin/io/github/reactivecircus/cocoon/compiler/CocoonIrGenerationExtension.kt

@@ -0,0 +1,33 @@
+package io.github.reactivecircus.cocoon.compiler
+
+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.declarations.IrModuleFragment
+import org.jetbrains.kotlin.name.CallableId
+import org.jetbrains.kotlin.name.ClassId
+
+internal class CocoonIrGenerationExtension(
+    private val annotationName: ClassId,
+    private val wrappingFunctionName: CallableId,
+    private val messageCollector: MessageCollector,
+) : IrGenerationExtension {
+    override fun generate(moduleFragment: IrModuleFragment, pluginContext: IrPluginContext) {
+        if (pluginContext.referenceClass(annotationName) == null) {
+            messageCollector.report(CompilerMessageSeverity.ERROR, "Could not find annotation class <$annotationName>.")
+            return
+        }
+        if (pluginContext.referenceFunctions(wrappingFunctionName).isEmpty()) {
+            messageCollector.report(
+                CompilerMessageSeverity.ERROR,
+                "Could not find wrapping function <$wrappingFunctionName>.",
+            )
+            return
+        }
+        moduleFragment.transform(
+            CocoonFunctionTransformer(pluginContext, messageCollector, annotationName, wrappingFunctionName),
+            null,
+        )
+    }
+}

build-logic/cocoon/cocoon-compiler-plugin/src/main/kotlin/io/github/reactivecircus/cocoon/compiler/utils.kt

@@ -0,0 +1,11 @@
+package io.github.reactivecircus.cocoon.compiler
+
+import org.jetbrains.kotlin.name.CallableId
+import org.jetbrains.kotlin.name.ClassId
+import org.jetbrains.kotlin.name.FqName
+
+internal fun String.toClassId(): ClassId =
+    FqName(this).run { ClassId(parent(), shortName()) }
+
+internal fun String.toCallableId(): CallableId =
+    FqName(this).run { CallableId(parent(), shortName()) }

build-logic/cocoon/cocoon-compiler-plugin/src/main/resources/META-INF/services/org.jetbrains.kotlin.compiler.plugin.CommandLineProcessor

@@ -0,0 +1 @@
+io.github.reactivecircus.cocoon.compiler.CocoonCommandLineProcessor

build-logic/cocoon/cocoon-compiler-plugin/src/main/resources/META-INF/services/org.jetbrains.kotlin.compiler.plugin.CompilerPluginRegistrar

@@ -0,0 +1 @@
+io.github.reactivecircus.cocoon.compiler.CocoonCompilerPluginRegistrar