Additional customization for DeepCopyIrTreeWithSymbols

1. Allow overriding 'getNonTransformedLoop'
2. Allow descriptor remapping on deep copy

Author: dnpetrov

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

@@ -36,7 +36,7 @@ inline fun <reified T : IrElement> T.deepCopyWithSymbols(): T {
 }
 
 
-class DeepCopyIrTreeWithSymbols(private val symbolRemapper: SymbolRemapper) : IrElementTransformerVoid() {
+open class DeepCopyIrTreeWithSymbols(private val symbolRemapper: SymbolRemapper) : IrElementTransformerVoid() {
     private fun mapDeclarationOrigin(origin: IrDeclarationOrigin) = origin
     private fun mapStatementOrigin(origin: IrStatementOrigin?) = origin
 
@@ -455,7 +455,7 @@ class DeepCopyIrTreeWithSymbols(private val symbolRemapper: SymbolRemapper) : Ir
     private fun getTransformedLoop(irLoop: IrLoop): IrLoop =
             transformedLoops.getOrElse(irLoop) { getNonTransformedLoop(irLoop) }
 
-    private fun getNonTransformedLoop(irLoop: IrLoop): IrLoop =
+    protected open fun getNonTransformedLoop(irLoop: IrLoop): IrLoop =
             throw AssertionError("Outer loop was not transformed: ${irLoop.render()}")
 
     override fun visitWhileLoop(loop: IrWhileLoop): IrWhileLoop =

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

@@ -24,7 +24,9 @@ import org.jetbrains.kotlin.ir.symbols.impl.*
 import org.jetbrains.kotlin.ir.visitors.IrElementVisitorVoid
 import org.jetbrains.kotlin.ir.visitors.acceptChildrenVoid
 
-class DeepCopySymbolsRemapper : IrElementVisitorVoid, SymbolRemapper {
+class DeepCopySymbolsRemapper(
+        val descriptorsRemapper: DescriptorsRemapper = DescriptorsRemapper.DEFAULT
+) : IrElementVisitorVoid, SymbolRemapper {
     private val classes = hashMapOf<IrClassSymbol, IrClassSymbol>()
     private val constructors = hashMapOf<IrConstructorSymbol, IrConstructorSymbol>()
     private val enumEntries = hashMapOf<IrEnumEntrySymbol, IrEnumEntrySymbol>()
@@ -47,52 +49,72 @@ class DeepCopySymbolsRemapper : IrElementVisitorVoid, SymbolRemapper {
     }
 
     override fun visitClass(declaration: IrClass) {
-        remapSymbol(classes, declaration) { IrClassSymbolImpl(it.descriptor) }
+        remapSymbol(classes, declaration) {
+            IrClassSymbolImpl(descriptorsRemapper.remapDeclaredClass(it.descriptor))
+        }
         declaration.acceptChildrenVoid(this)
     }
 
     override fun visitConstructor(declaration: IrConstructor) {
-        remapSymbol(constructors, declaration) { IrConstructorSymbolImpl(it.descriptor) }
+        remapSymbol(constructors, declaration) {
+            IrConstructorSymbolImpl(descriptorsRemapper.remapDeclaredConstructor(it.descriptor))
+        }
         declaration.acceptChildrenVoid(this)
     }
 
     override fun visitEnumEntry(declaration: IrEnumEntry) {
-        remapSymbol(enumEntries, declaration) { IrEnumEntrySymbolImpl(it.descriptor) }
+        remapSymbol(enumEntries, declaration) {
+            IrEnumEntrySymbolImpl(descriptorsRemapper.remapDeclaredEnumEntry(it.descriptor))
+        }
         declaration.acceptChildrenVoid(this)
     }
 
     override fun visitExternalPackageFragment(declaration: IrExternalPackageFragment) {
-        remapSymbol(externalPackageFragments, declaration) { IrExternalPackageFragmentSymbolImpl(it.descriptor) }
+        remapSymbol(externalPackageFragments, declaration) {
+            IrExternalPackageFragmentSymbolImpl(descriptorsRemapper.remapDeclaredExternalPackageFragment(it.descriptor))
+        }
         declaration.acceptChildrenVoid(this)
     }
 
     override fun visitField(declaration: IrField) {
-        remapSymbol(fields, declaration) { IrFieldSymbolImpl(it.descriptor) }
+        remapSymbol(fields, declaration) {
+            IrFieldSymbolImpl(descriptorsRemapper.remapDeclaredField(it.descriptor))
+        }
         declaration.acceptChildrenVoid(this)
     }
 
     override fun visitFile(declaration: IrFile) {
-        remapSymbol(files, declaration) { IrFileSymbolImpl(it.descriptor) }
+        remapSymbol(files, declaration) {
+            IrFileSymbolImpl(descriptorsRemapper.remapDeclaredFilePackageFragment(it.descriptor))
+        }
         declaration.acceptChildrenVoid(this)
     }
 
     override fun visitSimpleFunction(declaration: IrSimpleFunction) {
-        remapSymbol(functions, declaration) { IrSimpleFunctionSymbolImpl(it.descriptor) }
+        remapSymbol(functions, declaration) {
+            IrSimpleFunctionSymbolImpl(descriptorsRemapper.remapDeclaredSimpleFunction(it.descriptor))
+        }
         declaration.acceptChildrenVoid(this)
     }
 
     override fun visitTypeParameter(declaration: IrTypeParameter) {
-        remapSymbol(typeParameters, declaration) { IrTypeParameterSymbolImpl(it.descriptor) }
+        remapSymbol(typeParameters, declaration) {
+            IrTypeParameterSymbolImpl(descriptorsRemapper.remapDeclaredTypeParameter(it.descriptor))
+        }
         declaration.acceptChildrenVoid(this)
     }
 
     override fun visitValueParameter(declaration: IrValueParameter) {
-        remapSymbol(valueParameters, declaration) { IrValueParameterSymbolImpl(it.descriptor) }
+        remapSymbol(valueParameters, declaration) {
+            IrValueParameterSymbolImpl(descriptorsRemapper.remapDeclaredValueParameter(it.descriptor))
+        }
         declaration.acceptChildrenVoid(this)
     }
 
     override fun visitVariable(declaration: IrVariable) {
-        remapSymbol(variables, declaration) { IrVariableSymbolImpl(it.descriptor) }
+        remapSymbol(variables, declaration) {
+            IrVariableSymbolImpl(descriptorsRemapper.remapDeclaredVariable(it.descriptor))
+        }
         declaration.acceptChildrenVoid(this)
     }
 

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

@@ -0,0 +1,34 @@
+/*
+ * Copyright 2010-2017 JetBrains s.r.o.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.jetbrains.kotlin.ir.util
+
+import org.jetbrains.kotlin.descriptors.*
+
+interface DescriptorsRemapper {
+    fun remapDeclaredClass(descriptor: ClassDescriptor): ClassDescriptor = descriptor
+    fun remapDeclaredConstructor(descriptor: ClassConstructorDescriptor): ClassConstructorDescriptor = descriptor
+    fun remapDeclaredEnumEntry(descriptor: ClassDescriptor): ClassDescriptor = descriptor
+    fun remapDeclaredExternalPackageFragment(descriptor: PackageFragmentDescriptor): PackageFragmentDescriptor = descriptor
+    fun remapDeclaredField(descriptor: PropertyDescriptor): PropertyDescriptor = descriptor
+    fun remapDeclaredFilePackageFragment(descriptor: PackageFragmentDescriptor): PackageFragmentDescriptor = descriptor
+    fun remapDeclaredSimpleFunction(descriptor: FunctionDescriptor): FunctionDescriptor = descriptor
+    fun remapDeclaredTypeParameter(descriptor: TypeParameterDescriptor): TypeParameterDescriptor = descriptor
+    fun remapDeclaredValueParameter(descriptor: ParameterDescriptor): ParameterDescriptor = descriptor
+    fun remapDeclaredVariable(descriptor: VariableDescriptor): VariableDescriptor = descriptor
+
+    object DEFAULT : DescriptorsRemapper
+}