Kotlin recipe DSL: preserve scan/generate phases in K2 synthesis

rewrite-kotlin/src/main/kotlin/org/openrewrite/RecipeDsl.kt

@@ -713,3 +713,61 @@ public fun buildImperativeVisitor(
     builder.block()
     return builder.build("", "", emptySet(), null).getVisitor()
 }
+
+private fun buildImperativeRecipe(block: RecipeBuilder.() -> Unit): Recipe {
+    val builder = RecipeBuilder()
+    builder.block()
+    return builder.build("", "", emptySet(), null)
+}
+
+/**
+ * Scanning counterparts to [buildImperativeVisitor], one per [ScanningRecipe]
+ * lifecycle method. The K2 plugin's synthesized `<Name>$KtRecipe` extends
+ * [ScanningRecipe] and routes each call here so all phases survive synthesis.
+ * The accumulator is threaded by the framework, so the generated class stays
+ * field-less; an edit-only block leaves [buildImperativeRecipe] returning a
+ * plain [Recipe], for which the scanner/generate helpers no-op.
+ */
+public fun buildImperativeInitialValue(
+    ctx: ExecutionContext,
+    block: RecipeBuilder.() -> Unit,
+): Any {
+    val recipe = buildImperativeRecipe(block)
+    @Suppress("UNCHECKED_CAST")
+    return if (recipe is ScanningRecipe<*>) (recipe as ScanningRecipe<Any>).getInitialValue(ctx) else Unit
+}
+
+public fun buildImperativeScanner(
+    acc: Any,
+    block: RecipeBuilder.() -> Unit,
+): TreeVisitor<*, ExecutionContext> {
+    val recipe = buildImperativeRecipe(block)
+    @Suppress("UNCHECKED_CAST")
+    return if (recipe is ScanningRecipe<*>) (recipe as ScanningRecipe<Any>).getScanner(acc)
+    else TreeVisitor.noop<Tree, ExecutionContext>()
+}
+
+public fun buildImperativeEditVisitor(
+    acc: Any,
+    block: RecipeBuilder.() -> Unit,
+): TreeVisitor<*, ExecutionContext> {
+    val recipe = buildImperativeRecipe(block)
+    @Suppress("UNCHECKED_CAST")
+    return if (recipe is ScanningRecipe<*>) (recipe as ScanningRecipe<Any>).getVisitor(acc)
+    else recipe.getVisitor()
+}
+
+public fun buildImperativeGenerate(
+    acc: Any,
+    generatedInThisCycle: Collection<SourceFile>,
+    ctx: ExecutionContext,
+    block: RecipeBuilder.() -> Unit,
+): Collection<SourceFile> {
+    val recipe = buildImperativeRecipe(block)
+    @Suppress("UNCHECKED_CAST")
+    return if (recipe is ScanningRecipe<*>) {
+        (recipe as ScanningRecipe<Any>).generate(acc, generatedInThisCycle, ctx).toList()
+    } else {
+        emptyList()
+    }
+}

rewrite-kotlin/src/main/kotlin/org/openrewrite/kotlin/recipe/internal/RecipeIrGenerationExtension.kt

@@ -33,11 +33,13 @@ import org.jetbrains.kotlin.descriptors.Modality
 import org.jetbrains.kotlin.ir.IrElement
 import org.jetbrains.kotlin.ir.builders.declarations.addConstructor
 import org.jetbrains.kotlin.ir.builders.declarations.addFunction
+import org.jetbrains.kotlin.ir.builders.declarations.addValueParameter
 import org.jetbrains.kotlin.ir.builders.declarations.buildClass
 import org.jetbrains.kotlin.ir.builders.irBlockBody
 import org.jetbrains.kotlin.ir.builders.irCall
 import org.jetbrains.kotlin.ir.builders.irCallConstructor
 import org.jetbrains.kotlin.ir.builders.irDelegatingConstructorCall
