Add support for script base class with parameters

including possibility to describe which script params to pass to superclass

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

@@ -16,23 +16,22 @@
 
 package org.jetbrains.kotlin.codegen;
 
-import com.intellij.util.ArrayUtil;
+import kotlin.Pair;
 import kotlin.jvm.functions.Function0;
 import org.jetbrains.annotations.NotNull;
 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.ConstructorDescriptor;
-import org.jetbrains.kotlin.descriptors.ScriptDescriptor;
-import org.jetbrains.kotlin.descriptors.ValueParameterDescriptor;
+import org.jetbrains.kotlin.descriptors.*;
+import org.jetbrains.kotlin.name.Name;
 import org.jetbrains.kotlin.psi.*;
 import org.jetbrains.kotlin.resolve.BindingContext;
 import org.jetbrains.kotlin.resolve.descriptorUtil.DescriptorUtilsKt;
 import org.jetbrains.kotlin.resolve.jvm.diagnostics.JvmDeclarationOrigin;
 import org.jetbrains.kotlin.resolve.jvm.diagnostics.JvmDeclarationOriginKt;
 import org.jetbrains.kotlin.resolve.jvm.jvmSignature.JvmMethodSignature;
+import org.jetbrains.kotlin.types.KotlinType;
 import org.jetbrains.org.objectweb.asm.MethodVisitor;
 import org.jetbrains.org.objectweb.asm.Type;
 import org.jetbrains.org.objectweb.asm.commons.InstructionAdapter;
