[FIR] Add opt-in to replaceIsTrailingLambda

#KT-66553

compiler/fir/raw-fir/light-tree2fir/src/org/jetbrains/kotlin/fir/lightTree/converter/LightTreeRawFirExpressionBuilder.kt

@@ -1515,6 +1515,8 @@ class LightTreeRawFirExpressionBuilder(
                 LABELED_EXPRESSION,
                 ANNOTATED_EXPRESSION,
                 -> container += getAsFirExpression<FirAnonymousFunctionExpression>(node).apply {
+                    // TODO(KT-66553) remove and set in builder
+                    @OptIn(RawFirApi::class)
                     replaceIsTrailingLambda(newIsTrailingLambda = true)
                 }
             }

compiler/fir/raw-fir/psi2fir/src/org/jetbrains/kotlin/fir/builder/PsiRawFirBuilder.kt

@@ -905,6 +905,8 @@ open class PsiRawFirBuilder(
                         buildOrLazyExpression((argument as? PsiElement)?.toFirSourceElement()) { argument.toFirExpression() }
                     arguments += when (argument) {
                         is KtLambdaArgument -> argumentExpression.apply {
+                            // TODO(KT-66553) remove and set in builder
+                            @OptIn(RawFirApi::class)
                             (this as? FirAnonymousFunctionExpression)?.replaceIsTrailingLambda(true)
                         }
                         else -> argumentExpression

compiler/fir/tree/gen/org/jetbrains/kotlin/fir/expressions/FirAnonymousFunctionExpression.kt

@@ -37,6 +37,7 @@ abstract class FirAnonymousFunctionExpression : FirExpression() {
 
     abstract override fun replaceAnnotations(newAnnotations: List<FirAnnotation>)
 
+    @RawFirApi
     abstract fun replaceIsTrailingLambda(newIsTrailingLambda: Boolean)
 
     abstract override fun <D> transformAnnotations(transformer: FirTransformer<D>, data: D): FirAnonymousFunctionExpression

compiler/fir/tree/src/org/jetbrains/kotlin/fir/expressions/RawFirApi.kt

@@ -0,0 +1,9 @@
+/*
+ * Copyright 2010-2024 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.fir.expressions
+
+@RequiresOptIn("API is only for raw FIR building.")
+annotation class RawFirApi

compiler/fir/tree/src/org/jetbrains/kotlin/fir/expressions/impl/FirAnonymousFunctionExpressionImpl.kt

@@ -9,6 +9,7 @@ import org.jetbrains.kotlin.KtSourceElement
 import org.jetbrains.kotlin.fir.declarations.FirAnonymousFunction
 import org.jetbrains.kotlin.fir.expressions.FirAnnotation
 import org.jetbrains.kotlin.fir.expressions.FirAnonymousFunctionExpression
+import org.jetbrains.kotlin.fir.expressions.RawFirApi
 import org.jetbrains.kotlin.fir.expressions.UnresolvedExpressionTypeAccess
 import org.jetbrains.kotlin.fir.types.ConeKotlinType
 import org.jetbrains.kotlin.fir.types.coneTypeOrNull
@@ -42,6 +43,7 @@ internal class FirAnonymousFunctionExpressionImpl(
         anonymousFunction.replaceAnnotations(newAnnotations)
     }
 
+    @RawFirApi
     override fun replaceIsTrailingLambda(newIsTrailingLambda: Boolean) {
         isTrailingLambda = newIsTrailingLambda
     }

compiler/fir/tree/tree-generator/src/org/jetbrains/kotlin/fir/tree/generator/NodeConfigurator.kt

@@ -332,7 +332,9 @@ object NodeConfigurator : AbstractFieldConfigurator<FirTreeBuilder>(FirTreeBuild
 
         anonymousFunctionExpression.configure {
             +field(anonymousFunction).withTransform()
-            +booleanField("isTrailingLambda", withReplace = true)
+            +booleanField("isTrailingLambda", withReplace = true).apply {
+                replaceOptInAnnotation = rawFirApi
+            }
         }
 
         typeParameter.configure {

compiler/fir/tree/tree-generator/src/org/jetbrains/kotlin/fir/tree/generator/Types.kt

@@ -124,4 +124,5 @@ val firTransformerType = generatedType("visitors", "FirTransformer")
 
 val resolveStateAccessAnnotation = type("fir.declarations", "ResolveStateAccess", kind = TypeKind.Class)
 val unresolvedExpressionTypeAccessAnnotation = type("fir.expressions", "UnresolvedExpressionTypeAccess", kind = TypeKind.Class)
+val rawFirApi = type("fir.expressions", "RawFirApi", kind = TypeKind.Class)
 val firBuilderDslAnnotation = type("fir.builder", "FirBuilderDsl", kind = TypeKind.Class)

compiler/fir/tree/tree-generator/src/org/jetbrains/kotlin/fir/tree/generator/model/Field.kt

@@ -103,6 +103,10 @@ class FieldWithDefault(override val origin: Field) : Field(), AbstractFieldWithD
         get() = origin.optInAnnotation
         set(_) {}
 
+    override var replaceOptInAnnotation: ClassRef<*>?
+        get() = origin.replaceOptInAnnotation
+        set(_) {}
+
     override var defaultValueInImplementation: String? = origin.defaultValueInImplementation
     override var defaultValueInBuilder: String? = null
     override var isMutable: Boolean = origin.isMutable

compiler/fir/tree/tree-generator/src/org/jetbrains/kotlin/fir/tree/generator/printer/utils.kt

@@ -78,6 +78,7 @@ fun SmartPrinter.replaceFunctionDeclaration(
             implementationKind == ImplementationKind.AbstractClass || implementationKind == ImplementationKind.SealedClass
         },
         override = override,
+        optInAnnotation = field.replaceOptInAnnotation,
     )
 }
 

generators/tree-generator-common/src/org/jetbrains/kotlin/generators/tree/AbstractField.kt

@@ -27,6 +27,7 @@ abstract class AbstractField<Field : AbstractField<Field>> {
     open val arbitraryImportables: MutableList<Importable> = mutableListOf()
 
     open var optInAnnotation: ClassRef<*>? = null
+    open var replaceOptInAnnotation: ClassRef<*>? = null
 
     abstract var isMutable: Boolean
     open val withGetter: Boolean get() = false
@@ -102,6 +103,7 @@ abstract class AbstractField<Field : AbstractField<Field>> {
         copy.isLateinit = isLateinit
         copy.arbitraryImportables += arbitraryImportables
         copy.optInAnnotation = optInAnnotation
+        copy.replaceOptInAnnotation = replaceOptInAnnotation
         copy.isMutable = isMutable
         copy.deprecation = deprecation
         copy.visibility = visibility

generators/tree-generator-common/src/org/jetbrains/kotlin/generators/tree/printer/printUtils.kt

@@ -117,7 +117,12 @@ fun SmartPrinter.printFunctionDeclaration(
     override: Boolean = false,
     isInline: Boolean = false,
     allParametersOnSeparateLines: Boolean = false,
+    optInAnnotation: ClassRef<*>? = null,
 ) {
+    optInAnnotation?.let {
+        println("@", it.render())
+    }
+
     if (visibility != Visibility.PUBLIC) {
         print(visibility.name.toLowerCaseAsciiOnly(), " ")
     }