Implement script environment variables in backend, add simple test

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

@@ -32,9 +32,11 @@ import org.jetbrains.kotlin.resolve.annotations.AnnotationUtilKt;
 import org.jetbrains.kotlin.resolve.calls.model.ResolvedCall;
 import org.jetbrains.kotlin.resolve.calls.util.UnderscoreUtilKt;
 import org.jetbrains.kotlin.resolve.constants.ConstantValue;
+import org.jetbrains.kotlin.resolve.jvm.AsmTypes;
 import org.jetbrains.kotlin.resolve.jvm.diagnostics.JvmDeclarationOriginKt;
 import org.jetbrains.kotlin.resolve.jvm.jvmSignature.JvmMethodGenericSignature;
 import org.jetbrains.kotlin.resolve.jvm.jvmSignature.JvmMethodSignature;
+import org.jetbrains.kotlin.resolve.lazy.descriptors.ScriptEnvironmentPropertyDescriptor;
 import org.jetbrains.kotlin.serialization.deserialization.descriptors.DeserializedPropertyDescriptor;
 import org.jetbrains.kotlin.storage.LockBasedStorageManager;
 import org.jetbrains.kotlin.types.ErrorUtils;
@@ -486,13 +488,13 @@ public class PropertyCodegen {
         return false;
     }
 
