[FIR2IR] Refactor signatures of Fir2IrConversionScope.withXxx functions

Usually functions like `withSomething` are considered to return the value
  of lambda which is passed to it, not one of its arguments

compiler/fir/fir2ir/src/org/jetbrains/kotlin/fir/backend/Fir2IrConversionScope.kt

@@ -20,12 +20,13 @@ class Fir2IrConversionScope {
     private val containingFirClassStack = mutableListOf<FirClass>()
     private val currentlyGeneratedDelegatedConstructors = mutableMapOf<IrClass, IrConstructor>()
 
-    fun <T : IrDeclarationParent?> withParent(parent: T, f: T.() -> Unit): T {
-        if (parent == null) return parent
+    fun <T : IrDeclarationParent, R> withParent(parent: T, f: T.() -> R): R {
         parentStack += parent
-        parent.f()
-        parentStack.removeAt(parentStack.size - 1)
-        return parent
+        try {
+            return parent.f()
+        } finally {
+            parentStack.removeAt(parentStack.size - 1)
+        }
     }
 
     fun <T> forDelegatingConstructorCall(constructor: IrConstructor, irClass: IrClass, f: () -> T): T {
@@ -67,29 +68,35 @@ class Fir2IrConversionScope {
 
     private val functionStack = mutableListOf<IrFunction>()
 
-    fun <T : IrFunction> withFunction(function: T, f: T.() -> Unit): T {
+    fun <T : IrFunction, R> withFunction(function: T, f: T.() -> R): R {
         functionStack += function
-        function.f()
-        functionStack.removeAt(functionStack.size - 1)
-        return function
+        try {
+            return function.f()
+        } finally {
+            functionStack.removeAt(functionStack.size - 1)
+        }
     }
 
     private val propertyStack = mutableListOf<Pair<IrProperty, FirProperty?>>()
 
-    fun withProperty(property: IrProperty, firProperty: FirProperty? = null, f: IrProperty.() -> Unit): IrProperty {
+    fun <R> withProperty(property: IrProperty, firProperty: FirProperty? = null, f: IrProperty.() -> R): R {
         propertyStack += (property to firProperty)
-        property.f()
-        propertyStack.removeAt(propertyStack.size - 1)
-        return property
+        try {
+            return property.f()
+        } finally {
+            propertyStack.removeAt(propertyStack.size - 1)
+        }
     }
 
     private val classStack = mutableListOf<IrClass>()
 
-    fun withClass(klass: IrClass, f: IrClass.() -> Unit): IrClass {
+    fun <R> withClass(klass: IrClass, f: IrClass.() -> R): R {
         classStack += klass
-        klass.f()
-        classStack.removeAt(classStack.size - 1)
-        return klass
+        return try {
+            klass.f()
+        } finally {
+            classStack.removeAt(classStack.size - 1)
+        }
     }
 
     private val whenSubjectVariableStack = mutableListOf<IrVariable>()

compiler/fir/fir2ir/src/org/jetbrains/kotlin/fir/backend/Fir2IrVisitor.kt

@@ -105,13 +105,15 @@ class Fir2IrVisitor(
     }
 
     override fun visitFile(file: FirFile, data: Any?): IrFile {
-        return conversionScope.withParent(declarationStorage.getIrFile(file)) {
+        val irFile = declarationStorage.getIrFile(file)
+        conversionScope.withParent(irFile) {
             file.declarations.forEach {
                 it.toIrDeclaration()
             }
             annotationGenerator.generate(this, file)
             metadata = FirMetadataSource.File(listOf(file))
         }
+        return irFile
     }
 
     private fun FirDeclaration.toIrDeclaration(): IrDeclaration =
@@ -188,9 +190,10 @@ class Fir2IrVisitor(
             // NB: for implicit types it is possible that local class is already cached
             val irClass = classifierStorage.getCachedIrClass(regularClass)?.apply { this.parent = irParent }
             if (irClass != null) {
-                return conversionScope.withParent(irClass) {
+                conversionScope.withParent(irClass) {
                     memberGenerator.convertClassContent(irClass, regularClass)
                 }
+                return irClass
             }
             converter.processLocalClassAndNestedClasses(regularClass, irParent)
         }
@@ -198,9 +201,10 @@ class Fir2IrVisitor(
         if (regularClass.isSealed) {
             irClass.sealedSubclasses = regularClass.getIrSymbolsForSealedSubclasses()
         }
-        return conversionScope.withParent(irClass) {
+        conversionScope.withParent(irClass) {
             memberGenerator.convertClassContent(irClass, regularClass)
         }
+        return irClass
     }
 
     override fun visitScript(script: FirScript, data: Any?): IrElement {

compiler/fir/fir2ir/src/org/jetbrains/kotlin/fir/backend/generators/FakeOverrideGenerator.kt

@@ -35,12 +35,12 @@ class FakeOverrideGenerator(
     private val basePropertySymbols: MutableMap<IrProperty, List<FirPropertySymbol>> = mutableMapOf()
     private val baseStaticFieldSymbols: MutableMap<IrField, List<FirFieldSymbol>> = mutableMapOf()
 
-    private fun IrSimpleFunction.withFunction(f: IrSimpleFunction.() -> Unit): IrSimpleFunction {
-        return conversionScope.withFunction(this, f)
+    private fun IrSimpleFunction.withFunction(f: IrSimpleFunction.() -> Unit) {
+        conversionScope.withFunction(this, f)
     }
 
-    private fun IrProperty.withProperty(f: IrProperty.() -> Unit): IrProperty {
-        return conversionScope.withProperty(this, firProperty = null, f)
+    private fun IrProperty.withProperty(f: IrProperty.() -> Unit) {
+        conversionScope.withProperty(this, firProperty = null, f)
     }
 
     private fun IrType.containsErrorType(): Boolean {