@@ -147,13 +146,40 @@ public class ScriptCodegen extends MemberCodegen<KtScript> {
 
             Type classType = typeMapper.mapType(scriptDescriptor);
 
-            iv.load(0, classType);
-            ClassDescriptor superclass = DescriptorUtilsKt.getSuperClassOrAny(scriptDescriptor);
-            assert hasDefaultConstructor(superclass);
-            iv.invokespecial(
-                    typeMapper.mapSupertype(superclass.getDefaultType(), null).getInternalName(),
-                    "<init>", "()V", false);
+            ClassDescriptor superclass = DescriptorUtilsKt.getSuperClassNotAny(scriptDescriptor);
 
+            List<ValueParameterDescriptor> valueParameters = scriptDescriptor.getUnsubstitutedPrimaryConstructor().getValueParameters();
+
+            if (superclass == null) {
+                iv.load(0, classType);
+                iv.invokespecial("java/lang/Object", "<init>", "()V", false);
+            }
+            else {
+                List<Pair<Name, KotlinType>> superclassParamsMap = scriptDescriptor.getScriptParametersToPassToSuperclass();
+                ConstructorDescriptor ctorDesc = DescriptorUtilsKt.getConstructorByParamsMap(superclass, superclassParamsMap);
+                assert ctorDesc != null;
+
+                iv.load(0, classType);
+
+                for (Pair<Name, KotlinType> superclassParam: superclassParamsMap) {
+                    ValueParameterDescriptor valueParam = null;
+                    for (ValueParameterDescriptor vpd: valueParameters) {
+                        if (vpd.getName().equals(superclassParam.getFirst())) {
+                            valueParam = vpd;
+                            break;
+                        }
+                    }
+                    assert valueParam != null;
+                    iv.load(valueParam.getIndex() + 1, typeMapper.mapType(valueParam.getType()));
+                }
+
+                CallableMethod ctorMethod = typeMapper.mapToCallableMethod(ctorDesc, false);
+                String sig = ctorMethod.getAsmMethod().getDescriptor();
+
+                iv.invokespecial(
+                        typeMapper.mapSupertype(superclass.getDefaultType(), null).getInternalName(),
+                        "<init>", sig, false);
+            }
             iv.load(0, classType);
 
             FrameMap frameMap = new FrameMap();
@@ -166,7 +192,6 @@ public class ScriptCodegen extends MemberCodegen<KtScript> {
             Type[] argTypes = jvmSignature.getAsmMethod().getArgumentTypes();
             int add = 0;
 
-            List<ValueParameterDescriptor> valueParameters = scriptDescriptor.getUnsubstitutedPrimaryConstructor().getValueParameters();
             for (int i = 0; i < valueParameters.size(); i++) {
                 ValueParameterDescriptor parameter = valueParameters.get(i);
                 frameMap.enter(parameter, argTypes[i + add]);
@@ -206,15 +231,6 @@ public class ScriptCodegen extends MemberCodegen<KtScript> {
         mv.visitEnd();
     }
 
-    private boolean hasDefaultConstructor(ClassDescriptor superclass) {
-        for (ConstructorDescriptor ctor : superclass.getConstructors()) {
-            if (ctor.getValueParameters().isEmpty()) {
-                return true;
-            }
-        }
-        return false;
-    }
-
     private void genFieldsForParameters(@NotNull ScriptDescriptor script, @NotNull ClassBuilder classBuilder) {
         for (ScriptDescriptor earlierScript : context.getEarlierScripts()) {
             Type earlierClassName = typeMapper.mapType(earlierScript);

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

@@ -27,6 +27,7 @@ import org.jetbrains.kotlin.resolve.lazy.data.KtScriptInfo
 import org.jetbrains.kotlin.resolve.lazy.declarations.ClassMemberDeclarationProvider
 import org.jetbrains.kotlin.resolve.source.toSourceElement
 import org.jetbrains.kotlin.script.KotlinScriptDefinitionProvider
+import org.jetbrains.kotlin.script.ScriptParameter
 import org.jetbrains.kotlin.script.ScriptPriorities
 import org.jetbrains.kotlin.types.KotlinType
 import org.jetbrains.kotlin.types.TypeSubstitutor
@@ -77,7 +78,18 @@ class LazyScriptDescriptor(
     override fun getInjectedSupertypes(): List<KotlinType>? {
         val file = scriptInfo.script.getContainingKtFile()
         val scriptDefinition = KotlinScriptDefinitionProvider.getInstance(file.project).findScriptDefinition(file)
-        val superclasses = scriptDefinition.getScriptSuperclasses(this)
+        val superclasses = scriptDefinition.getScriptSupertypes(this)
         return if (superclasses.isEmpty()) null else superclasses
     }
+
+    override fun getScriptParametersToPassToSuperclass(): List<Pair<Name, KotlinType>> {
+        val file = scriptInfo.script.getContainingKtFile()
+        val scriptDefinition = KotlinScriptDefinitionProvider.getInstance(file.project).findScriptDefinition(file)
+        val scriptParams = scriptDefinition.getScriptParameters(this)
+        return scriptDefinition.getScriptParametersToPassToSuperclass(this).map { name ->
+            Pair(name,
+                 scriptParams.find { it.name == name }?.type
+                    ?: throw RuntimeException("Unknown script parameter '$name' is specified as a script base class parameter"))
+        }
+    }
 }

compiler/frontend/src/org/jetbrains/kotlin/script/KotlinScriptDefinition.kt

@@ -34,7 +34,8 @@ import kotlin.reflect.KClass
 
 interface KotlinScriptDefinition {
     fun getScriptParameters(scriptDescriptor: ScriptDescriptor): List<ScriptParameter>
-    fun getScriptSuperclasses(scriptDescriptor: ScriptDescriptor): List<KotlinType> = emptyList()
+    fun getScriptSupertypes(scriptDescriptor: ScriptDescriptor): List<KotlinType> = emptyList()
+    fun getScriptParametersToPassToSuperclass(scriptDescriptor: ScriptDescriptor): List<Name> = emptyList()
     fun isScript(file: PsiFile): Boolean
     fun getScriptName(script: KtScript): Name
 }

compiler/tests-common/org/jetbrains/kotlin/scripts/TestScriptDefinitions.kt

@@ -38,7 +38,14 @@ class ReflectedParamClassTestScriptDefinition(extension: String, val name: Strin
             listOf(makeReflectedClassScriptParameter(scriptDescriptor, Name.identifier(name), parameter))
 }
 
-class ReflectedSuperclassTestScriptDefinition(extension: String, parameters: List<ScriptParameter>, val superclass: KClass<out Any>) : SimpleParamsTestScriptDefinition(extension, parameters) {
-    override fun getScriptSuperclasses(scriptDescriptor: ScriptDescriptor): List<KotlinType> =
+open class ReflectedSuperclassTestScriptDefinition(extension: String, parameters: List<ScriptParameter>, val superclass: KClass<out Any>) : SimpleParamsTestScriptDefinition(extension, parameters) {
+    override fun getScriptSupertypes(scriptDescriptor: ScriptDescriptor): List<KotlinType> =
             listOf(getKotlinType(scriptDescriptor, superclass))
 }
+
+class ReflectedSuperclassWithParamsTestScriptDefinition(extension: String, parameters: List<ScriptParameter>, superclass: KClass<out Any>, val superclassParameters: List<ScriptParameter>)
+    : ReflectedSuperclassTestScriptDefinition(extension, parameters, superclass)
+{
+    override fun getScriptParametersToPassToSuperclass(scriptDescriptor: ScriptDescriptor): List<Name> =
+            superclassParameters.map { it.name }
+}

compiler/tests/org/jetbrains/kotlin/scripts/ScriptTest.kt

@@ -80,6 +80,13 @@ class ScriptTest {
         aClass!!.getConstructor(Integer.TYPE).newInstance(4)
     }
 
+    @Test
+    fun testScriptWithBaseClassWithParam() {
+        val aClass = compileScript("fib_dsl.kts", ReflectedSuperclassWithParamsTestScriptDefinition(".kts", numIntParam() + numIntParam("passthrough"), TestDSLClassWithParam::class, numIntParam("passthrough")), runIsolated = false)
+        Assert.assertNotNull(aClass)
+        aClass!!.getConstructor(Integer.TYPE, Integer.TYPE).newInstance(4, 1)
+    }
+
     @Test
     fun testScriptWithInterface() {
         val aClass = compileScript("fib_dsl.kts", ReflectedSuperclassTestScriptDefinition(".kts", numIntParam(), TestDSLInterface::class), runIsolated = false)
@@ -139,8 +146,8 @@ class ScriptTest {
         cp.forEach { addJvmClasspathRoot(it) }
     }
 
-    private fun numIntParam(): List<ScriptParameter> {
-        return listOf(ScriptParameter(Name.identifier("num"), DefaultBuiltIns.Instance.intType))
+    private fun numIntParam(name: String = "num"): List<ScriptParameter> {
+        return listOf(ScriptParameter(Name.identifier(name), DefaultBuiltIns.Instance.intType))
     }
 }
 

compiler/tests/org/jetbrains/kotlin/scripts/TestScriptDsl.kt

@@ -26,3 +26,6 @@ interface TestDSLInterface
 
 fun TestDSLInterface.fibCombine(f: (Int) -> Int, n: Int) = if (n < 2) 1 else f(n - 1) + f(n - 2)
 
+open class TestDSLClassWithParam(val offset: Int)
+
+fun TestDSLClassWithParam.fibCombine(f: (Int) -> Int, n: Int) = if (n < 2) offset else f(n - 1) + f(n - 2)

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

@@ -16,7 +16,12 @@
 
 package org.jetbrains.kotlin.descriptors;
 
+import kotlin.Pair;
 import org.jetbrains.annotations.NotNull;
+import org.jetbrains.kotlin.name.Name;
+import org.jetbrains.kotlin.types.KotlinType;
+
+import java.util.List;
 
 public interface ScriptDescriptor extends ClassDescriptor {
     int getPriority();
@@ -24,4 +29,6 @@ public interface ScriptDescriptor extends ClassDescriptor {
     @NotNull
     @Override
     ConstructorDescriptor getUnsubstitutedPrimaryConstructor();
+
+    List<Pair<Name, KotlinType>> getScriptParametersToPassToSuperclass();
 }

core/descriptors/src/org/jetbrains/kotlin/resolve/DescriptorUtils.kt

@@ -25,6 +25,7 @@ import org.jetbrains.kotlin.incremental.components.LookupLocation
 import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.name.FqNameUnsafe
+import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.resolve.DescriptorUtils
 import org.jetbrains.kotlin.resolve.constants.EnumValue
 import org.jetbrains.kotlin.types.KotlinType
@@ -246,3 +247,16 @@ fun CallableDescriptor.overriddenTreeUniqueAsSequence(useOriginal: Boolean): Seq
 
     return doBuildOverriddenTreeAsSequence()
 }
+
+fun ClassDescriptor.getConstructorByParams(params: List<KotlinType>): ConstructorDescriptor? =
+    getConstructors().firstOrNull {
+        when {
+            it.valueParameters.isEmpty() && params.isEmpty() -> true
+            it.valueParameters.size != params.size -> false
+            else -> it.valueParameters.zip(params) { ctorP, p -> ctorP.type == p }.all { it }
+        }
+    }
+
+// TODO: inline and remove as soon as all usage sies are converted to kotlin
+fun getConstructorByParamsMap(classDescriptor: ClassDescriptor, params: List<Pair<Name, KotlinType>>): ConstructorDescriptor? =
+        classDescriptor.getConstructorByParams(params.map { it.second })