Implement support for additional receivers in the backend

compiler/backend/src/org/jetbrains/kotlin/codegen/ExpressionCodegen.java

@@ -2749,7 +2749,13 @@ public class ExpressionCodegen extends KtVisitor<StackValue, StackValue> impleme
             }
             else {
                 cur = getNotNullParentContextForMethod(cur);
-                result = cur.getOuterExpression(result, false);
+                // for now the script codegen only passes this branch, since the method context for script constructor is defined using function context
+                if (cur instanceof ScriptContext && !(thisOrOuterClass instanceof ScriptDescriptor)) {
+                    return ((ScriptContext) cur).getOuterReceiverExpression(result, thisOrOuterClass);
+                }
+                else {
+                    result = cur.getOuterExpression(result, false);
+                }
             }
 
             cur = cur.getEnclosingClassContext();

compiler/backend/src/org/jetbrains/kotlin/codegen/ScriptCodegen.kt

@@ -10,8 +10,8 @@ import org.jetbrains.kotlin.codegen.context.CodegenContext
 import org.jetbrains.kotlin.codegen.context.MethodContext
 import org.jetbrains.kotlin.codegen.context.ScriptContext
 import org.jetbrains.kotlin.codegen.state.GenerationState
+import org.jetbrains.kotlin.descriptors.ClassDescriptor
 import org.jetbrains.kotlin.descriptors.ScriptDescriptor
-import org.jetbrains.kotlin.descriptors.ValueParameterDescriptor
 import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.resolve.BindingContext
 import org.jetbrains.kotlin.resolve.descriptorUtil.*
