[IR] Update inheritors of SymbolTable which override some descriptor-related methods

compiler/ir/backend.jvm/entrypoint/src/org/jetbrains/kotlin/backend/jvm/SymbolTableWithBuiltInsDeduplication.kt

@@ -11,6 +11,7 @@ import org.jetbrains.kotlin.ir.IrBuiltIns
 import org.jetbrains.kotlin.ir.ObsoleteDescriptorBasedAPI
 import org.jetbrains.kotlin.ir.declarations.IrFactory
 import org.jetbrains.kotlin.ir.symbols.IrClassSymbol
+import org.jetbrains.kotlin.ir.util.DescriptorSymbolTableExtension
 import org.jetbrains.kotlin.ir.util.IdSignatureComposer
 import org.jetbrains.kotlin.ir.util.SymbolTable
 import org.jetbrains.kotlin.psi2ir.descriptors.IrBuiltInsOverDescriptors
@@ -44,28 +45,34 @@ class SymbolTableWithBuiltInsDeduplication(
         }
     }
 
-    /**
+    override fun createDescriptorExtension(): DescriptorSymbolTableExtension {
-     * Gets or creates the [IrClassSymbol] for [descriptor], or for the built-in descriptor with the same name if [descriptor] is a
+        return Extension()
-     * duplicate built-in.
+    }
-     *
-     * Note that not all built-in symbols may have been bound or created by the time [irBuiltIns] has been bound. However, [referenceClass]
-     * will create a symbol in such a case (via `super.referenceClass`) and [org.jetbrains.kotlin.ir.util.DeclarationStubGenerator] will
-     * create a stub for the symbol if [referenceClass] was invoked from the stub generator.
-     */
-    @ObsoleteDescriptorBasedAPI
-    override fun referenceClass(descriptor: ClassDescriptor): IrClassSymbol {
-        val irBuiltIns = this.irBuiltIns ?: return super.descriptorExtension.referenceClass(descriptor)
 
-        // We need to find out whether `descriptor` is possibly a built-in symbol before it's actually retrieved to break recursion as
+    private inner class Extension : DescriptorSymbolTableExtension(this) {
-        // `irBuiltIns.findClass` uses `referenceClass` recursively.
+        /**
-        val builtInDescriptor = irBuiltIns.findBuiltInClassDescriptor(descriptor)
+         * Gets or creates the [IrClassSymbol] for [descriptor], or for the built-in descriptor with the same name if [descriptor] is a
-        if (builtInDescriptor != null) {
+         * duplicate built-in.
-            // We need to delegate to the supertype implementation here to break recursion. `findBuiltInClassDescriptor` will return
+         *
-            // `descriptor` even if `descriptor` was found via `findBuiltInClassDescriptor`.
+         * Note that not all built-in symbols may have been bound or created by the time [irBuiltIns] has been bound. However, [referenceClass]
-            return super.descriptorExtension.referenceClass(builtInDescriptor)
+         * will create a symbol in such a case (via `super.referenceClass`) and [org.jetbrains.kotlin.ir.util.DeclarationStubGenerator] will
+         * create a stub for the symbol if [referenceClass] was invoked from the stub generator.
+         */
+        @ObsoleteDescriptorBasedAPI
+        override fun referenceClass(descriptor: ClassDescriptor): IrClassSymbol {
+            val irBuiltIns = this@SymbolTableWithBuiltInsDeduplication.irBuiltIns ?: return super.referenceClass(descriptor)
+
+            // We need to find out whether `descriptor` is possibly a built-in symbol before it's actually retrieved to break recursion as
+            // `irBuiltIns.findClass` uses `referenceClass` recursively.
+            val builtInDescriptor = irBuiltIns.findBuiltInClassDescriptor(descriptor)
+            if (builtInDescriptor != null) {
+                // We need to delegate to the supertype implementation here to break recursion. `findBuiltInClassDescriptor` will return
+                // `descriptor` even if `descriptor` was found via `findBuiltInClassDescriptor`.
+                return super.referenceClass(builtInDescriptor)
+            }
+
+            return super.referenceClass(descriptor)
         }
-
-        return super.descriptorExtension.referenceClass(descriptor)
     }
 
     private fun IrBuiltInsOverDescriptors.findBuiltInClassDescriptor(descriptor: ClassDescriptor): ClassDescriptor? {

compiler/ir/ir.psi2ir/src/org/jetbrains/kotlin/psi2ir/generators/fragments/FragmentCompilerSymbolTableDecorator.kt

@@ -12,6 +12,7 @@ import org.jetbrains.kotlin.descriptors.impl.AbstractReceiverParameterDescriptor
 import org.jetbrains.kotlin.ir.declarations.IrFactory
 import org.jetbrains.kotlin.ir.symbols.IrValueParameterSymbol
 import org.jetbrains.kotlin.ir.symbols.IrValueSymbol
+import org.jetbrains.kotlin.ir.util.DescriptorSymbolTableExtension
 import org.jetbrains.kotlin.ir.util.IdSignatureComposer
 import org.jetbrains.kotlin.ir.util.NameProvider
 import org.jetbrains.kotlin.ir.util.SymbolTable
@@ -28,48 +29,53 @@ class FragmentCompilerSymbolTableDecorator(
     nameProvider: NameProvider = NameProvider.DEFAULT,
 ) : SymbolTable(signatureComposer, irFactory, nameProvider) {
 
-    override fun referenceValueParameter(descriptor: ParameterDescriptor): IrValueParameterSymbol {
+    override fun createDescriptorExtension(): DescriptorSymbolTableExtension {
-        val fi = fragmentInfo ?: return super.descriptorExtension.referenceValueParameter(descriptor)
+        return ExtensionDecorator()
-
-        if (descriptor !is ReceiverParameterDescriptor) return super.descriptorExtension.referenceValueParameter(descriptor)
-
-        val finderPredicate = when (val receiverValue = descriptor.value) {
-            is ExtensionReceiver, is ContextReceiver -> { (targetDescriptor, _): EvaluatorFragmentParameterInfo ->
-                receiverValue == (targetDescriptor as? ReceiverParameterDescriptor)?.value
-            }
-            is ThisClassReceiver -> { (targetDescriptor, _): EvaluatorFragmentParameterInfo ->
-                receiverValue.classDescriptor == targetDescriptor.original
-            }
-            else -> TODO("Unimplemented")
-        }
-
-        val parameterPosition =
-            fi.parameters.indexOfFirst(finderPredicate)
-        if (parameterPosition > -1) {
-            return super.descriptorExtension.referenceValueParameter(fi.methodDescriptor.valueParameters[parameterPosition])
-        }
-        return super.descriptorExtension.referenceValueParameter(descriptor)
     }
 
+    private inner class ExtensionDecorator : DescriptorSymbolTableExtension(this) {
+        override fun referenceValueParameter(descriptor: ParameterDescriptor): IrValueParameterSymbol {
+            val fi = fragmentInfo ?: return super.referenceValueParameter(descriptor)
 
-    fun referenceValue(value: ValueDescriptor): IrValueSymbol {
+            if (descriptor !is ReceiverParameterDescriptor) return super.referenceValueParameter(descriptor)
-        val fi = fragmentInfo ?: return super.descriptorExtension.referenceValue(value)
 
-        val finderPredicate = when (value) {
+            val finderPredicate = when (val receiverValue = descriptor.value) {
-            is AbstractReceiverParameterDescriptor -> { (targetDescriptor, _): EvaluatorFragmentParameterInfo ->
+                is ExtensionReceiver, is ContextReceiver -> { (targetDescriptor, _): EvaluatorFragmentParameterInfo ->
-                value.containingDeclaration == targetDescriptor
+                    receiverValue == (targetDescriptor as? ReceiverParameterDescriptor)?.value
+                }
+                is ThisClassReceiver -> { (targetDescriptor, _): EvaluatorFragmentParameterInfo ->
+                    receiverValue.classDescriptor == targetDescriptor.original
+                }
+                else -> TODO("Unimplemented")
             }
-            else -> { (targetDescriptor, _): EvaluatorFragmentParameterInfo ->
+
-                targetDescriptor == value
+            val parameterPosition =
+                fi.parameters.indexOfFirst(finderPredicate)
+            if (parameterPosition > -1) {
+                return super.referenceValueParameter(fi.methodDescriptor.valueParameters[parameterPosition])
             }
+            return super.referenceValueParameter(descriptor)
         }
 
-        val parameterPosition =
+        override fun referenceValue(value: ValueDescriptor): IrValueSymbol {
-            fi.parameters.indexOfFirst(finderPredicate)
+            val fi = fragmentInfo ?: return super.referenceValue(value)
-        if (parameterPosition > -1) {
-            return super.descriptorExtension.referenceValueParameter(fi.methodDescriptor.valueParameters[parameterPosition])
-        }
 
-        return super.descriptorExtension.referenceValue(value)
+            val finderPredicate = when (value) {
+                is AbstractReceiverParameterDescriptor -> { (targetDescriptor, _): EvaluatorFragmentParameterInfo ->
+                    value.containingDeclaration == targetDescriptor
+                }
+                else -> { (targetDescriptor, _): EvaluatorFragmentParameterInfo ->
+                    targetDescriptor == value
+                }
+            }
+
+            val parameterPosition =
+                fi.parameters.indexOfFirst(finderPredicate)
+            if (parameterPosition > -1) {
+                return super.referenceValueParameter(fi.methodDescriptor.valueParameters[parameterPosition])
+            }
+
+            return super.referenceValue(value)
+        }
     }
 }

compiler/ir/ir.tree/src/org/jetbrains/kotlin/ir/declarations/lazy/IrLazySymbolTable.kt

@@ -9,9 +9,7 @@ import org.jetbrains.kotlin.descriptors.*
 import org.jetbrains.kotlin.ir.IrLock
 import org.jetbrains.kotlin.ir.ObsoleteDescriptorBasedAPI
 import org.jetbrains.kotlin.ir.symbols.*
-import org.jetbrains.kotlin.ir.util.DeclarationStubGenerator
+import org.jetbrains.kotlin.ir.util.*
-import org.jetbrains.kotlin.ir.util.ReferenceSymbolTable
-import org.jetbrains.kotlin.ir.util.SymbolTable
 
 @OptIn(ObsoleteDescriptorBasedAPI::class)
 class IrLazySymbolTable(private val originalTable: SymbolTable) : ReferenceSymbolTable by originalTable {
@@ -21,73 +19,104 @@ class IrLazySymbolTable(private val originalTable: SymbolTable) : ReferenceSymbo
     /*Don't force builtins class linking before unbound symbols linking: otherwise stdlib compilation will failed*/
     var stubGenerator: DeclarationStubGenerator? = null
 
-    fun referenceClass(descriptor: ClassDescriptor): IrClassSymbol {
+    @ObsoleteDescriptorBasedAPI
-        synchronized(lock) {
+    override val descriptorExtension: DescriptorBasedReferenceSymbolTableExtension = ExtensionWrapper()
-            return originalTable.descriptorExtension.referenceClass(descriptor).also {
+
-                if (!it.isBound) {
+    private inner class ExtensionWrapper : DescriptorSymbolTableExtension(originalTable) {
-                    stubGenerator?.generateClassStub(descriptor)
+        private val delegate get() = originalTable.descriptorExtension
+
+        override fun referenceClass(descriptor: ClassDescriptor): IrClassSymbol {
+            synchronized(lock) {
+                return delegate.referenceClass(descriptor).also {
+                    if (!it.isBound) {
+                        stubGenerator?.generateClassStub(descriptor)
+                    }
                 }
             }
         }
-    }
 
-    fun referenceTypeAlias(descriptor: TypeAliasDescriptor): IrTypeAliasSymbol {
+        override fun referenceTypeAlias(descriptor: TypeAliasDescriptor): IrTypeAliasSymbol {
-        synchronized(lock) {
+            synchronized(lock) {
-            return originalTable.descriptorExtension.referenceTypeAlias(descriptor).also {
+                return delegate.referenceTypeAlias(descriptor).also {
-                if (!it.isBound) {
+                    if (!it.isBound) {
-                    stubGenerator?.generateTypeAliasStub(descriptor)
+                        stubGenerator?.generateTypeAliasStub(descriptor)
+                    }
                 }
             }
         }
-    }
 
-    fun referenceConstructor(descriptor: ClassConstructorDescriptor): IrConstructorSymbol {
+        override fun referenceConstructor(descriptor: ClassConstructorDescriptor): IrConstructorSymbol {
-        synchronized(lock) {
+            synchronized(lock) {
-            return originalTable.descriptorExtension.referenceConstructor(descriptor).also {
+                return delegate.referenceConstructor(descriptor).also {
-                if (!it.isBound) {
+                    if (!it.isBound) {
-                    stubGenerator?.generateConstructorStub(descriptor)
+                        stubGenerator?.generateConstructorStub(descriptor)
+                    }
                 }
             }
         }
-    }
 
-    fun referenceEnumEntry(descriptor: ClassDescriptor): IrEnumEntrySymbol {
+        override fun referenceEnumEntry(descriptor: ClassDescriptor): IrEnumEntrySymbol {
-        synchronized(lock) {
+            synchronized(lock) {
-            return originalTable.descriptorExtension.referenceEnumEntry(descriptor).also {
+                return delegate.referenceEnumEntry(descriptor).also {
-                if (!it.isBound) {
+                    if (!it.isBound) {
-                    stubGenerator?.generateEnumEntryStub(descriptor)
+                        stubGenerator?.generateEnumEntryStub(descriptor)
+                    }
                 }
             }
         }
-    }
 
-    fun referenceSimpleFunction(descriptor: FunctionDescriptor): IrSimpleFunctionSymbol {
+        override fun referenceSimpleFunction(descriptor: FunctionDescriptor): IrSimpleFunctionSymbol {
-        synchronized(lock) {
+            synchronized(lock) {
-            return originalTable.descriptorExtension.referenceSimpleFunction(descriptor).also {
+                return delegate.referenceSimpleFunction(descriptor).also {
-                if (!it.isBound) {
+                    if (!it.isBound) {
-                    stubGenerator?.generateFunctionStub(descriptor)
+                        stubGenerator?.generateFunctionStub(descriptor)
+                    }
                 }
             }
         }
-    }
 
-    fun referenceProperty(descriptor: PropertyDescriptor): IrPropertySymbol {
+        override fun referenceProperty(descriptor: PropertyDescriptor): IrPropertySymbol {
-        synchronized(lock) {
+            synchronized(lock) {
-            return originalTable.descriptorExtension.referenceProperty(descriptor).also {
+                return delegate.referenceProperty(descriptor).also {
-                if (!it.isBound) {
+                    if (!it.isBound) {
-                    stubGenerator?.generatePropertyStub(descriptor)
+                        stubGenerator?.generatePropertyStub(descriptor)
+                    }
                 }
             }
         }
-    }
 
-    fun referenceTypeParameter(classifier: TypeParameterDescriptor): IrTypeParameterSymbol {
+        override fun referenceTypeParameter(classifier: TypeParameterDescriptor): IrTypeParameterSymbol {
-        synchronized(lock) {
+            synchronized(lock) {
-            return originalTable.descriptorExtension.referenceTypeParameter(classifier).also {
+                return delegate.referenceTypeParameter(classifier).also {
-                if (!it.isBound) {
+                    if (!it.isBound) {
-                    stubGenerator?.generateOrGetTypeParameterStub(classifier)
+                        stubGenerator?.generateOrGetTypeParameterStub(classifier)
+                    }
                 }
             }
         }
+
+        override fun referenceValueParameter(descriptor: ParameterDescriptor): IrValueParameterSymbol {
+            return delegate.referenceValueParameter(descriptor)
+        }
+
+        override fun referenceValue(value: ValueDescriptor): IrValueSymbol {
+            return delegate.referenceValue(value)
+        }
+
+        override fun referenceScript(descriptor: ScriptDescriptor): IrScriptSymbol {
+            return delegate.referenceScript(descriptor)
+        }
+
+        override fun referenceField(descriptor: PropertyDescriptor): IrFieldSymbol {
+            return delegate.referenceField(descriptor)
+        }
+
+        override fun referenceDeclaredFunction(descriptor: FunctionDescriptor): IrSimpleFunctionSymbol {
+            return delegate.referenceDeclaredFunction(descriptor)
+        }
+
+        override fun referenceScopedTypeParameter(classifier: TypeParameterDescriptor): IrTypeParameterSymbol {
+            return delegate.referenceScopedTypeParameter(classifier)
+        }
     }
 }

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

@@ -27,7 +27,7 @@ typealias DescriptorBasedReferenceSymbolTableExtension = ReferenceSymbolTableExt
         >
 
 @ObsoleteDescriptorBasedAPI
-class DescriptorSymbolTableExtension(table: SymbolTable) : SymbolTableExtension<
+open class DescriptorSymbolTableExtension(table: SymbolTable) : SymbolTableExtension<
         DeclarationDescriptor, ClassDescriptor, TypeAliasDescriptor, ScriptDescriptor, FunctionDescriptor,
         ClassConstructorDescriptor, PropertyDescriptor, ParameterDescriptor, TypeParameterDescriptor>(table)
 {
@@ -405,8 +405,7 @@ class DescriptorSymbolTableExtension(table: SymbolTable) : SymbolTableExtension<
         }
     }
 
-    // TODO: convert to extension
+    open fun referenceValue(value: ValueDescriptor): IrValueSymbol {
-    fun referenceValue(value: ValueDescriptor): IrValueSymbol {
         return when (value) {
             is ParameterDescriptor -> valueParameterSlice.referenced(value) { error("Undefined parameter referenced: $value") }
             is VariableDescriptor -> variableSlice.referenced(value) { error("Undefined variable referenced: $value") }

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

@@ -35,7 +35,12 @@ open class SymbolTable(
 
     @Suppress("LeakingThis")
     @OptIn(ObsoleteDescriptorBasedAPI::class)
-    override val descriptorExtension: DescriptorSymbolTableExtension = DescriptorSymbolTableExtension(this)
+    override val descriptorExtension: DescriptorSymbolTableExtension = createDescriptorExtension()
+
+    @OptIn(ObsoleteDescriptorBasedAPI::class)
+    protected open fun createDescriptorExtension(): DescriptorSymbolTableExtension {
+        return DescriptorSymbolTableExtension(this)
+    }
 
     // ------------------------------------ script ------------------------------------