JVM IR: minor, refactor replaceThisByStaticReference

Allow to replace "this" references not only inside IrBody, but inside
any IrElement. This will be useful in JvmStatic in object lowering where
we want not only to replace usages inside the function body, but also in
default values of its parameters.

compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/JvmCachedDeclarations.kt

@@ -24,7 +24,6 @@ import org.jetbrains.kotlin.ir.builders.declarations.buildFun
 import org.jetbrains.kotlin.ir.builders.irCall
 import org.jetbrains.kotlin.ir.builders.setSourceRange
 import org.jetbrains.kotlin.ir.declarations.*
-import org.jetbrains.kotlin.ir.expressions.IrExpressionBody
 import org.jetbrains.kotlin.ir.util.*
 import org.jetbrains.kotlin.load.java.JavaDescriptorVisibilities
 import org.jetbrains.kotlin.load.java.JvmAbi
@@ -135,9 +134,8 @@ class JvmCachedDeclarations(
                     annotations = oldField.annotations
                 }
 
-                initializer = oldField.initializer
-                    ?.replaceThisByStaticReference(this@JvmCachedDeclarations, oldParent, oldParent.thisReceiver!!)
-                    ?.patchDeclarationParents(this) as IrExpressionBody?
+                initializer = oldField.initializer?.patchDeclarationParents(this)
+                oldField.replaceThisByStaticReference(this@JvmCachedDeclarations, oldParent, oldParent.thisReceiver!!)
                 origin = if (irProperty.parentAsClass.isCompanion) JvmLoweredDeclarationOrigin.COMPANION_PROPERTY_BACKING_FIELD else origin
             }
         }

compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/ir/IrUtils.kt

@@ -20,6 +20,7 @@ import org.jetbrains.kotlin.descriptors.ClassKind
 import org.jetbrains.kotlin.descriptors.DescriptorVisibilities
 import org.jetbrains.kotlin.descriptors.Modality
 import org.jetbrains.kotlin.descriptors.deserialization.PLATFORM_DEPENDENT_ANNOTATION_FQ_NAME
+import org.jetbrains.kotlin.ir.IrElement
 import org.jetbrains.kotlin.ir.UNDEFINED_OFFSET
 import org.jetbrains.kotlin.ir.builders.IrBuilderWithScope
 import org.jetbrains.kotlin.ir.builders.Scope
@@ -39,6 +40,7 @@ import org.jetbrains.kotlin.ir.types.impl.IrStarProjectionImpl
 import org.jetbrains.kotlin.ir.types.impl.makeTypeProjection
 import org.jetbrains.kotlin.ir.util.*
 import org.jetbrains.kotlin.ir.visitors.IrElementTransformerVoid
+import org.jetbrains.kotlin.ir.visitors.transformChildrenVoid
 import org.jetbrains.kotlin.load.java.JavaDescriptorVisibilities
 import org.jetbrains.kotlin.load.java.JvmAbi
 import org.jetbrains.kotlin.metadata.jvm.deserialization.JvmProtoBufUtil
@@ -223,25 +225,23 @@ fun IrExpression.isSmartcastFromHigherThanNullable(context: JvmBackendContext):
     }
 }
 
-fun IrBody.replaceThisByStaticReference(
+fun IrElement.replaceThisByStaticReference(
     cachedDeclarations: JvmCachedDeclarations,
     irClass: IrClass,
     oldThisReceiverParameter: IrValueParameter
-): IrBody =
-    transform(object : IrElementTransformerVoid() {
-        override fun visitGetValue(expression: IrGetValue): IrExpression {
+) {
+    transformChildrenVoid(object : IrElementTransformerVoid() {
+        override fun visitGetValue(expression: IrGetValue): IrExpression =
             if (expression.symbol == oldThisReceiverParameter.symbol) {
-                val instanceField = cachedDeclarations.getPrivateFieldForObjectInstance(irClass)
-                return IrGetFieldImpl(
+                IrGetFieldImpl(
                     expression.startOffset,
                     expression.endOffset,
-                    instanceField.symbol,
+                    cachedDeclarations.getPrivateFieldForObjectInstance(irClass).symbol,
                     irClass.defaultType
                 )
-            }
-            return super.visitGetValue(expression)
-        }
-    }, null)
+            } else super.visitGetValue(expression)
+    })
+}
 
 // TODO: Interface Parameters
 //

compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/lower/JvmStaticAnnotationLowering.kt

@@ -139,11 +139,8 @@ private class SingletonObjectJvmStaticLowering(val context: JvmBackendContext) :
             val replacement = createReplacement(context, function as IrSimpleFunction)
             // Set dispatch receiver parameter for body move operation.
             replacement.dispatchReceiverParameter = function.dispatchReceiverParameter
-            replacement.body = function.moveBodyTo(replacement)?.replaceThisByStaticReference(
-                context.cachedDeclarations,
-                irClass,
-                function.dispatchReceiverParameter!!
-            )
+            replacement.body = function.moveBodyTo(replacement)
+            replacement.replaceThisByStaticReference(context.cachedDeclarations, irClass, function.dispatchReceiverParameter!!)
             // Clear dispatch receiver parameter again after body move operation.
             replacement.dispatchReceiverParameter = null
             irClass.declarations.remove(function)

compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/lower/MoveCompanionObjectFieldsLowering.kt

@@ -7,7 +7,6 @@ package org.jetbrains.kotlin.backend.jvm.lower
 
 import org.jetbrains.kotlin.backend.common.ClassLoweringPass
 import org.jetbrains.kotlin.backend.common.FileLoweringPass
-import org.jetbrains.kotlin.backend.common.lower.createIrBuilder
 import org.jetbrains.kotlin.backend.common.phaser.makeIrFilePhase
 import org.jetbrains.kotlin.backend.jvm.JvmBackendContext
 import org.jetbrains.kotlin.backend.jvm.ir.createJvmIrBuilder
@@ -105,9 +104,8 @@ private class MoveOrCopyCompanionObjectFieldsLowering(val context: JvmBackendCon
             val newSymbol = IrAnonymousInitializerSymbolImpl(newParent.symbol)
             IrAnonymousInitializerImpl(startOffset, endOffset, origin, newSymbol, isStatic = true).apply {
                 parent = newParent
-                body = this@with.body
-                    .replaceThisByStaticReference(context.cachedDeclarations, oldParent, oldParent.thisReceiver!!)
-                    .patchDeclarationParents(newParent) as IrBlockBody
+                body = this@with.body.patchDeclarationParents(newParent)
+                replaceThisByStaticReference(context.cachedDeclarations, oldParent, oldParent.thisReceiver!!)
             }
         }