@@ -66,7 +66,9 @@ class ScriptCodegen private constructor(
             classBuilder: ClassBuilder,
             methodContext: MethodContext
     ) {
-        val jvmSignature = typeMapper.mapScriptSignature(scriptDescriptor, scriptContext.earlierScripts)
+        val scriptDefinition = scriptContext.script.kotlinScriptDefinition.value
+
+        val jvmSignature = typeMapper.mapScriptSignature(scriptDescriptor, scriptContext.earlierScripts, scriptDefinition.implicitReceivers)
 
         if (state.replSpecific.shouldGenerateScriptResultValue) {
             val resultFieldInfo = scriptContext.resultFieldInfo
@@ -102,7 +104,10 @@ class ScriptCodegen private constructor(
 
                 iv.load(0, classType)
 
-                val valueParamStart = if (scriptContext.earlierScripts.isEmpty()) 1 else 2 // this + array of earlier scripts if not empty
+                fun Int.incrementIf(cond: Boolean): Int = if (cond) plus(1) else this
+                val valueParamStart = 1
+                    .incrementIf(scriptContext.earlierScripts.isNotEmpty())
+                    .incrementIf(scriptDefinition.implicitReceivers.isNotEmpty())
 
                 val valueParameters = scriptDescriptor.unsubstitutedPrimaryConstructor.valueParameters
                 for (superclassParam in ctorDesc.valueParameters) {
@@ -122,19 +127,31 @@ class ScriptCodegen private constructor(
             val frameMap = FrameMap()
             frameMap.enterTemp(OBJECT_TYPE)
 
+            fun genFieldFromArrayElement(descriptor: ClassDescriptor, paramIndex: Int, elementIndex: Int, name: String) {
+                val elementClassType = typeMapper.mapClass(descriptor)
+                iv.load(0, classType)
+                iv.load(paramIndex, elementClassType)
+                iv.aconst(elementIndex)
+                iv.aload(OBJECT_TYPE)
+                iv.checkcast(elementClassType)
+                iv.putfield(classType.internalName, name, elementClassType.descriptor)
+            }
 
             if (!scriptContext.earlierScripts.isEmpty()) {
                 val scriptsParamIndex = frameMap.enterTemp(AsmUtil.getArrayType(OBJECT_TYPE))
 
-                var earlierScriptIndex = 0
+                scriptContext.earlierScripts.forEachIndexed { earlierScriptIndex, earlierScript ->
-                for (earlierScript in scriptContext.earlierScripts) {
+                    val name = scriptContext.getScriptFieldName(earlierScript)
-                    val earlierClassType = typeMapper.mapClass(earlierScript)
+                    genFieldFromArrayElement(earlierScript, scriptsParamIndex, earlierScriptIndex, name)
-                    iv.load(0, classType)
+                }
-                    iv.load(scriptsParamIndex, earlierClassType)
+            }
-                    iv.aconst(earlierScriptIndex++)
+
-                    iv.aload(OBJECT_TYPE)
+            if (scriptDefinition.implicitReceivers.isNotEmpty()) {
-                    iv.checkcast(earlierClassType)
+                val receiversParamIndex = frameMap.enterTemp(AsmUtil.getArrayType(OBJECT_TYPE))
-                    iv.putfield(classType.internalName, scriptContext.getScriptFieldName(earlierScript), earlierClassType.descriptor)
+
+                scriptContext.receiverDescriptors.forEachIndexed { receiverIndex, receiver ->
+                    val name = scriptContext.getImplicitReceiverName(receiverIndex)
+                    genFieldFromArrayElement(receiver, receiversParamIndex, receiverIndex, name)
                 }
             }
 
@@ -151,9 +168,24 @@ class ScriptCodegen private constructor(
 
     private fun genFieldsForParameters(classBuilder: ClassBuilder) {
         for (earlierScript in scriptContext.earlierScripts) {
-            val earlierClassName = typeMapper.mapType(earlierScript)
+            classBuilder.newField(
-            val access = ACC_PUBLIC or ACC_FINAL
+                NO_ORIGIN,
-            classBuilder.newField(NO_ORIGIN, access, scriptContext.getScriptFieldName(earlierScript), earlierClassName.descriptor, null, null)
+                ACC_PUBLIC or ACC_FINAL,
+                scriptContext.getScriptFieldName(earlierScript),
+                typeMapper.mapType(earlierScript).descriptor,
+                null,
+                null
+            )
+        }
+        for (receiverIndex in scriptContext.receiverDescriptors.indices) {
+            classBuilder.newField(
+                NO_ORIGIN,
+                ACC_PUBLIC or ACC_FINAL,
+                scriptContext.getImplicitReceiverName(receiverIndex),
+                scriptContext.getImplicitReceiverType(receiverIndex)!!.descriptor,
+                null,
+                null
+            )
         }
     }
 

compiler/backend/src/org/jetbrains/kotlin/codegen/context/ScriptContext.kt

@@ -16,21 +16,30 @@
 
 package org.jetbrains.kotlin.codegen.context
 
+import org.jetbrains.kotlin.codegen.AsmUtil
 import org.jetbrains.kotlin.codegen.FieldInfo
 import org.jetbrains.kotlin.codegen.OwnerKind
-import org.jetbrains.kotlin.codegen.state.GenerationState
+import org.jetbrains.kotlin.codegen.StackValue
 import org.jetbrains.kotlin.codegen.state.KotlinTypeMapper
 import org.jetbrains.kotlin.descriptors.ClassDescriptor
 import org.jetbrains.kotlin.descriptors.ScriptDescriptor
+import org.jetbrains.kotlin.descriptors.findClassAcrossModuleDependencies
+import org.jetbrains.kotlin.name.ClassId
+import org.jetbrains.kotlin.name.FqName
+import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.psi.KtAnonymousInitializer
-import org.jetbrains.kotlin.psi.KtDeclaration
 import org.jetbrains.kotlin.psi.KtExpression
 import org.jetbrains.kotlin.psi.KtScript
 import org.jetbrains.kotlin.resolve.DescriptorToSourceUtils
+import org.jetbrains.kotlin.resolve.descriptorUtil.fqNameSafe
+import org.jetbrains.kotlin.resolve.descriptorUtil.module
 import org.jetbrains.kotlin.resolve.jvm.AsmTypes
+import org.jetbrains.kotlin.resolve.lazy.descriptors.LazyScriptDescriptor
+import org.jetbrains.org.objectweb.asm.Type
+import kotlin.reflect.KClass
 
 class ScriptContext(
-    typeMapper: KotlinTypeMapper,
+    val typeMapper: KotlinTypeMapper,
     val scriptDescriptor: ScriptDescriptor,
     val earlierScripts: List<ScriptDescriptor>,
     contextDescriptor: ClassDescriptor,
@@ -41,14 +50,14 @@ class ScriptContext(
     val resultFieldInfo: FieldInfo
         get() {
             assert(state.replSpecific.shouldGenerateScriptResultValue) { "Should not be called unless 'scriptResultFieldName' is set" }
-            val state = state
             val scriptResultFieldName = state.replSpecific.scriptResultFieldName!!
             return FieldInfo.createForHiddenField(state.typeMapper.mapClass(scriptDescriptor), AsmTypes.OBJECT_TYPE, scriptResultFieldName)
         }
 
+    val script = DescriptorToSourceUtils.getSourceFromDescriptor(scriptDescriptor) as KtScript?
+            ?: error("Declaration should be present for script: $scriptDescriptor")
+
     init {
-        val script = DescriptorToSourceUtils.getSourceFromDescriptor(scriptDescriptor) as KtScript?
-                ?: error("Declaration should be present for script: $scriptDescriptor")
         val lastDeclaration = script.declarations.lastOrNull()
         if (lastDeclaration is KtAnonymousInitializer) {
             this.lastStatement = lastDeclaration.body
@@ -57,6 +66,29 @@ class ScriptContext(
         }
     }
 
+    fun getImplicitReceiverName(index: Int): String = "\$\$implicitReceiver$index"
+
+    fun getImplicitReceiverType(index: Int): Type? {
+        val receivers = script.kotlinScriptDefinition.value.implicitReceivers
+        val kClass = receivers.getOrNull(index)?.classifier as? KClass<*>
+        return kClass?.java?.classId?.let(AsmUtil::asmTypeByClassId)
+    }
+
+    fun getOuterReceiverExpression(prefix: StackValue?, thisOrOuterClass: ClassDescriptor): StackValue {
+        receiverDescriptors.forEachIndexed { index, outerReceiver ->
+            if (outerReceiver == thisOrOuterClass) {
+                return getImplicitReceiverType(index)?.let { type ->
+                    val owner = typeMapper.mapType(scriptDescriptor)
+                    StackValue.field(type, owner, getImplicitReceiverName(index), false, prefix ?: StackValue.LOCAL_0)
+                } ?: error("Invalid script receiver: ${thisOrOuterClass.fqNameSafe}")
+            }
+        }
+        error("Script receiver not found: ${thisOrOuterClass.fqNameSafe}")
+    }
+
+    val receiverDescriptors: List<ClassDescriptor>
+        get() = scriptDescriptor.implicitReceivers
+
     fun getScriptFieldName(scriptDescriptor: ScriptDescriptor): String {
         val index = earlierScripts.indexOf(scriptDescriptor)
         if (index < 0) {
@@ -69,3 +101,6 @@ class ScriptContext(
         return "Script: " + contextDescriptor.name.asString()
     }
 }
+
+private val Class<*>.classId: ClassId
+    get() = enclosingClass?.classId?.createNestedClassId(Name.identifier(simpleName)) ?: ClassId.topLevel(FqName(name))

compiler/backend/src/org/jetbrains/kotlin/codegen/state/KotlinTypeMapper.java

@@ -10,6 +10,7 @@ import com.intellij.psi.PsiElement;
 import kotlin.Pair;
 import kotlin.Unit;
 import kotlin.collections.CollectionsKt;
+import kotlin.reflect.KType;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 import org.jetbrains.kotlin.builtins.BuiltInsPackageFragment;
@@ -1563,7 +1564,11 @@ public class KotlinTypeMapper {
     }
 
     @NotNull
-    public JvmMethodSignature mapScriptSignature(@NotNull ScriptDescriptor script, @NotNull List<ScriptDescriptor> importedScripts) {
+    public JvmMethodSignature mapScriptSignature(
+            @NotNull ScriptDescriptor script,
+            @NotNull List<ScriptDescriptor> importedScripts,
+            List<? extends KType> implicitReceivers
+    ) {
         JvmSignatureWriter sw = new BothSignatureWriter(BothSignatureWriter.Mode.METHOD);
 
         sw.writeParametersStart();
@@ -1572,6 +1577,10 @@ public class KotlinTypeMapper {
             writeParameter(sw, DescriptorUtilsKt.getModule(script).getBuiltIns().getArray().getDefaultType(), null);
         }
 
+        if (implicitReceivers.size() > 0) {
+            writeParameter(sw, DescriptorUtilsKt.getModule(script).getBuiltIns().getArray().getDefaultType(), null);
+        }
+
         for (ValueParameterDescriptor valueParameter : script.getUnsubstitutedPrimaryConstructor().getValueParameters()) {
             writeParameter(sw, valueParameter.getType(), /* callableDescriptor = */ null);
         }

compiler/frontend/src/org/jetbrains/kotlin/diagnostics/Errors.java

@@ -98,6 +98,7 @@ public interface Errors {
     DiagnosticFactory2<PsiElement, String, String> API_NOT_AVAILABLE = DiagnosticFactory2.create(ERROR);
 
     DiagnosticFactory1<PsiElement, FqName> MISSING_DEPENDENCY_CLASS = DiagnosticFactory1.create(ERROR);
+    DiagnosticFactory1<PsiElement, FqName> MISSING_SCRIPT_RECEIVER_CLASS = DiagnosticFactory1.create(ERROR);
     DiagnosticFactory1<PsiElement, String> PRE_RELEASE_CLASS = DiagnosticFactory1.create(ERROR);
     DiagnosticFactory2<PsiElement, String, IncompatibleVersionErrorData<?>> INCOMPATIBLE_CLASS = DiagnosticFactory2.create(ERROR);
 

compiler/frontend/src/org/jetbrains/kotlin/diagnostics/rendering/DefaultErrorMessages.java

@@ -363,6 +363,7 @@ public class DefaultErrorMessages {
         MAP.put(API_NOT_AVAILABLE, "This declaration is only available since Kotlin {0} and cannot be used with the specified API version {1}", STRING, STRING);
 
         MAP.put(MISSING_DEPENDENCY_CLASS, "Cannot access class ''{0}''. Check your module classpath for missing or conflicting dependencies", TO_STRING);
+        MAP.put(MISSING_SCRIPT_RECEIVER_CLASS, "Cannot access implicit script receiver class ''{0}''. Check your module classpath for missing or conflicting dependencies", TO_STRING);
         MAP.put(PRE_RELEASE_CLASS, "{0} is compiled by a pre-release version of Kotlin and cannot be loaded by this version of the compiler", TO_STRING);
         MAP.put(INCOMPATIBLE_CLASS,
                 "{0} was compiled with an incompatible version of Kotlin. {1}",

compiler/frontend/src/org/jetbrains/kotlin/resolve/lazy/descriptors/LazyScriptDescriptor.kt

@@ -16,10 +16,7 @@
 
 package org.jetbrains.kotlin.resolve.lazy.descriptors
 
-import org.jetbrains.kotlin.descriptors.DeclarationDescriptor
+import org.jetbrains.kotlin.descriptors.*
-import org.jetbrains.kotlin.descriptors.DeclarationDescriptorVisitor
-import org.jetbrains.kotlin.descriptors.ScriptDescriptor
-import org.jetbrains.kotlin.descriptors.findClassAcrossModuleDependencies
 import org.jetbrains.kotlin.diagnostics.Errors
 import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.name.FqName
@@ -92,16 +89,26 @@ class LazyScriptDescriptor(
     override fun computeSupertypes() =
         listOf(ScriptHelper.getInstance().getKotlinType(this, scriptDefinition.template)).ifEmpty { listOf(builtIns.anyType) }
 
-    private val scriptOuterScope: () -> LexicalScope = resolveSession.storageManager.createLazyValue {
+    private val scriptImplicitReceivers: () -> List<ClassDescriptor> = resolveSession.storageManager.createLazyValue {
-        val receivers = scriptDefinition.implicitReceivers
+        scriptDefinition.implicitReceivers.mapNotNull { receiver ->
-        var outerScope = super.getOuterScope()
-        for (receiver in receivers.asReversed()) {
             val receiverClassId = receiver.classifier?.let { it as? KClass<*> }?.java?.classId
-            val receiverClassDescriptor = receiverClassId?.let { module.findClassAcrossModuleDependencies(it) }
+            receiverClassId?.let {
-            if (receiverClassDescriptor == null) {
+                module.findClassAcrossModuleDependencies(it)
-                resolveSession.trace.report(Errors.MISSING_DEPENDENCY_CLASS.on(scriptInfo.script, receiverClassId?.asSingleFqName() ?: FqName(receiver.toString())))
+                        ?: also {
-                break
+                            // TODO: use PositioningStrategies to highlight some specific place in case of error, instead of treating the whole file as invalid
+                            resolveSession.trace.report(
+                                Errors.MISSING_SCRIPT_RECEIVER_CLASS.on(scriptInfo.script, receiverClassId.asSingleFqName())
+                            )
+                        }
             }
+        }
+    }
+
+    override fun getImplicitReceivers(): List<ClassDescriptor> = scriptImplicitReceivers()
+
+    private val scriptOuterScope: () -> LexicalScope = resolveSession.storageManager.createLazyValue {
+        var outerScope = super.getOuterScope()
+        for (receiverClassDescriptor in implicitReceivers.asReversed()) {
             outerScope = LexicalScopeImpl(
                 outerScope,
                 receiverClassDescriptor,

compiler/psi/src/org/jetbrains/kotlin/psi/KtScript.java

@@ -33,7 +33,7 @@ import java.util.Objects;
 import static kotlin.LazyThreadSafetyMode.PUBLICATION;
 
 public class KtScript extends KtNamedDeclarationStub<KotlinScriptStub> implements KtDeclarationContainer {
-    private final Lazy<KotlinScriptDefinition> kotlinScriptDefinition = LazyKt.lazy(PUBLICATION, () -> Objects.requireNonNull(
+    public final Lazy<KotlinScriptDefinition> kotlinScriptDefinition = LazyKt.lazy(PUBLICATION, () -> Objects.requireNonNull(
             KotlinScriptDefinitionProviderKt.getScriptDefinition(getContainingKtFile()),
             () -> "Should not parse a script without definition: " + getContainingKtFile().getVirtualFile().getPath()
     ));

core/descriptors/src/org/jetbrains/kotlin/descriptors/ScriptDescriptor.java

@@ -18,10 +18,15 @@ package org.jetbrains.kotlin.descriptors;
 
 import org.jetbrains.annotations.NotNull;
 
+import java.util.List;
+
 public interface ScriptDescriptor extends ClassDescriptor {
     int getPriority();
 
     @NotNull
     @Override
     ClassConstructorDescriptor getUnsubstitutedPrimaryConstructor();
+
+    @NotNull
+    List<ClassDescriptor> getImplicitReceivers();
 }

libraries/scripting/jvm/src/kotlin/script/experimental/jvm/runners/BasicJvmScriptRunner.kt

@@ -26,7 +26,12 @@ open class BasicJvmScriptRunner<ScriptBase : Any>(val baseClass: KClass<ScriptBa
                     if (scriptObject !is Class<*>)
                         ResultWithDiagnostics.Failure(ScriptDiagnostic("expecting class in this implementation, got ${scriptObject?.javaClass}"))
                     else {
-                        scriptObject.getConstructor().newInstance()
+                        val receivers = scriptEvaluationEnvironment.getOrNull(ScriptEvaluationEnvironmentParams.implicitReceivers)
+                        if (receivers == null) {
+                            scriptObject.getConstructor().newInstance()
+                        } else {
+                            scriptObject.getConstructor(Array<Any>::class.java).newInstance(receivers.toTypedArray())
+                        }
 
                         ResultWithDiagnostics.Success(EvaluationResult(null, scriptEvaluationEnvironment))
                     }