+import org.jetbrains.kotlin.ir.builders.irGet
 import org.jetbrains.kotlin.ir.builders.irInt
 import org.jetbrains.kotlin.ir.builders.irReturn
 import org.jetbrains.kotlin.ir.builders.irString
@@ -69,6 +71,7 @@ import org.jetbrains.kotlin.ir.symbols.IrValueSymbol
 import org.jetbrains.kotlin.ir.types.IrType
 import org.jetbrains.kotlin.ir.types.classFqName
 import org.jetbrains.kotlin.ir.types.defaultType
+import org.jetbrains.kotlin.ir.types.typeWith
 import org.jetbrains.kotlin.ir.util.addChild
 import org.jetbrains.kotlin.ir.util.createThisReceiverParameter
 import org.jetbrains.kotlin.ir.util.deepCopyWithSymbols
@@ -134,6 +137,8 @@ internal class RecipeIrGenerationExtension : IrGenerationExtension {
 
         const val EDIT_SCOPE_REWRITE_FQN = "org.openrewrite.EditScope.rewrite"
         const val RECIPE_BUILDER_EDIT_FQN = "org.openrewrite.RecipeBuilder.edit"
+        const val RECIPE_BUILDER_SCAN_FQN = "org.openrewrite.RecipeBuilder.scan"
+        const val RECIPE_BUILDER_GENERATE_FQN = "org.openrewrite.RecipeBuilder.generate"
 
         /**
          * K2 FIR2IR represents `expr!!` as a call to this synthetic intrinsic,
@@ -216,7 +221,32 @@ internal class RecipeIrGenerationExtension : IrGenerationExtension {
          * DSL's anonymous-Recipe path.
          */
         val buildImperativeVisitorSymbol: IrSimpleFunctionSymbol?,
-    )
+        // ScanningRecipe + its lifecycle methods and the buildImperative* helpers,
+        // for the scan/generate synthesis path. Any null (older rewrite-kotlin
+        // without the helpers) leaves the recipe as a runtime `recipe(...)` call.
+        val scanningRecipeClassSymbol: IrClassSymbol?,
+        val scanningRecipeNoArgCtorSymbol: IrConstructorSymbol?,
+        val scanningGetInitialValue: IrSimpleFunction?,
+        val scanningGetScanner: IrSimpleFunction?,
+        val scanningGetVisitor: IrSimpleFunction?,
+        val scanningGenerate: IrSimpleFunction?,
+        val buildImperativeInitialValueSymbol: IrSimpleFunctionSymbol?,
+        val buildImperativeScannerSymbol: IrSimpleFunctionSymbol?,
+        val buildImperativeEditVisitorSymbol: IrSimpleFunctionSymbol?,
+        val buildImperativeGenerateSymbol: IrSimpleFunctionSymbol?,
+    ) {
+        fun canSynthesizeScanningRecipe(): Boolean =
+            scanningRecipeClassSymbol != null &&
+                scanningRecipeNoArgCtorSymbol != null &&
+                scanningGetInitialValue != null &&
+                scanningGetScanner != null &&
+                scanningGetVisitor != null &&
+                scanningGenerate != null &&
+                buildImperativeInitialValueSymbol != null &&
+                buildImperativeScannerSymbol != null &&
+                buildImperativeEditVisitorSymbol != null &&
+                buildImperativeGenerateSymbol != null
+    }
 
     private fun buildIrGenContext(pluginContext: IrPluginContext): RecipeIrGenContext? {
         val recipeClassId = ClassId.topLevel(FqName("org.openrewrite.Recipe"))
@@ -286,6 +316,36 @@ internal class RecipeIrGenerationExtension : IrGenerationExtension {
             ))
             .singleOrNull { it.owner.valueParameters.size == 1 }
 
