Support reflection on lambdas and function expressions

Write a special annotation containing the bytes for the Callable protobuf
message and deserialize it at runtime properly

 #KT-9005 Fixed

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

@@ -33,14 +33,22 @@ import org.jetbrains.kotlin.descriptors.*;
 import org.jetbrains.kotlin.descriptors.impl.SimpleFunctionDescriptorImpl;
 import org.jetbrains.kotlin.incremental.components.NoLookupLocation;
 import org.jetbrains.kotlin.load.java.JvmAbi;
+import org.jetbrains.kotlin.load.java.JvmAnnotationNames;
 import org.jetbrains.kotlin.load.kotlin.PackageClassUtils;
 import org.jetbrains.kotlin.psi.JetElement;
 import org.jetbrains.kotlin.resolve.BindingContext;
 import org.jetbrains.kotlin.resolve.DescriptorUtils;
 import org.jetbrains.kotlin.resolve.scopes.JetScope;
+import org.jetbrains.kotlin.serialization.DescriptorSerializer;
+import org.jetbrains.kotlin.serialization.ProtoBuf;
+import org.jetbrains.kotlin.serialization.SerializationUtil;
+import org.jetbrains.kotlin.serialization.StringTable;
+import org.jetbrains.kotlin.serialization.deserialization.NameResolver;
+import org.jetbrains.kotlin.serialization.jvm.BitEncoding;
 import org.jetbrains.kotlin.types.JetType;
 import org.jetbrains.kotlin.types.expressions.OperatorConventions;
 import org.jetbrains.kotlin.utils.UtilsPackage;
+import org.jetbrains.org.objectweb.asm.AnnotationVisitor;
 import org.jetbrains.org.objectweb.asm.MethodVisitor;
 import org.jetbrains.org.objectweb.asm.Type;
 import org.jetbrains.org.objectweb.asm.commons.InstructionAdapter;
@@ -219,6 +227,23 @@ public class ClosureCodegen extends MemberCodegen<JetElement> {
     @Override
     protected void generateKotlinAnnotation() {
         writeKotlinSyntheticClassAnnotation(v, syntheticClassKind);
+
+        DescriptorSerializer serializer =
+                DescriptorSerializer.createTopLevel(new JvmSerializerExtension(v.getSerializationBindings(), typeMapper));
+
+        ProtoBuf.Callable callableProto = serializer.callableProto(funDescriptor).build();
+
+        StringTable strings = serializer.getStringTable();
+        NameResolver nameResolver = new NameResolver(strings.serializeSimpleNames(), strings.serializeQualifiedNames());
+
+        AnnotationVisitor av = v.getVisitor().visitAnnotation(asmDescByFqNameWithoutInnerClasses(JvmAnnotationNames.KOTLIN_CALLABLE), true);
+        av.visit(JvmAnnotationNames.ABI_VERSION_FIELD_NAME, JvmAbi.VERSION);
+        AnnotationVisitor array = av.visitArray(JvmAnnotationNames.DATA_FIELD_NAME);
+        for (String string : BitEncoding.encodeBytes(SerializationUtil.serializeCallableData(nameResolver, callableProto))) {
+            array.visit(null, string);
+        }
+        array.visitEnd();
+        av.visitEnd();
     }
 
     @Override

compiler/serialization/src/org/jetbrains/kotlin/serialization/SerializationUtil.java

@@ -54,6 +54,19 @@ public class SerializationUtil {
         }
     }
 
+    @NotNull
+    public static byte[] serializeCallableData(@NotNull NameResolver nameResolver, @NotNull ProtoBuf.Callable callableProto) {
+        try {
+            ByteArrayOutputStream result = new ByteArrayOutputStream();
+            serializeNameResolver(result, nameResolver);
+            callableProto.writeTo(result);
+            return result.toByteArray();
+        }
+        catch (IOException e) {
+            throw UtilsPackage.rethrow(e);
+        }
+    }
+
     private static void serializeNameResolver(@NotNull OutputStream out, @NotNull NameResolver nameResolver) {
         serializeStringTable(out, nameResolver.getStringTable(), nameResolver.getQualifiedNameTable());
     }

compiler/testData/codegen/boxWithStdlib/reflection/lambdaClasses/parameterNamesAndNullability.kt

