[PSI2IR] Support context receivers on classes

2020-09-04 11:24:28 +07:00
parent 307f318c9e
commit aaabf5e1ca
23 changed files with 368 additions and 15 deletions
@@ -34,6 +34,7 @@ import org.jetbrains.kotlin.ir.symbols.impl.IrSimpleFunctionSymbolImpl
 import org.jetbrains.kotlin.ir.symbols.impl.IrValueParameterSymbolImpl
 import org.jetbrains.kotlin.ir.types.IrType
 import org.jetbrains.kotlin.name.Name
+import org.jetbrains.kotlin.psi.KtConstructorDelegationCall
 import org.jetbrains.kotlin.psi.KtElement
 import org.jetbrains.kotlin.psi.KtExpression
 import org.jetbrains.kotlin.psi.ValueArgument
@@ -92,6 +93,20 @@ fun StatementGenerator.generateReceiver(defaultStartOffset: Int, defaultEndOffse
                            context.symbolTable.referenceValueParameter(receiverClassDescriptor.thisAsReceiverParameter)
                        )
                }
+                is ExtensionClassReceiver -> {
+                    val receiverClassDescriptor = receiver.classDescriptor
+                    val thisAsReceiverParameter = receiverClassDescriptor.thisAsReceiverParameter
+                    val thisReceiver = IrGetValueImpl(
+                        defaultStartOffset, defaultEndOffset,
+                        thisAsReceiverParameter.type.toIrType(),
+                        context.symbolTable.referenceValue(thisAsReceiverParameter)
+                    )
+                    IrGetFieldImpl(
+                        defaultStartOffset, defaultEndOffset,
+                        context.symbolTable.referenceField(context.additionalDescriptorStorage.getField(receiver)),
+                        irReceiverType, thisReceiver
+                    )
+                }
                is ThisClassReceiver ->
                    generateThisOrSuperReceiver(receiver, receiver.classDescriptor)
                is SuperCallReceiverValue ->
@@ -210,7 +225,10 @@ fun StatementGenerator.generateCallReceiver(
        else -> {
            dispatchReceiverValue = generateReceiverOrNull(ktDefaultElement, dispatchReceiver)
            extensionReceiverValue = generateReceiverOrNull(ktDefaultElement, extensionReceiver)
-            contextReceiverValues = contextReceivers.mapNotNull { generateReceiverOrNull(ktDefaultElement, it) }
+            contextReceiverValues = if (ktDefaultElement is KtConstructorDelegationCall) contextReceivers.mapNotNull {
+                generateReceiverOrNull(ktDefaultElement, it as ExtensionClassReceiver)
+            }
+            else contextReceivers.mapNotNull { generateReceiverOrNull(ktDefaultElement, it) }
        }
    }

@@ -18,7 +18,9 @@ package org.jetbrains.kotlin.psi2ir.generators

 import org.jetbrains.kotlin.builtins.KotlinBuiltIns
 import org.jetbrains.kotlin.descriptors.*
+import org.jetbrains.kotlin.ir.UNDEFINED_OFFSET
 import org.jetbrains.kotlin.ir.builders.Scope
+import org.jetbrains.kotlin.ir.declarations.IrConstructor
 import org.jetbrains.kotlin.ir.declarations.IrFunction
 import org.jetbrains.kotlin.ir.expressions.*
 import org.jetbrains.kotlin.ir.expressions.impl.*