+        val scanningRecipeClassId = ClassId.topLevel(FqName("org.openrewrite.ScanningRecipe"))
+        val scanningRecipeClassSymbol = pluginContext.referenceClass(scanningRecipeClassId)
+        val scanningRecipeNoArgCtor = pluginContext.referenceConstructors(scanningRecipeClassId)
+            .singleOrNull { it.owner.valueParameters.isEmpty() }
+        val scanningMembers = scanningRecipeClassSymbol?.owner?.declarations
+            ?.filterIsInstance<IrSimpleFunction>().orEmpty()
+        val scanningGetInitialValue = scanningMembers.firstOrNull {
+            it.name.asString() == "getInitialValue" && it.valueParameters.size == 1
+        }
+        val scanningGetScanner = scanningMembers.firstOrNull {
+            it.name.asString() == "getScanner" && it.valueParameters.size == 1
+        }
+        val scanningGetVisitor = scanningMembers.firstOrNull {
+            it.name.asString() == "getVisitor" && it.valueParameters.size == 1
+        }
+        val scanningGenerate = scanningMembers.firstOrNull {
+            it.name.asString() == "generate" && it.valueParameters.size == 3
+        }
+
+        fun topLevelHelper(name: String, arity: Int): IrSimpleFunctionSymbol? = pluginContext
+            .referenceFunctions(org.jetbrains.kotlin.name.CallableId(
+                packageName = FqName("org.openrewrite"),
+                callableName = Name.identifier(name),
+            ))
+            .singleOrNull { it.owner.valueParameters.size == arity }
+        val buildImperativeInitialValueSymbol = topLevelHelper("buildImperativeInitialValue", 2)
+        val buildImperativeScannerSymbol = topLevelHelper("buildImperativeScanner", 2)
+        val buildImperativeEditVisitorSymbol = topLevelHelper("buildImperativeEditVisitor", 2)
+        val buildImperativeGenerateSymbol = topLevelHelper("buildImperativeGenerate", 4)
+
         return RecipeIrGenContext(
             pluginContext = pluginContext,
             recipeClassSymbol = recipeClassSymbol,
@@ -302,6 +362,16 @@ internal class RecipeIrGenerationExtension : IrGenerationExtension {
             methodInvocationRewriteKotlinNotNullSymbol = kotlinNotNullHelper,
             propertyAccessRewriteKotlinSymbol = propertyAccessHelper,
             buildImperativeVisitorSymbol = buildImperativeVisitorSymbol,
+            scanningRecipeClassSymbol = scanningRecipeClassSymbol,
+            scanningRecipeNoArgCtorSymbol = scanningRecipeNoArgCtor,
+            scanningGetInitialValue = scanningGetInitialValue,
+            scanningGetScanner = scanningGetScanner,
+            scanningGetVisitor = scanningGetVisitor,
+            scanningGenerate = scanningGenerate,
+            buildImperativeInitialValueSymbol = buildImperativeInitialValueSymbol,
+            buildImperativeScannerSymbol = buildImperativeScannerSymbol,
+            buildImperativeEditVisitorSymbol = buildImperativeEditVisitorSymbol,
+            buildImperativeGenerateSymbol = buildImperativeGenerateSymbol,
         )
     }
 
@@ -358,15 +428,29 @@ internal class RecipeIrGenerationExtension : IrGenerationExtension {
                         )
                     } else {
                         val imperativeBlock = findTrailingLambda(initializerExpr) ?: continue
-                        val helperSymbol = ctx.buildImperativeVisitorSymbol ?: continue
-                        buildImperativeRecipeClass(
-                            ctx = ctx,
-                            parentFile = file,
-                            propertyName = declaration.name,
-                            metadata = metadata,
-                            recipeBlock = imperativeBlock,
-                            helperSymbol = helperSymbol,
-                        )
+                        // A scan/generate recipe must extend ScanningRecipe or the
+                        // scan/generate phases are silently dropped; edit-only keeps
+                        // the lighter Recipe + getVisitor synthesis.
+                        if (recipeBlockDeclaresScanOrGenerate(imperativeBlock)) {
+                            if (!ctx.canSynthesizeScanningRecipe()) continue
+                            buildScanningRecipeClass(
+                                ctx = ctx,
+                                parentFile = file,
+                                propertyName = declaration.name,
+                                metadata = metadata,
+                                recipeBlock = imperativeBlock,
+                            )
+                        } else {
+                            val helperSymbol = ctx.buildImperativeVisitorSymbol ?: continue
+                            buildImperativeRecipeClass(
+                                ctx = ctx,
+                                parentFile = file,
+                                propertyName = declaration.name,
+                                metadata = metadata,
+                                recipeBlock = imperativeBlock,
+                                helperSymbol = helperSymbol,
+                            )
+                        }
                     }
                 }
                 RECIPES_FQN -> {
@@ -1845,6 +1929,160 @@ internal class RecipeIrGenerationExtension : IrGenerationExtension {
         }
     }
 