@@ -0,0 +1,50 @@
+import kotlin.reflect.*
+import kotlin.reflect.jvm.*
+import kotlin.test.assertEquals
+import kotlin.test.assertNull
+
+fun lambda() {
+    val f = { x: Int, y: String? -> }
+
+    val g = f.reflect()!!
+
+    // TODO: maybe change this name
+    assertEquals("<anonymous>", g.name)
+    assertEquals(listOf("x", "y"), g.parameters.map { it.name })
+    assertEquals(listOf(false, true), g.parameters.map { it.type.isMarkedNullable })
+}
+
+fun funExpr() {
+    val f = fun(x: Int, y: String?) {}
+
+    val g = f.reflect()!!
+
+    // TODO: maybe change this name
+    assertEquals("<no name provided>", g.name)
+    assertEquals(listOf("x", "y"), g.parameters.map { it.name })
+    assertEquals(listOf(false, true), g.parameters.map { it.type.isMarkedNullable })
+}
+
+fun extensionFunExpr() {
+    val f = fun String.(): String = this
+
+    val g = f.reflect()!!
+
+    assertEquals(KParameter.Kind.EXTENSION_RECEIVER, g.parameters.single().kind)
+    assertEquals(null, g.parameters.single().name)
+}
+
+fun customFunction() {
+    val f = object : Function<String> {}
+    assertNull(f.reflect())
+}
+
+fun box(): String {
+    lambda()
+    funExpr()
+    extensionFunExpr()
+
+    customFunction()
+
+    return "OK"
+}

compiler/tests/org/jetbrains/kotlin/codegen/generated/BlackBoxWithStdlibCodegenTestGenerated.java

@@ -3421,6 +3421,21 @@ public class BlackBoxWithStdlibCodegenTestGenerated extends AbstractBlackBoxCode
             }
         }
 
+        @TestMetadata("compiler/testData/codegen/boxWithStdlib/reflection/lambdaClasses")
+        @TestDataPath("$PROJECT_ROOT")
+        @RunWith(JUnit3RunnerWithInners.class)
+        public static class LambdaClasses extends AbstractBlackBoxCodegenTest {
+            public void testAllFilesPresentInLambdaClasses() throws Exception {
+                JetTestUtils.assertAllTestsPresentByMetadata(this.getClass(), new File("compiler/testData/codegen/boxWithStdlib/reflection/lambdaClasses"), Pattern.compile("^(.+)\\.kt$"), true);
+            }
+
+            @TestMetadata("parameterNamesAndNullability.kt")
+            public void testParameterNamesAndNullability() throws Exception {
+                String fileName = JetTestUtils.navigationMetadata("compiler/testData/codegen/boxWithStdlib/reflection/lambdaClasses/parameterNamesAndNullability.kt");
+                doTestWithStdlib(fileName);
+            }
+        }
+
         @TestMetadata("compiler/testData/codegen/boxWithStdlib/reflection/mapping")
         @TestDataPath("$PROJECT_ROOT")
         @RunWith(JUnit3RunnerWithInners.class)

core/descriptor.loader.java/src/org/jetbrains/kotlin/load/java/JvmAnnotationNames.java

