IR: IrFunctionSymbol.owner: IrFunction

This requires some class hierarchy juggling around
IrFunction and IrReturnableBlock.

compiler/ir/ir.tree/src/org/jetbrains/kotlin/ir/declarations/IrFunction.kt

@@ -25,7 +25,7 @@ import org.jetbrains.kotlin.ir.symbols.IrFunctionSymbol
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.types.KotlinType
 
-interface IrFunction : IrDeclaration, IrTypeParametersContainer, IrSymbolOwner, IrDeclarationParent {
+interface IrFunction : IrDeclaration, IrTypeParametersContainer, IrSymbolOwner, IrDeclarationParent, IrReturnTarget {
     override val descriptor: FunctionDescriptor
     override val symbol: IrFunctionSymbol
 

compiler/ir/ir.tree/src/org/jetbrains/kotlin/ir/declarations/IrReturnTarget.kt

@@ -0,0 +1,15 @@
+/*
+ * Copyright 2010-2018 JetBrains s.r.o. 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.ir.declarations
+
+import org.jetbrains.kotlin.descriptors.FunctionDescriptor
+import org.jetbrains.kotlin.ir.IrElement
+import org.jetbrains.kotlin.ir.symbols.IrReturnTargetSymbol
+
+interface IrReturnTarget : IrSymbolOwner {
+    val descriptor: FunctionDescriptor
+    override val symbol: IrReturnTargetSymbol
+}

compiler/ir/ir.tree/src/org/jetbrains/kotlin/ir/expressions/IrBlock.kt

@@ -16,7 +16,7 @@
 
 package org.jetbrains.kotlin.ir.expressions
 
-import org.jetbrains.kotlin.descriptors.FunctionDescriptor
+import org.jetbrains.kotlin.ir.declarations.IrReturnTarget
 import org.jetbrains.kotlin.ir.declarations.IrSymbolOwner
 import org.jetbrains.kotlin.ir.symbols.IrReturnableBlockSymbol
 
@@ -35,8 +35,7 @@ interface IrComposite : IrContainerExpression {
         get() = true
 }
 
-interface IrReturnableBlock : IrBlock, IrSymbolOwner {
+interface IrReturnableBlock : IrBlock, IrSymbolOwner, IrReturnTarget {
     override val symbol: IrReturnableBlockSymbol
-    val descriptor: FunctionDescriptor
     val sourceFileName: String
 }

compiler/ir/ir.tree/src/org/jetbrains/kotlin/ir/expressions/IrReturn.kt

@@ -17,12 +17,12 @@
 package org.jetbrains.kotlin.ir.expressions
 
 import org.jetbrains.kotlin.descriptors.FunctionDescriptor
-import org.jetbrains.kotlin.ir.symbols.IrFunctionSymbol
+import org.jetbrains.kotlin.ir.symbols.IrReturnTargetSymbol
 
 
 interface IrReturn : IrExpression {
     var value: IrExpression
     val returnTarget: FunctionDescriptor
-    val returnTargetSymbol: IrFunctionSymbol
+    val returnTargetSymbol: IrReturnTargetSymbol
 }
 

compiler/ir/ir.tree/src/org/jetbrains/kotlin/ir/expressions/impl/IrReturnImpl.kt

@@ -20,6 +20,7 @@ import org.jetbrains.kotlin.descriptors.FunctionDescriptor
 import org.jetbrains.kotlin.ir.expressions.IrExpression
 import org.jetbrains.kotlin.ir.expressions.IrReturn
 import org.jetbrains.kotlin.ir.symbols.IrFunctionSymbol
+import org.jetbrains.kotlin.ir.symbols.IrReturnTargetSymbol
 import org.jetbrains.kotlin.ir.symbols.impl.createFunctionSymbol
 import org.jetbrains.kotlin.ir.visitors.IrElementTransformer
 import org.jetbrains.kotlin.ir.visitors.IrElementVisitor
@@ -30,7 +31,7 @@ class IrReturnImpl(
     startOffset: Int,
     endOffset: Int,
     type: KotlinType,
-    override val returnTargetSymbol: IrFunctionSymbol,
+    override val returnTargetSymbol: IrReturnTargetSymbol,
     override var value: IrExpression
 ) : IrExpressionBase(startOffset, endOffset, type), IrReturn {
     constructor(startOffset: Int, endOffset: Int, returnTargetSymbol: IrFunctionSymbol, value: IrExpression) :

compiler/ir/ir.tree/src/org/jetbrains/kotlin/ir/symbols/IrSymbol.kt

@@ -59,11 +59,16 @@ interface IrValueSymbol : IrSymbol {
 interface IrValueParameterSymbol : IrValueSymbol, IrBindableSymbol<ParameterDescriptor, IrValueParameter>
 interface IrVariableSymbol : IrValueSymbol, IrBindableSymbol<VariableDescriptor, IrVariable>
 
-interface IrFunctionSymbol : IrSymbol {
+interface IrReturnTargetSymbol : IrSymbol {
     override val descriptor: FunctionDescriptor
+    override val owner: IrReturnTarget
+}
+
+interface IrFunctionSymbol : IrReturnTargetSymbol {
+    override val owner: IrFunction
 }
 
 interface IrConstructorSymbol : IrFunctionSymbol, IrBindableSymbol<ClassConstructorDescriptor, IrConstructor>
 interface IrSimpleFunctionSymbol : IrFunctionSymbol, IrBindableSymbol<FunctionDescriptor, IrSimpleFunction>
 
-interface IrReturnableBlockSymbol : IrFunctionSymbol, IrBindableSymbol<FunctionDescriptor, IrReturnableBlock>
+interface IrReturnableBlockSymbol : IrReturnTargetSymbol, IrBindableSymbol<FunctionDescriptor, IrReturnableBlock>

compiler/ir/ir.tree/src/org/jetbrains/kotlin/ir/util/DeepCopyIrTreeWithSymbols.kt

@@ -22,6 +22,9 @@ import org.jetbrains.kotlin.ir.declarations.*
 import org.jetbrains.kotlin.ir.declarations.impl.*
 import org.jetbrains.kotlin.ir.expressions.*
 import org.jetbrains.kotlin.ir.expressions.impl.*
+import org.jetbrains.kotlin.ir.symbols.IrFunctionSymbol
+import org.jetbrains.kotlin.ir.symbols.IrReturnTargetSymbol
+import org.jetbrains.kotlin.ir.symbols.IrReturnableBlockSymbol
 import org.jetbrains.kotlin.ir.symbols.IrSimpleFunctionSymbol
 import org.jetbrains.kotlin.ir.symbols.impl.IrAnonymousInitializerSymbolImpl
 import org.jetbrains.kotlin.ir.visitors.IrElementTransformerVoid
@@ -550,10 +553,17 @@ open class DeepCopyIrTreeWithSymbols(private val symbolRemapper: SymbolRemapper)
         IrReturnImpl(
             expression.startOffset, expression.endOffset,
             expression.type,
-            symbolRemapper.getReferencedFunction(expression.returnTargetSymbol),
+            symbolRemapper.getReferencedReturnTarget(expression.returnTargetSymbol),
             expression.value.transform()
         )
 
+    private fun SymbolRemapper.getReferencedReturnTarget(returnTarget: IrReturnTargetSymbol) =
+        when (returnTarget) {
+            is IrFunctionSymbol -> getReferencedFunction(returnTarget)
+            is IrReturnableBlockSymbol -> getReferencedReturnableBlock(returnTarget)
+            else -> throw AssertionError("Unexpected return target: ${returnTarget.javaClass} $returnTarget")
+        }
+
     override fun visitThrow(expression: IrThrow): IrThrow =
         IrThrowImpl(
             expression.startOffset, expression.endOffset,

compiler/ir/ir.tree/src/org/jetbrains/kotlin/ir/util/DeepCopySymbolsRemapper.kt

@@ -19,6 +19,8 @@ package org.jetbrains.kotlin.ir.util
 import org.jetbrains.kotlin.descriptors.DeclarationDescriptor
 import org.jetbrains.kotlin.ir.IrElement
 import org.jetbrains.kotlin.ir.declarations.*
+import org.jetbrains.kotlin.ir.expressions.IrBlock
+import org.jetbrains.kotlin.ir.expressions.IrReturnableBlock
 import org.jetbrains.kotlin.ir.symbols.*
 import org.jetbrains.kotlin.ir.symbols.impl.*
 import org.jetbrains.kotlin.ir.visitors.IrElementVisitorVoid
@@ -34,6 +36,7 @@ open class DeepCopySymbolsRemapper(
     private val fields = hashMapOf<IrFieldSymbol, IrFieldSymbol>()
     private val files = hashMapOf<IrFileSymbol, IrFileSymbol>()
     private val functions = hashMapOf<IrSimpleFunctionSymbol, IrSimpleFunctionSymbol>()
+    private val returnableBlocks = hashMapOf<IrReturnableBlockSymbol, IrReturnableBlockSymbol>()
     private val typeParameters = hashMapOf<IrTypeParameterSymbol, IrTypeParameterSymbol>()
     private val valueParameters = hashMapOf<IrValueParameterSymbol, IrValueParameterSymbol>()
     private val variables = hashMapOf<IrVariableSymbol, IrVariableSymbol>()
@@ -118,6 +121,15 @@ open class DeepCopySymbolsRemapper(
         declaration.acceptChildrenVoid(this)
     }
 
+    override fun visitBlock(expression: IrBlock) {
+        if (expression is IrReturnableBlock) {
+            remapSymbol(returnableBlocks, expression) {
+                IrReturnableBlockSymbolImpl(expression.descriptor)
+            }
+        }
+        expression.acceptChildrenVoid(this)
+    }
+
     private fun <T : IrSymbol> Map<T, T>.getDeclared(symbol: T) =
         getOrElse(symbol) {
             throw IllegalArgumentException("Non-remapped symbol $symbol ${symbol.descriptor}")
@@ -145,7 +157,6 @@ open class DeepCopySymbolsRemapper(
     override fun getReferencedVariable(symbol: IrVariableSymbol): IrVariableSymbol = variables.getReferenced(symbol)
     override fun getReferencedField(symbol: IrFieldSymbol): IrFieldSymbol = fields.getReferenced(symbol)
     override fun getReferencedConstructor(symbol: IrConstructorSymbol): IrConstructorSymbol = constructors.getReferenced(symbol)
-
     override fun getReferencedValue(symbol: IrValueSymbol): IrValueSymbol =
         when (symbol) {
             is IrValueParameterSymbol -> valueParameters.getReferenced(symbol)
@@ -160,6 +171,9 @@ open class DeepCopySymbolsRemapper(
             else -> throw IllegalArgumentException("Unexpected symbol $symbol ${symbol.descriptor}")
         }
 
+    override fun getReferencedReturnableBlock(symbol: IrReturnableBlockSymbol): IrReturnableBlockSymbol =
+        returnableBlocks.getReferenced(symbol)
+
     override fun getReferencedClassifier(symbol: IrClassifierSymbol): IrClassifierSymbol =
         when (symbol) {
             is IrClassSymbol -> classes.getReferenced(symbol)

compiler/ir/ir.tree/src/org/jetbrains/kotlin/ir/util/SymbolRemapper.kt

@@ -37,5 +37,6 @@ interface SymbolRemapper {
     fun getReferencedConstructor(symbol: IrConstructorSymbol): IrConstructorSymbol
     fun getReferencedValue(symbol: IrValueSymbol): IrValueSymbol
     fun getReferencedFunction(symbol: IrFunctionSymbol): IrFunctionSymbol
+    fun getReferencedReturnableBlock(symbol: IrReturnableBlockSymbol): IrReturnableBlockSymbol
     fun getReferencedClassifier(symbol: IrClassifierSymbol): IrClassifierSymbol
 }