-    private void generateGetter(@Nullable KtNamedDeclaration p, @NotNull PropertyDescriptor descriptor, @Nullable KtPropertyAccessor getter) {
+    public void generateGetter(@Nullable KtNamedDeclaration p, @NotNull PropertyDescriptor descriptor, @Nullable KtPropertyAccessor getter) {
         generateAccessor(p, getter, descriptor.getGetter() != null
                                     ? descriptor.getGetter()
                                     : DescriptorFactory.createDefaultGetter(descriptor, Annotations.Companion.getEMPTY()));
     }
 
-    private void generateSetter(@Nullable KtNamedDeclaration p, @NotNull PropertyDescriptor descriptor, @Nullable KtPropertyAccessor setter) {
+    public void generateSetter(@Nullable KtNamedDeclaration p, @NotNull PropertyDescriptor descriptor, @Nullable KtPropertyAccessor setter) {
         if (!descriptor.isVar()) return;
 
         generateAccessor(p, setter, descriptor.getSetter() != null
@@ -513,7 +515,10 @@ public class PropertyCodegen {
 
         FunctionGenerationStrategy strategy;
         if (accessor == null || !accessor.hasBody()) {
-            if (p instanceof KtProperty && ((KtProperty) p).hasDelegate()) {
+            if (accessorDescriptor.getCorrespondingProperty() instanceof ScriptEnvironmentPropertyDescriptor) {
+                strategy = new ScriptEnvPropertyAccessorStrategy(state, accessorDescriptor);
+            }
+            else if (p instanceof KtProperty && ((KtProperty) p).hasDelegate()) {
                 strategy = new DelegatedPropertyAccessorStrategy(state, accessorDescriptor);
             }
             else {
@@ -626,6 +631,42 @@ public class PropertyCodegen {
         }
     }
 
+    private static class ScriptEnvPropertyAccessorStrategy extends FunctionGenerationStrategy.CodegenBased {
+        public static final String MAP_IFACE_NAME = "java/util/Map";
+        public static final String MAP_FIELD_NAME = "environment";
+        private final PropertyAccessorDescriptor propertyAccessorDescriptor;
+
+        public ScriptEnvPropertyAccessorStrategy(@NotNull GenerationState state, @NotNull PropertyAccessorDescriptor descriptor) {
+            super(state);
+            this.propertyAccessorDescriptor = descriptor;
+        }
+
+        @Override
+        public void doGenerateBody(@NotNull ExpressionCodegen codegen, @NotNull JvmMethodSignature signature) {
+            InstructionAdapter v = codegen.v;
+
+            String intScriptName =
+                    state.getTypeMapper().mapClass((ClassifierDescriptor)(propertyAccessorDescriptor.getCorrespondingProperty().getContainingDeclaration())).getClassName();
+            v.visitVarInsn(Opcodes.ALOAD, 0);
+            v.visitFieldInsn(Opcodes.GETFIELD, intScriptName, MAP_FIELD_NAME, Type.getObjectType(MAP_IFACE_NAME).getDescriptor());
+            v.visitLdcInsn(propertyAccessorDescriptor.getCorrespondingProperty().getName().asString());
+            if (propertyAccessorDescriptor instanceof PropertyGetterDescriptor) {
+                v.visitMethodInsn(Opcodes.INVOKEINTERFACE, MAP_IFACE_NAME, "get",
+                                  Type.getMethodDescriptor(AsmTypes.OBJECT_TYPE, AsmTypes.OBJECT_TYPE), true);
+            }
+            else {
+                Type valueType = state.getTypeMapper().mapType(propertyAccessorDescriptor.getCorrespondingProperty());
+                v.load(1, valueType);
+                StackValue.coerce(valueType, AsmTypes.OBJECT_TYPE, v);
+                v.visitMethodInsn(Opcodes.INVOKEINTERFACE, MAP_IFACE_NAME, "set",
+                                  Type.getMethodDescriptor(Type.BOOLEAN_TYPE, AsmTypes.OBJECT_TYPE, AsmTypes.OBJECT_TYPE), true);
+            }
+            Type returnType = state.getTypeMapper().mapReturnType(propertyAccessorDescriptor);
+            StackValue.coerce(AsmTypes.OBJECT_TYPE, returnType, v);
+            v.areturn(returnType);
+        }
+    }
+
     public void genDelegate(@NotNull PropertyDescriptor delegate, @NotNull PropertyDescriptor delegateTo, @NotNull StackValue field) {
         ClassDescriptor toClass = (ClassDescriptor) delegateTo.getContainingDeclaration();
 

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

@@ -15,6 +15,7 @@ import org.jetbrains.kotlin.descriptors.ScriptDescriptor
 import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.resolve.BindingContext
 import org.jetbrains.kotlin.resolve.descriptorUtil.*
+import org.jetbrains.kotlin.resolve.jvm.AsmTypes
 import org.jetbrains.kotlin.resolve.jvm.diagnostics.JvmDeclarationOrigin
 import org.jetbrains.kotlin.resolve.jvm.diagnostics.*
 import org.jetbrains.org.objectweb.asm.Type
@@ -53,6 +54,10 @@ class ScriptCodegen private constructor(
 
     override fun generateSyntheticPartsBeforeBody() {
         generatePropertyMetadataArrayFieldIfNeeded(classAsmType)
+        scriptContext.scriptDescriptor.scriptEnvironmentProperties.forEach {
+            propertyCodegen.generateGetter(null, it, null)
+            propertyCodegen.generateSetter(null, it, null)
+        }
     }
 
     override fun generateSyntheticPartsAfterBody() {}
@@ -68,7 +73,12 @@ class ScriptCodegen private constructor(
     ) {
         val scriptDefinition = scriptContext.script.kotlinScriptDefinition.value
 
-        val jvmSignature = typeMapper.mapScriptSignature(scriptDescriptor, scriptContext.earlierScripts, scriptDefinition.implicitReceivers)
+        val jvmSignature = typeMapper.mapScriptSignature(
+            scriptDescriptor,
+            scriptContext.earlierScripts,
+            scriptDefinition.implicitReceivers,
+            scriptDefinition.environmentVariables
+        )
 
         if (state.replSpecific.shouldGenerateScriptResultValue) {
             val resultFieldInfo = scriptContext.resultFieldInfo
@@ -108,6 +118,7 @@ class ScriptCodegen private constructor(
                 val valueParamStart = 1
                     .incrementIf(scriptContext.earlierScripts.isNotEmpty())
                     .incrementIf(scriptDefinition.implicitReceivers.isNotEmpty())
+                    .incrementIf(scriptDefinition.environmentVariables.isNotEmpty())
 
                 val valueParameters = scriptDescriptor.unsubstitutedPrimaryConstructor.valueParameters
                 for (superclassParam in ctorDesc.valueParameters) {
@@ -155,6 +166,14 @@ class ScriptCodegen private constructor(
                 }
             }
 
+            if (scriptDefinition.environmentVariables.isNotEmpty()) {
+                val envParamIndex = frameMap.enterTemp(AsmTypes.OBJECT_TYPE)
+                val mapType = Type.getObjectType("java/util/Map")
+                iv.load(0, classType)
+                iv.load(envParamIndex, mapType)
+                iv.putfield(classType.internalName, "environment", mapType.descriptor)
+            }
+
             val codegen = ExpressionCodegen(mv, frameMap, Type.VOID_TYPE, methodContext, state, this)
 
             generateInitializers { codegen }
@@ -187,6 +206,16 @@ class ScriptCodegen private constructor(
                 null
             )
         }
+        if (scriptContext.scriptDescriptor.scriptEnvironmentProperties.isNotEmpty()) {
+            classBuilder.newField(
+                NO_ORIGIN,
+                ACC_PUBLIC or ACC_FINAL,
+                "environment",
+                Type.getObjectType("java/util/Map").descriptor,
+                null,
+                null
+            )
+        }
     }
 
     private fun genMembers() {

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

@@ -23,7 +23,6 @@ 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
@@ -32,9 +31,8 @@ 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.kotlin.resolve.lazy.descriptors.ScriptEnvironmentDescriptor
 import org.jetbrains.org.objectweb.asm.Type
 import kotlin.reflect.KClass
 
@@ -75,6 +73,9 @@ class ScriptContext(
     }
 
     fun getOuterReceiverExpression(prefix: StackValue?, thisOrOuterClass: ClassDescriptor): StackValue {
+        if (thisOrOuterClass is ScriptEnvironmentDescriptor) {
+            return prefix ?: StackValue.LOCAL_0
+        }
         receiverDescriptors.forEachIndexed { index, outerReceiver ->
             if (outerReceiver == thisOrOuterClass) {
                 return getImplicitReceiverType(index)?.let { type ->

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

@@ -1586,7 +1586,8 @@ public class KotlinTypeMapper {
     public JvmMethodSignature mapScriptSignature(
             @NotNull ScriptDescriptor script,
             @NotNull List<ScriptDescriptor> importedScripts,
-            List<? extends KType> implicitReceivers
+            List<? extends KType> implicitReceivers,
+            List<? extends Pair<String, ? extends KType>> environmentVariables
     ) {
         JvmSignatureWriter sw = new BothSignatureWriter(BothSignatureWriter.Mode.METHOD);
 
@@ -1600,6 +1601,10 @@ public class KotlinTypeMapper {
             writeParameter(sw, DescriptorUtilsKt.getModule(script).getBuiltIns().getArray().getDefaultType(), null);
         }
 
+        if (environmentVariables.size() > 0) {
+            writeParameter(sw, DescriptorUtilsKt.getModule(script).getBuiltIns().getMap().getDefaultType(), null);
+        }
+
         for (ValueParameterDescriptor valueParameter : script.getUnsubstitutedPrimaryConstructor().getValueParameters()) {
             writeParameter(sw, valueParameter.getType(), /* callableDescriptor = */ null);
         }

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

@@ -102,7 +102,7 @@ class LazyScriptDescriptor(
         }
     }
 
-    private fun findTypeDescriptor(type: KType, errorDiagnostic: DiagnosticFactory1<PsiElement, FqName>): ClassDescriptor? {
+    internal fun findTypeDescriptor(type: KType, errorDiagnostic: DiagnosticFactory1<PsiElement, FqName>): ClassDescriptor? {
         val receiverClassId = type.classifier?.let { it as? KClass<*> }?.classId
         return receiverClassId?.let {
             module.findClassAcrossModuleDependencies(it)
@@ -119,15 +119,11 @@ class LazyScriptDescriptor(
 
     override fun getImplicitReceivers(): List<ClassDescriptor> = scriptImplicitReceivers()
 
-    private val scriptEnvironmentProperties: () -> List<Pair<String, ClassDescriptor>> = resolveSession.storageManager.createLazyValue {
+    private val scriptEnvironment: () -> ScriptEnvironmentDescriptor = resolveSession.storageManager.createLazyValue {
-        scriptDefinition.environmentVariables.mapNotNull { (name, type) ->
+        ScriptEnvironmentDescriptor(this)
-            findTypeDescriptor(type, Errors.MISSING_SCRIPT_ENVIRONMENT_PROPERTY_CLASS)?.let {
-                name to it
-            }
-        }
     }
 
-    fun getScriptEnvironmentProperties(): List<Pair<String, ClassDescriptor>> = scriptEnvironmentProperties()
+    override fun getScriptEnvironmentProperties(): List<PropertyDescriptor> = scriptEnvironment().properties
 
     private val scriptOuterScope: () -> LexicalScope = resolveSession.storageManager.createLazyValue {
         var outerScope = super.getOuterScope()
@@ -160,7 +156,7 @@ private val KClass<*>.classId: ClassId
 private val KType.classId: ClassId?
     get() = classifier?.let { it as? KClass<*> }?.classId
 
-private class ScriptEnvironmentDescriptor(script: LazyScriptDescriptor) :
+class ScriptEnvironmentDescriptor(script: LazyScriptDescriptor) :
     MutableClassDescriptor(
         script, ClassKind.CLASS, false, false,
         Name.special("<synthetic script environment for ${script.name}>"), SourceElement.NO_SOURCE, LockBasedStorageManager.NO_LOCKS
@@ -176,23 +172,32 @@ private class ScriptEnvironmentDescriptor(script: LazyScriptDescriptor) :
     private val memberScope by lazy {
         ScriptEnvironmentMemberScope(
             script.name.identifier,
-            script.getScriptEnvironmentProperties().map {
+            properties
-                makeEnvironmentPropertyDescriptor(
-                    Name.identifier(it.first),
-                    it.second,
-                    true,
-                    script
-                )
-            }
         )
     }
 
     override fun getUnsubstitutedMemberScope(): MemberScope = memberScope
+
+    val properties by lazy {
+        script.scriptDefinition.environmentVariables.mapNotNull { (name, type) ->
+            script.findTypeDescriptor(type, Errors.MISSING_SCRIPT_ENVIRONMENT_PROPERTY_CLASS)?.let {
+                name to it
+            }
+        }.map { (key, value) ->
+            ScriptEnvironmentPropertyDescriptor(
+                Name.identifier(key),
+                value,
+                thisAsReceiverParameter,
+                true,
+                script
+            )
+        }
+    }
 }
 
 private class ScriptEnvironmentMemberScope(
     private val scriptId: String,
-    private val environmentProperties: List<PropertyDescriptorImpl>
+    private val environmentProperties: List<PropertyDescriptor>
 ) : MemberScopeImpl() {
     override fun getContributedDescriptors(
         kindFilter: DescriptorKindFilter,
@@ -208,27 +213,35 @@ private class ScriptEnvironmentMemberScope(
     }
 }
 
-private fun ScriptEnvironmentDescriptor.makeEnvironmentPropertyDescriptor(name: Name, typeDescriptor: ClassDescriptor, isVar: Boolean, script: LazyScriptDescriptor) =
+class ScriptEnvironmentPropertyDescriptor(
-    PropertyDescriptorImpl.create(
+    name: Name,
-        script.containingDeclaration,
+    typeDescriptor: ClassDescriptor,
-        Annotations.EMPTY, Modality.FINAL, Visibilities.PUBLIC,
+    receiver: ReceiverParameterDescriptor?,
-        isVar,
+    isVar: Boolean,
-        name,
+    script: LazyScriptDescriptor
-        CallableMemberDescriptor.Kind.SYNTHESIZED,
+) : PropertyDescriptorImpl(
-        SourceElement.NO_SOURCE,
+    script,
-        /* lateInit = */ false, /* isConst = */ false, /* isExpect = */ false, /* isActual = */ false, /* isExternal = */ false,
+    null,
-        /* isDelegated = */ true
+    Annotations.EMPTY, Modality.OPEN, Visibilities.PUBLIC,
-    ).also {
+    isVar,
-        it.setType(typeDescriptor.defaultType, emptyList<TypeParameterDescriptorImpl>(), thisAsReceiverParameter, null as KotlinType?)
+    name,
-        it.initialize(
+    CallableMemberDescriptor.Kind.SYNTHESIZED,
-            it.makePropertyGetterDescriptor(),
+    SourceElement.NO_SOURCE,
-            if (!isVar) null else it.makePropertySetterDescriptor()
+    /* lateInit = */ false, /* isConst = */ false, /* isExpect = */ false, /* isActual = */ false, /* isExternal = */ false,
+    /* isDelegated = */ true
+) {
+    init {
+        setType(typeDescriptor.defaultType, emptyList<TypeParameterDescriptorImpl>(), receiver, null as KotlinType?)
+        initialize(
+            makePropertyGetterDescriptor(),
+            if (!isVar) null else makePropertySetterDescriptor()
         )
     }
+}
 
 private fun PropertyDescriptorImpl.makePropertyGetterDescriptor() =
     PropertyGetterDescriptorImpl(
-        this, Annotations.EMPTY, Modality.FINAL, Visibilities.PUBLIC,
+        this, Annotations.EMPTY, Modality.OPEN, Visibilities.PUBLIC,
         /* isDefault = */ false, /* isExternal = */ false, /* isInline = */ false,
         CallableMemberDescriptor.Kind.SYNTHESIZED, null, SourceElement.NO_SOURCE
     ).also {
@@ -237,7 +250,7 @@ private fun PropertyDescriptorImpl.makePropertyGetterDescriptor() =
 
 private fun PropertyDescriptorImpl.makePropertySetterDescriptor() =
     PropertySetterDescriptorImpl(
-        this, Annotations.EMPTY, Modality.FINAL, Visibilities.PUBLIC,
+        this, Annotations.EMPTY, Modality.OPEN, Visibilities.PUBLIC,
         /* isDefault = */ false, /* isExternal = */ false, /* isInline = */ false,
         CallableMemberDescriptor.Kind.SYNTHESIZED, null, SourceElement.NO_SOURCE
     ).also {

compiler/testData/codegen/customScript/simpleEnvVars.kts

@@ -0,0 +1,8 @@
+
+// KOTLIN_SCRIPT_DEFINITION: org.jetbrains.kotlin.codegen.TestScriptWithSimpleEnvVars
+
+// envVar: stringVar1 = abracadabra
+
+val res = stringVar1.drop(4)
+
+// expected: res = cadabra

compiler/tests-common/tests/org/jetbrains/kotlin/codegen/AbstractCustomScriptCodegenTest.kt

@@ -212,3 +212,13 @@ object TestScriptWithReceiversConfiguration : ArrayList<Pair<TypedKey<*>, Any?>>
 @KotlinScriptDefaultCompilationConfiguration(TestScriptWithReceiversConfiguration::class)
 abstract class TestScriptWithReceivers
 
+object TestScriptWithSimpleEnvVarsConfiguration : ArrayList<Pair<TypedKey<*>, Any?>>(
+    listOf(
+        ScriptCompileConfigurationProperties.contextVariables("stringVar1" to String::class.starProjectedType)
+    )
+)
+
+@Suppress("unused")
+@KotlinScript
+@KotlinScriptDefaultCompilationConfiguration(TestScriptWithSimpleEnvVarsConfiguration::class)
+abstract class TestScriptWithSimpleEnvVars

compiler/tests/org/jetbrains/kotlin/codegen/CustomScriptCodegenTestGenerated.java

@@ -29,6 +29,11 @@ public class CustomScriptCodegenTestGenerated extends AbstractCustomScriptCodege
         KotlinTestUtils.assertAllTestsPresentByMetadata(this.getClass(), new File("compiler/testData/codegen/customScript"), Pattern.compile("^(.+)\\.kts$"), TargetBackend.ANY, true);
     }
 
+    @TestMetadata("simpleEnvVars.kts")
+    public void testSimpleEnvVars() throws Exception {
+        runTest("compiler/testData/codegen/customScript/simpleEnvVars.kts");
+    }
+
     @TestMetadata("stringReceiver.kts")
     public void testStringReceiver() throws Exception {
         runTest("compiler/testData/codegen/customScript/stringReceiver.kts");

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

@@ -29,4 +29,7 @@ public interface ScriptDescriptor extends ClassDescriptor {
 
     @NotNull
     List<ClassDescriptor> getImplicitReceivers();
+
+    @NotNull
+    List<PropertyDescriptor> getScriptEnvironmentProperties();
 }