@@ -30,6 +30,7 @@ import java.util.Set;
 public final class JvmAnnotationNames {
     public static final FqName KOTLIN_CLASS = KotlinClass.CLASS_NAME.getFqNameForClassNameWithoutDollars();
     public static final FqName KOTLIN_PACKAGE = new FqName("kotlin.jvm.internal.KotlinPackage");
+    public static final FqName KOTLIN_CALLABLE = new FqName("kotlin.jvm.internal.KotlinCallable");
 
     public static final FqName KOTLIN_SIGNATURE = new FqName("kotlin.jvm.KotlinSignature");
     public static final FqName OLD_KOTLIN_SIGNATURE = new FqName("jet.runtime.typeinfo.KotlinSignature");

core/descriptors.runtime/src/org/jetbrains/kotlin/load/kotlin/reflect/RuntimeModuleData.kt

@@ -34,10 +34,14 @@ import org.jetbrains.kotlin.load.kotlin.JavaClassDataFinder
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.platform.JavaToKotlinClassMap
 import org.jetbrains.kotlin.resolve.jvm.JavaDescriptorResolver
+import org.jetbrains.kotlin.serialization.deserialization.DeserializationComponents
 import org.jetbrains.kotlin.serialization.deserialization.LocalClassResolver
 import org.jetbrains.kotlin.storage.LockBasedStorageManager
 
-public class RuntimeModuleData private constructor(public val module: ModuleDescriptor, public val localClassResolver: LocalClassResolver) {
+public class RuntimeModuleData private constructor(public val deserialization: DeserializationComponents) {
+    public val module: ModuleDescriptor get() = deserialization.moduleDescriptor
+    public val localClassResolver: LocalClassResolver get() = deserialization.localClassResolver
+
     companion object {
         public fun create(classLoader: ClassLoader): RuntimeModuleData {
             val storageManager = LockBasedStorageManager()
@@ -63,7 +67,7 @@ public class RuntimeModuleData private constructor(public val module: ModuleDesc
             module.setDependencies(module, KotlinBuiltIns.getInstance().getBuiltInsModule())
             module.initialize(javaDescriptorResolver.packageFragmentProvider)
 
-            return RuntimeModuleData(module, deserializationComponentsForJava.components.localClassResolver)
+            return RuntimeModuleData(deserializationComponentsForJava.components)
         }
     }
 }

core/deserialization/src/org/jetbrains/kotlin/serialization/deserialization/MemberDeserializer.kt

@@ -125,7 +125,7 @@ public class MemberDeserializer(private val c: DeserializationContext) {
         if (Flags.HAS_CONSTANT.get(flags)) {
             property.setCompileTimeInitializer(
                     c.storageManager.createNullableLazyValue {
-                        val container = c.containingDeclaration.asProtoContainer()
+                        val container = c.containingDeclaration.asProtoContainer()!!
                         c.components.annotationAndConstantLoader.loadPropertyConstant(container, proto, c.nameResolver, property.getReturnType())
                     }
             )
@@ -178,9 +178,11 @@ public class MemberDeserializer(private val c: DeserializationContext) {
             return Annotations.EMPTY
         }
         return DeserializedAnnotationsWithPossibleTargets(c.storageManager) {
-            c.components.annotationAndConstantLoader.loadCallableAnnotations(
+            c.containingDeclaration.asProtoContainer()?.let {
-                    c.containingDeclaration.asProtoContainer(), proto, c.nameResolver, kind
+                c.components.annotationAndConstantLoader.loadCallableAnnotations(
-            )
+                        it, proto, c.nameResolver, kind
+                )
+            }.orEmpty()
         }
     }
 
@@ -191,11 +193,13 @@ public class MemberDeserializer(private val c: DeserializationContext) {
     ): Annotations {
         return DeserializedAnnotationsWithPossibleTargets(c.storageManager) {
             if (proto.hasReceiverType()) {
-                val container = c.containingDeclaration.asProtoContainer()
+                c.containingDeclaration.asProtoContainer()?.let {
-                c.components.annotationAndConstantLoader
+                    c.components.annotationAndConstantLoader
-                        .loadExtensionReceiverParameterAnnotations(container, proto, c.nameResolver, receiverTargetedKind)
+                            .loadExtensionReceiverParameterAnnotations(it, proto, c.nameResolver, receiverTargetedKind)
-                        .map { AnnotationWithTarget(it, AnnotationUseSiteTarget.RECEIVER) }
+                            .map { AnnotationWithTarget(it, AnnotationUseSiteTarget.RECEIVER) }
-            } else emptyList()
+                }.orEmpty()
+            }
+            else emptyList()
         }
     }
 
@@ -206,7 +210,7 @@ public class MemberDeserializer(private val c: DeserializationContext) {
         return callable.getValueParameterList().mapIndexed { i, proto ->
             ValueParameterDescriptorImpl(
                     callableDescriptor, null, i,
-                    getParameterAnnotations(containerOfCallable, callable, kind, proto),
+                    containerOfCallable?.let { getParameterAnnotations(it, callable, kind, proto) } ?: Annotations.EMPTY,
                     c.nameResolver.getName(proto.getName()),
                     c.typeDeserializer.type(proto.getType()),
                     Flags.DECLARES_DEFAULT_VALUE.get(proto.getFlags()),
@@ -227,9 +231,9 @@ public class MemberDeserializer(private val c: DeserializationContext) {
         }
     }
 
-    private fun DeclarationDescriptor.asProtoContainer(): ProtoContainer = when(this) {
+    private fun DeclarationDescriptor.asProtoContainer(): ProtoContainer? = when(this) {
         is PackageFragmentDescriptor -> ProtoContainer(null, fqName)
         is DeserializedClassDescriptor -> ProtoContainer(classProto, null)
-        else -> error("Only members in classes or package fragments should be serialized: $this")
+        else -> null // TODO: support annotations on lambdas and their parameters
     }
 }

core/reflection.jvm/src/kotlin/reflect/jvm/internal/KFunctionFromReferenceImpl.kt

@@ -76,5 +76,6 @@ object EmptyContainerForLocal : KDeclarationContainerImpl() {
 
     override fun getFunctions(name: Name): Collection<FunctionDescriptor> = fail()
 
-    private fun fail() = throw KotlinReflectionInternalError("Introspecting local functions is not yet supported in Kotlin reflection")
+    private fun fail() = throw KotlinReflectionInternalError("Introspecting local functions, lambdas and function expressions " +
+                                                             "is not yet fully supported in Kotlin reflection")
 }

core/reflection.jvm/src/kotlin/reflect/jvm/internal/moduleByClassLoader.kt

@@ -44,7 +44,7 @@ private class WeakClassLoaderBox(classLoader: ClassLoader) {
             ref.get()?.let { it.toString() } ?: "<null>"
 }
 
-private fun Class<*>.getOrCreateModule(): RuntimeModuleData {
+internal fun Class<*>.getOrCreateModule(): RuntimeModuleData {
     val classLoader = this.safeClassLoader
     val key = WeakClassLoaderBox(classLoader)
 

core/reflection.jvm/src/kotlin/reflect/jvm/reflectLambda.kt

@@ -0,0 +1,48 @@
+/*
+ * Copyright 2010-2015 JetBrains s.r.o.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package kotlin.reflect.jvm
+
+import org.jetbrains.kotlin.descriptors.FunctionDescriptor
+import org.jetbrains.kotlin.serialization.ProtoBuf
+import org.jetbrains.kotlin.serialization.deserialization.DeserializationContext
+import org.jetbrains.kotlin.serialization.deserialization.MemberDeserializer
+import org.jetbrains.kotlin.serialization.deserialization.NameResolver
+import org.jetbrains.kotlin.serialization.jvm.BitEncoding
+import org.jetbrains.kotlin.serialization.jvm.JvmProtoBufUtil
+import kotlin.jvm.internal.KotlinCallable
+import kotlin.reflect.KFunction
+import kotlin.reflect.jvm.internal.EmptyContainerForLocal
+import kotlin.reflect.jvm.internal.KFunctionImpl
+import kotlin.reflect.jvm.internal.getOrCreateModule
+
+/**
+ * This is an experimental API. Given a class for a compiled Kotlin lambda or a function expression,
+ * returns a [KFunction] instance providing introspection capabilities for that lambda or function expression and its parameters.
+ * Not all features are currently supported, in particular [KCallable.call] and [KCallable.callBy] will fail at the moment.
+ */
+public fun <R> Function<R>.reflect(): KFunction<R>? {
+    val callable = javaClass.getAnnotation(KotlinCallable::class.java) ?: return null
+    val input = BitEncoding.decodeBytes(callable.data).inputStream()
+    val nameResolver = NameResolver.read(input)
+    val proto = ProtoBuf.Callable.parseFrom(input, JvmProtoBufUtil.EXTENSION_REGISTRY)
+    val moduleData = javaClass.getOrCreateModule()
+    val context = DeserializationContext(moduleData.deserialization, nameResolver, moduleData.module,
+                                         parentTypeDeserializer = null, typeParameters = listOf())
+    val descriptor = MemberDeserializer(context).loadCallable(proto) as FunctionDescriptor
+    @suppress("UNCHECKED_CAST")
+    return KFunctionImpl(EmptyContainerForLocal, descriptor) as KFunction<R>
+}

core/runtime.jvm/src/kotlin/jvm/internal/KotlinCallable.java

@@ -0,0 +1,27 @@
+/*
+ * Copyright 2010-2015 JetBrains s.r.o.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package kotlin.jvm.internal;
+
+import java.lang.annotation.Retention;
+import java.lang.annotation.RetentionPolicy;
+
+@Retention(RetentionPolicy.RUNTIME)
+public @interface KotlinCallable {
+    int abiVersion();
+
+    String[] data();
+}