+    // ------------------------------------------------------------------
+    // Scanning-shape class synthesis (for `scan { }` / `generate { }` recipes).
+    // ------------------------------------------------------------------
+
+    /**
+     * True when the recipe trailing lambda declares a `scan { }` or bare
+     * `generate { }` phase. Inspects only the block's top-level statements (never
+     * user code inside phase lambdas), walking each dispatch-receiver chain to a
+     * `RecipeBuilder.scan` / `.generate` head.
+     */
+    private fun recipeBlockDeclaresScanOrGenerate(recipeBlock: IrFunctionExpression): Boolean {
+        val stmts = (recipeBlock.function.body as? IrBlockBody)?.statements ?: return false
+        for (stmt in stmts) {
+            // `Scan.generate` returns Scan<A> (non-Unit), so K2 wraps the
+            // statement in IMPLICIT_COERCION_TO_UNIT; peel it to reach the call.
+            val rawStmt = (stmt as? IrReturn)?.value ?: stmt
+            val unwrapped = if (rawStmt is IrTypeOperatorCall &&
+                rawStmt.operator == IrTypeOperator.IMPLICIT_COERCION_TO_UNIT
+            ) {
+                rawStmt.argument
+            } else {
+                rawStmt
+            }
+            var cursor: IrCall? = unwrapped as? IrCall
+            while (cursor != null) {
+                when (cursor.symbol.owner.kotlinFqName.asString()) {
+                    RECIPE_BUILDER_SCAN_FQN, RECIPE_BUILDER_GENERATE_FQN -> return true
+                }
+                cursor = cursor.dispatchReceiver as? IrCall
+            }
+        }
+        return false
+    }
+
+    /**
+     * Synthesize a field-less `<Name>$KtRecipe` extending `ScanningRecipe<Any>`
+     * for a scan/generate recipe, overriding all four lifecycle methods to
+     * delegate to the `org.openrewrite.buildImperative*` runtime helpers. Unlike
+     * [buildImperativeRecipeClass] (plain `Recipe`, `getVisitor()` only), this
+     * preserves every phase.
+     */
+    private fun buildScanningRecipeClass(
+        ctx: RecipeIrGenContext,
+        parentFile: IrFile,
+        propertyName: Name,
+        metadata: RecipeMetadata,
+        recipeBlock: IrFunctionExpression,
+    ): IrClass {
+        val anyType = ctx.pluginContext.irBuiltIns.anyType
+        val cls = ctx.pluginContext.irFactory.buildClass {
+            name = Name.identifier("${propertyName.asString()}\$KtRecipe")
+            kind = ClassKind.CLASS
+            modality = Modality.FINAL
+            visibility = DescriptorVisibilities.PUBLIC
+        }
+        cls.parent = parentFile
+        cls.createThisReceiverParameter()
+        cls.superTypes = listOf(ctx.scanningRecipeClassSymbol!!.typeWith(anyType))
+
+        cls.addConstructor {
+            isPrimary = true
+            returnType = cls.symbol.defaultType
+            visibility = DescriptorVisibilities.PUBLIC
+        }.apply {
+            body = DeclarationIrBuilder(ctx.pluginContext, symbol).irBlockBody {
+                +irDelegatingConstructorCall(ctx.scanningRecipeNoArgCtorSymbol!!.owner)
+            }
+        }
+
+        addMetadataOverrides(cls, ctx, metadata)
+        addScanningLifecycleOverrides(cls, ctx, recipeBlock, anyType)
+        return cls
+    }
+
+    private fun addScanningLifecycleOverrides(
+        cls: IrClass,
+        ctx: RecipeIrGenContext,
+        recipeBlock: IrFunctionExpression,
+        anyType: IrType,
+    ) {
+        val getInitialValue = ctx.scanningGetInitialValue!!
+        val getScanner = ctx.scanningGetScanner!!
+        val getVisitor = ctx.scanningGetVisitor!!
+        val generate = ctx.scanningGenerate!!
+
+        cls.addFunction(
+            name = "getInitialValue",
+            returnType = anyType,
+            modality = Modality.OPEN,
+            visibility = DescriptorVisibilities.PUBLIC,
+        ).apply {
+            val ctxParam = addValueParameter("ctx", getInitialValue.valueParameters[0].type)
+            overriddenSymbols = listOf(getInitialValue.symbol)
+            body = DeclarationIrBuilder(ctx.pluginContext, symbol).irBlockBody {
+                val call = irCall(ctx.buildImperativeInitialValueSymbol!!, type = anyType)
+                call.arguments[0] = irGet(ctxParam)
+                call.arguments[1] = recipeBlock.deepCopyWithSymbols(initialParent = this@apply)
+                +irReturn(call)
+            }
+        }
+
+        cls.addFunction(
+            name = "getScanner",
+            returnType = getScanner.returnType,
+            modality = Modality.OPEN,
+            visibility = DescriptorVisibilities.PUBLIC,
+        ).apply {
+            val accParam = addValueParameter("acc", anyType)
+            overriddenSymbols = listOf(getScanner.symbol)
+            body = DeclarationIrBuilder(ctx.pluginContext, symbol).irBlockBody {
+                val call = irCall(ctx.buildImperativeScannerSymbol!!, type = getScanner.returnType)
+                call.arguments[0] = irGet(accParam)
+                call.arguments[1] = recipeBlock.deepCopyWithSymbols(initialParent = this@apply)
+                +irReturn(call)
+            }
+        }
+
+        cls.addFunction(
+            name = "getVisitor",
+            returnType = getVisitor.returnType,
+            modality = Modality.OPEN,
+            visibility = DescriptorVisibilities.PUBLIC,
+        ).apply {
+            val accParam = addValueParameter("acc", anyType)
+            overriddenSymbols = listOf(getVisitor.symbol)
+            body = DeclarationIrBuilder(ctx.pluginContext, symbol).irBlockBody {
+                val call = irCall(ctx.buildImperativeEditVisitorSymbol!!, type = getVisitor.returnType)
+                call.arguments[0] = irGet(accParam)
+                call.arguments[1] = recipeBlock.deepCopyWithSymbols(initialParent = this@apply)
+                +irReturn(call)
+            }
+        }
+
+        cls.addFunction(
+            name = "generate",
+            returnType = generate.returnType,
+            modality = Modality.OPEN,
+            visibility = DescriptorVisibilities.PUBLIC,
+        ).apply {
+            val accParam = addValueParameter("acc", anyType)
+            val generatedParam = addValueParameter("generatedInThisCycle", generate.valueParameters[1].type)
+            val ctxParam = addValueParameter("ctx", generate.valueParameters[2].type)
+            overriddenSymbols = listOf(generate.symbol)
+            body = DeclarationIrBuilder(ctx.pluginContext, symbol).irBlockBody {
+                val call = irCall(ctx.buildImperativeGenerateSymbol!!, type = generate.returnType)
+                call.arguments[0] = irGet(accParam)
+                call.arguments[1] = irGet(generatedParam)
+                call.arguments[2] = irGet(ctxParam)
+                call.arguments[3] = recipeBlock.deepCopyWithSymbols(initialParent = this@apply)
+                +irReturn(call)
+            }
+        }
+    }
+
     // ------------------------------------------------------------------
     // Composite-shape class synthesis (for `recipes(...)` properties).
     // ------------------------------------------------------------------