From 1c2fbb61fe9cef1a3d1366715d9712ea6d228227 Mon Sep 17 00:00:00 2001
From: Mikhail Glukhikh <mikhail.glukhikh@jetbrains.com>
Date: Fri, 5 Jun 2020 09:41:23 +0300
Subject: [PATCH] [IR] Introduce & use declaration-based
 SymbolTable.withReferenceScope

---
 .../src/org/jetbrains/kotlin/ir/util/SymbolTable.kt | 13 +++++++++++--
 .../compiler/backend/ir/GeneratorHelpers.kt         |  2 +-
 2 files changed, 12 insertions(+), 3 deletions(-)

diff --git a/compiler/ir/ir.tree/src/org/jetbrains/kotlin/ir/util/SymbolTable.kt b/compiler/ir/ir.tree/src/org/jetbrains/kotlin/ir/util/SymbolTable.kt
index bd1c8f735ab..ecd6d5df4ac 100644
--- a/compiler/ir/ir.tree/src/org/jetbrains/kotlin/ir/util/SymbolTable.kt
+++ b/compiler/ir/ir.tree/src/org/jetbrains/kotlin/ir/util/SymbolTable.kt
@@ -63,8 +63,10 @@ interface ReferenceSymbolTable {
     fun referenceTypeAliasFromLinker(descriptor: TypeAliasDescriptor, sig: IdSignature): IrTypeAliasSymbol
 
     fun enterScope(owner: DeclarationDescriptor)
+    fun enterScope(owner: IrDeclaration)
 
     fun leaveScope(owner: DeclarationDescriptor)
+    fun leaveScope(owner: IrDeclaration)
 }
 
 open class SymbolTable(
@@ -955,7 +957,7 @@ open class SymbolTable(
         scopedSymbolTables.forEach { it.enterScope(owner) }
     }
 
-    fun enterScope(owner: IrDeclaration) {
+    override fun enterScope(owner: IrDeclaration) {
         enterScope(owner.descriptor)
     }
 
@@ -963,7 +965,7 @@ open class SymbolTable(
         scopedSymbolTables.forEach { it.leaveScope(owner) }
     }
 
-    fun leaveScope(owner: IrDeclaration) {
+    override fun leaveScope(owner: IrDeclaration) {
         leaveScope(owner.descriptor)
     }
 
@@ -1026,6 +1028,13 @@ inline fun <T, D : DeclarationDescriptor> ReferenceSymbolTable.withReferenceScop
     return result
 }
 
+inline fun <T, D : IrDeclaration> ReferenceSymbolTable.withReferenceScope(owner: D, block: ReferenceSymbolTable.(D) -> T): T {
+    enterScope(owner)
+    val result = block(owner)
+    leaveScope(owner)
+    return result
+}
+
 val SymbolTable.allUnbound: Set<IrSymbol>
     get() {
         val r = mutableSetOf<IrSymbol>()
diff --git a/plugins/kotlin-serialization/kotlin-serialization-compiler/src/org/jetbrains/kotlinx/serialization/compiler/backend/ir/GeneratorHelpers.kt b/plugins/kotlin-serialization/kotlin-serialization-compiler/src/org/jetbrains/kotlinx/serialization/compiler/backend/ir/GeneratorHelpers.kt
index e08bc587575..6cb2292ec6c 100644
--- a/plugins/kotlin-serialization/kotlin-serialization-compiler/src/org/jetbrains/kotlinx/serialization/compiler/backend/ir/GeneratorHelpers.kt
+++ b/plugins/kotlin-serialization/kotlin-serialization-compiler/src/org/jetbrains/kotlinx/serialization/compiler/backend/ir/GeneratorHelpers.kt
@@ -136,7 +136,7 @@ interface IrBuilderExtension {
 
     fun <T : IrDeclaration> T.buildWithScope(builder: (T) -> Unit): T =
         also { irDeclaration ->
-            compilerContext.symbolTable.withReferenceScope(irDeclaration.descriptor) {
+            compilerContext.symbolTable.withReferenceScope(irDeclaration) {
                 builder(irDeclaration)
             }
         }