@@ -175,12 +177,15 @@ class BodyGenerator(
    fun getLoop(expression: KtExpression): IrLoop? =
        loopTable[expression]

-    fun generatePrimaryConstructorBody(ktClassOrObject: KtPureClassOrObject): IrBody {
+    fun generatePrimaryConstructorBody(ktClassOrObject: KtPureClassOrObject, irConstructor: IrConstructor): IrBody {
        val irBlockBody = context.irFactory.createBlockBody(ktClassOrObject.pureStartOffset, ktClassOrObject.pureEndOffset)

        generateSuperConstructorCall(irBlockBody, ktClassOrObject)

        val classDescriptor = (scopeOwner as ClassConstructorDescriptor).containingDeclaration
+        if (classDescriptor.contextReceivers.isNotEmpty()) {
+            generateSetAdditionalFieldForPrimaryConstructorBody(classDescriptor, irConstructor, irBlockBody)
+        }
        irBlockBody.statements.add(
            IrInstanceInitializerCallImpl(
                ktClassOrObject.pureStartOffset, ktClassOrObject.pureEndOffset,
@@ -329,4 +334,28 @@ class BodyGenerator(
            pregenerateCall(constructorCall)
        )

+    private fun generateSetAdditionalFieldForPrimaryConstructorBody(
+        classDescriptor: ClassDescriptor,
+        irConstructor: IrConstructor,
+        irBlockBody: IrBlockBody
+    ) {
+        val thisAsReceiverParameter = classDescriptor.thisAsReceiverParameter
+        val receiver = IrGetValueImpl(
+            UNDEFINED_OFFSET, UNDEFINED_OFFSET,
+            thisAsReceiverParameter.type.toIrType(),
+            context.symbolTable.referenceValue(thisAsReceiverParameter)
+        )
+        for ((index, receiverDescriptor) in classDescriptor.contextReceivers.withIndex()) {
+            val irValueParameter = irConstructor.valueParameters[index]
+            irBlockBody.statements.add(
+                IrSetFieldImpl(
+                    UNDEFINED_OFFSET, UNDEFINED_OFFSET,
+                    context.symbolTable.referenceField(context.additionalDescriptorStorage.getField(receiverDescriptor.value)),
+                    receiver,
+                    IrGetValueImpl(UNDEFINED_OFFSET, UNDEFINED_OFFSET, irValueParameter.type, irValueParameter.symbol),
+                    context.irBuiltIns.unitType
+                )
+            )
+        }
+    }
 }
@@ -20,6 +20,11 @@ import org.jetbrains.kotlin.backend.common.CodegenUtil
 import org.jetbrains.kotlin.builtins.KotlinBuiltIns
 import org.jetbrains.kotlin.descriptors.*
 import org.jetbrains.kotlin.ir.ObsoleteDescriptorBasedAPI
+import org.jetbrains.kotlin.descriptors.annotations.AnnotationUseSiteTarget
+import org.jetbrains.kotlin.descriptors.annotations.Annotations
+import org.jetbrains.kotlin.descriptors.impl.FieldDescriptorImpl
+import org.jetbrains.kotlin.descriptors.impl.PropertyDescriptorImpl
+import org.jetbrains.kotlin.ir.UNDEFINED_OFFSET
 import org.jetbrains.kotlin.ir.declarations.*
 import org.jetbrains.kotlin.ir.descriptors.IrImplementingDelegateDescriptorImpl
 import org.jetbrains.kotlin.ir.expressions.IrBlockBody
@@ -35,6 +40,8 @@ import org.jetbrains.kotlin.ir.util.SYNTHETIC_OFFSET
 import org.jetbrains.kotlin.ir.util.createIrClassFromDescriptor
 import org.jetbrains.kotlin.ir.util.declareSimpleFunctionWithOverrides
 import org.jetbrains.kotlin.ir.util.properties
+import org.jetbrains.kotlin.ir.util.referenceFunction
+import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.psi.KtClassOrObject
 import org.jetbrains.kotlin.psi.KtDelegatedSuperTypeEntry
 import org.jetbrains.kotlin.psi.KtEnumEntry
@@ -47,6 +54,7 @@ import org.jetbrains.kotlin.psi.synthetics.SyntheticClassOrObjectDescriptor
 import org.jetbrains.kotlin.psi.synthetics.findClassDescriptor
 import org.jetbrains.kotlin.resolve.BindingContext
 import org.jetbrains.kotlin.resolve.DelegationResolver
+import org.jetbrains.kotlin.resolve.DescriptorFactory
 import org.jetbrains.kotlin.resolve.DescriptorUtils
 import org.jetbrains.kotlin.resolve.descriptorUtil.propertyIfAccessor
 import org.jetbrains.kotlin.resolve.descriptorUtil.setSingleOverridden
@@ -89,6 +97,8 @@ class ClassGenerator(
                classDescriptor.thisAsReceiverParameter.type.toIrType()
            )

+            generateFieldsForAdditionalReceivers(irClass, classDescriptor)
+
            val irPrimaryConstructor = generatePrimaryConstructor(irClass, ktClassOrObject)
            if (irPrimaryConstructor != null) {
                generateDeclarationsForPrimaryConstructorParameters(irClass, irPrimaryConstructor, ktClassOrObject)
@@ -423,6 +433,41 @@ class ClassGenerator(
        EnumClassMembersGenerator(declarationGenerator).generateSpecialMembers(irClass)
    }

+    private fun generateFieldsForAdditionalReceivers(irClass: IrClass, classDescriptor: ClassDescriptor) {
+        for ((fieldIndex, receiverDescriptor) in classDescriptor.contextReceivers.withIndex()) {
+            val descriptor = PropertyDescriptorImpl.create(
+                classDescriptor,
+                Annotations.EMPTY,
+                Modality.FINAL,
+                DescriptorVisibilities.DEFAULT_VISIBILITY,
+                false,
+                Name.identifier("additionalReceiverField$fieldIndex"),
+                CallableMemberDescriptor.Kind.DECLARATION,
+                SourceElement.NO_SOURCE,
+                false, false,
+                classDescriptor.isExpect,
+                false, false, false
+            )
+            descriptor.setType(
+                receiverDescriptor.type, emptyList(), DescriptorUtils.getDispatchReceiverParameterIfNeeded(classDescriptor), null,
+                emptyList()
+            )
+            descriptor.initialize(
+                null, null,
+                FieldDescriptorImpl(Annotations.EMPTY, descriptor),
+                null
+            )
+            context.additionalDescriptorStorage.put(receiverDescriptor.value, descriptor)
+            val irField = context.symbolTable.declareField(
+                UNDEFINED_OFFSET, UNDEFINED_OFFSET,
+                IrDeclarationOrigin.FIELD_FOR_CLASS_CONTEXT_RECEIVER,
+                descriptor, descriptor.type.toIrType(),
+                descriptor.visibility
+            )
+            irClass.addMember(irField)
+        }
+    }
+
    private fun generatePrimaryConstructor(irClass: IrClass, ktClassOrObject: KtPureClassOrObject): IrConstructor? {
        val classDescriptor = irClass.descriptor
        val primaryConstructorDescriptor = classDescriptor.unsubstitutedPrimaryConstructor ?: return null
@@ -459,7 +504,7 @@ class ClassGenerator(
        // generate real body declarations
        ktClassOrObject.body?.let { ktClassBody ->
            ktClassBody.declarations.mapNotNullTo(irClass.declarations) { ktDeclaration ->
-                declarationGenerator.generateClassMemberDeclaration(ktDeclaration, irClass)
+                declarationGenerator.generateClassMemberDeclaration(ktDeclaration, irClass, ktClassOrObject)
            }
        }

@@ -85,12 +85,16 @@ class DeclarationGenerator(override val context: GeneratorContext) : Generator {
        return generateClassOrObjectDeclaration(syntheticDeclaration)
    }

-    fun generateClassMemberDeclaration(ktDeclaration: KtDeclaration, irClass: IrClass): IrDeclaration? =
+    fun generateClassMemberDeclaration(
+        ktDeclaration: KtDeclaration,
+        irClass: IrClass,
+        ktClassOrObject: KtPureClassOrObject
+    ): IrDeclaration? =
        when (ktDeclaration) {
            is KtAnonymousInitializer ->
                AnonymousInitializerGenerator(this).generateAnonymousInitializerDeclaration(ktDeclaration, irClass)
            is KtSecondaryConstructor ->
-                FunctionGenerator(this).generateSecondaryConstructor(ktDeclaration)
+                FunctionGenerator(this).generateSecondaryConstructor(ktDeclaration, ktClassOrObject)
            is KtEnumEntry ->
                generateEnumEntryDeclaration(ktDeclaration)
            else ->
@@ -0,0 +1,29 @@
+/*
+ * Copyright 2010-2020 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * 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.psi2ir.generators
+
+import org.jetbrains.kotlin.descriptors.PropertyDescriptor
+import org.jetbrains.kotlin.descriptors.VariableDescriptor
+import org.jetbrains.kotlin.psi.KtExpression
+import org.jetbrains.kotlin.resolve.scopes.receivers.ReceiverValue
+
+class DescriptorStorageForAdditionalReceivers {
+    private val fieldStorage: MutableMap<ReceiverValue, PropertyDescriptor> = mutableMapOf()
+    private val variableStorage: MutableMap<KtExpression, VariableDescriptor> = mutableMapOf()
+
+    fun put(receiverValue: ReceiverValue, descriptor: PropertyDescriptor) {
+        fieldStorage[receiverValue] = descriptor
+    }
+
+    fun put(expression: KtExpression, descriptor: VariableDescriptor) {
+        variableStorage[expression] = descriptor
+    }
+
+    fun getField(receiverValue: ReceiverValue) =
+        fieldStorage[receiverValue] ?: error("No field descriptor for receiver value $receiverValue")
+
+    fun getVariable(expression: KtExpression) = variableStorage[expression] ?: error("No variable descriptor for receiver $expression")
+}
@@ -234,19 +234,29 @@ class FunctionGenerator(declarationGenerator: DeclarationGenerator) : Declaratio
        primaryConstructorDescriptor: ClassConstructorDescriptor,
        ktClassOrObject: KtPureClassOrObject
    ): IrConstructor =
-        declareConstructor(ktClassOrObject, ktClassOrObject.primaryConstructor ?: ktClassOrObject, primaryConstructorDescriptor) {
+        declareConstructor(
+            ktClassOrObject,
+            ktClassOrObject.primaryConstructor ?: ktClassOrObject,
+            ktClassOrObject.contextReceiverTypeReferences,
+            primaryConstructorDescriptor
+        ) { irConstructor ->
            if (
                primaryConstructorDescriptor.isExpect ||
                primaryConstructorDescriptor.constructedClass.isEffectivelyExternal()
            )
                null
            else
-                generatePrimaryConstructorBody(ktClassOrObject)
+                generatePrimaryConstructorBody(ktClassOrObject, irConstructor)
        }

-    fun generateSecondaryConstructor(ktConstructor: KtSecondaryConstructor): IrConstructor {
+    fun generateSecondaryConstructor(ktConstructor: KtSecondaryConstructor, ktClassOrObject: KtPureClassOrObject): IrConstructor {
        val constructorDescriptor = getOrFail(BindingContext.CONSTRUCTOR, ktConstructor) as ClassConstructorDescriptor
-        return declareConstructor(ktConstructor, ktConstructor, constructorDescriptor) {
+        return declareConstructor(
+            ktConstructor,
+            ktConstructor,
+            ktClassOrObject.contextReceiverTypeReferences,
+            constructorDescriptor
+        ) {
            when {
                constructorDescriptor.constructedClass.isEffectivelyExternal() ->
                    null
@@ -263,8 +273,9 @@ class FunctionGenerator(declarationGenerator: DeclarationGenerator) : Declaratio
    private inline fun declareConstructor(
        ktConstructorElement: KtPureElement,
        ktParametersElement: KtPureElement,
+        ktContextReceiversElements: List<KtPureElement>,
        constructorDescriptor: ClassConstructorDescriptor,
-        generateBody: BodyGenerator.() -> IrBody?
+        generateBody: BodyGenerator.(IrConstructor) -> IrBody?
    ): IrConstructor {
        val startOffset = ktConstructorElement.getStartOffsetOfConstructorDeclarationKeywordOrNull() ?: ktConstructorElement.pureStartOffset
        val endOffset = ktConstructorElement.pureEndOffset
@@ -279,8 +290,8 @@ class FunctionGenerator(declarationGenerator: DeclarationGenerator) : Declaratio
                metadata = DescriptorMetadataSource.Function(it.descriptor)
            }
        }.buildWithScope { irConstructor ->
-            generateValueParameterDeclarations(irConstructor, ktParametersElement, null, emptyList())
-            irConstructor.body = createBodyGenerator(irConstructor.symbol).generateBody()
+            generateValueParameterDeclarations(irConstructor, ktParametersElement, null, ktContextReceiversElements)
+            irConstructor.body = createBodyGenerator(irConstructor.symbol).generateBody(irConstructor)
            irConstructor.returnType = constructorDescriptor.returnType.toIrType()
        }
    }
@@ -71,6 +71,8 @@ class GeneratorContext private constructor(
    // TODO: inject a correct StorageManager instance, or store NotFoundClasses inside ModuleDescriptor
    val reflectionTypes = ReflectionTypes(moduleDescriptor, NotFoundClasses(LockBasedStorageManager.NO_LOCKS, moduleDescriptor))

+    val additionalDescriptorStorage: DescriptorStorageForAdditionalReceivers = DescriptorStorageForAdditionalReceivers()
+
    val samTypeApproximator = SamTypeApproximator(moduleDescriptor.builtIns, languageVersionSettings)

    fun createFileScopeContext(ktFile: KtFile): GeneratorContext {
@@ -68,6 +68,7 @@ interface IrDeclarationOrigin {
    object SYNTHETIC_HELPER_FOR_ENUM_VALUES : IrDeclarationOriginImpl("SYNTHETIC_HELPER_FOR_ENUM_VALUES", isSynthetic = true)
    object FIELD_FOR_ENUM_VALUES : IrDeclarationOriginImpl("FIELD_FOR_ENUM_VALUES", isSynthetic = true)
    object FIELD_FOR_OBJECT_INSTANCE : IrDeclarationOriginImpl("FIELD_FOR_OBJECT_INSTANCE")
+    object FIELD_FOR_CLASS_CONTEXT_RECEIVER : IrDeclarationOriginImpl("FIELD_FOR_CLASS_CONTEXT_RECEIVER", isSynthetic = true)

    object ADAPTER_FOR_CALLABLE_REFERENCE : IrDeclarationOriginImpl("ADAPTER_FOR_CALLABLE_REFERENCE", isSynthetic = true)
    object ADAPTER_PARAMETER_FOR_CALLABLE_REFERENCE : IrDeclarationOriginImpl("ADAPTER_PARAMETER_FOR_CALLABLE_REFERENCE")
@@ -40,7 +40,7 @@ class IrConstructorCallImpl(
            val constructorDescriptor = constructorSymbol.descriptor
            val classTypeParametersCount = constructorDescriptor.constructedClass.original.declaredTypeParameters.size
            val totalTypeParametersCount = constructorDescriptor.typeParameters.size
-            val valueParametersCount = constructorDescriptor.valueParameters.size
+            val valueParametersCount = constructorDescriptor.valueParameters.size + constructorDescriptor.contextReceiverParameters.size
            return IrConstructorCallImpl(
                startOffset, endOffset,
                type,
@@ -48,7 +48,7 @@ class IrDelegatingConstructorCallImpl(
            type: IrType,
            symbol: IrConstructorSymbol,
            typeArgumentsCount: Int = symbol.descriptor.typeParametersCount,
-            valueArgumentsCount: Int = symbol.descriptor.valueParameters.size
+            valueArgumentsCount: Int = symbol.descriptor.valueParameters.size + symbol.descriptor.contextReceiverParameters.size
        ) = IrDelegatingConstructorCallImpl(startOffset, endOffset, type, symbol, typeArgumentsCount, valueArgumentsCount)

        fun fromSymbolOwner(