Codegen for callable reference expressions

#KT-1183 In Progress
2013-04-15 19:47:10 +04:00
parent 1eeaaad05d
commit 054e5fb5e7
46 changed files with 877 additions and 7 deletions
@@ -36,7 +36,7 @@ import org.jetbrains.jet.codegen.binding.CodegenBinding;
 import org.jetbrains.jet.codegen.binding.MutableClosure;
 import org.jetbrains.jet.codegen.context.*;
 import org.jetbrains.jet.codegen.intrinsics.IntrinsicMethod;
-import org.jetbrains.jet.codegen.signature.JvmPropertyAccessorSignature;
+import org.jetbrains.jet.codegen.signature.JvmMethodSignature;
 import org.jetbrains.jet.codegen.state.GenerationState;
 import org.jetbrains.jet.codegen.state.JetTypeMapper;
 import org.jetbrains.jet.codegen.state.JetTypeMapperMode;
@@ -67,12 +67,16 @@ import java.util.*;
 import static org.jetbrains.asm4.Opcodes.*;
 import static org.jetbrains.jet.codegen.AsmUtil.*;
 import static org.jetbrains.jet.codegen.CodegenUtil.*;
+import static org.jetbrains.jet.codegen.FunctionTypesUtil.functionTypeToImpl;
 import static org.jetbrains.jet.codegen.FunctionTypesUtil.getFunctionImplClassName;
 import static org.jetbrains.jet.codegen.binding.CodegenBinding.*;
 import static org.jetbrains.jet.lang.resolve.BindingContext.*;
 import static org.jetbrains.jet.lang.resolve.BindingContextUtils.descriptorToDeclaration;
 import static org.jetbrains.jet.lang.resolve.BindingContextUtils.getNotNull;
+import static org.jetbrains.jet.lang.resolve.calls.tasks.ExplicitReceiverKind.RECEIVER_ARGUMENT;
+import static org.jetbrains.jet.lang.resolve.calls.tasks.ExplicitReceiverKind.THIS_OBJECT;
 import static org.jetbrains.jet.lang.resolve.java.AsmTypeConstants.*;
+import static org.jetbrains.jet.lang.resolve.scopes.receivers.ReceiverValue.NO_RECEIVER;

 public class ExpressionCodegen extends JetVisitor<StackValue, StackValue> implements LocalLookup {

@@ -2320,6 +2324,164 @@ public class ExpressionCodegen extends JetVisitor<StackValue, StackValue> implem
        return lookupLocalIndex(declarationDescriptor);
    }

+    @Override
+    public StackValue visitCallableReferenceExpression(JetCallableReferenceExpression expression, StackValue data) {
+        // TODO: properties
+        final FunctionDescriptor functionDescriptor = bindingContext.get(CALLABLE_REFERENCE, expression);
+        assert functionDescriptor != null : "Callable reference is not resolved to descriptor: " + expression.getText();
+
+        final ResolvedCall<? extends CallableDescriptor> resolvedCall = bindingContext.get(RESOLVED_CALL, expression.getCallableReference());
+        assert resolvedCall != null : "Callable reference is not resolved: " + functionDescriptor + " " + expression.getText();
+
+        JetType kFunctionType = bindingContext.get(EXPRESSION_TYPE, expression);
+        assert kFunctionType != null : "Callable reference is not type checked: " + expression.getText();
+        ClassDescriptor kFunctionImpl = functionTypeToImpl(kFunctionType);
+        assert kFunctionImpl != null : "Impl type is not found for the function type: " + kFunctionType;
+
+        JvmClassName closureSuperClass = JvmClassName.byType(typeMapper.mapType(kFunctionImpl));
+
+        ClosureCodegen closureCodegen = new ClosureCodegen(state, expression, functionDescriptor, null, closureSuperClass, context, this,
+                new FunctionGenerationStrategy() {
+                    @Override
+                    public void generateBody(
+                           @NotNull MethodVisitor mv,
+                           @NotNull JvmMethodSignature signature,
+                           @NotNull MethodContext context,
+                           @NotNull FrameMap frameMap
+                    ) {
+                        /*
+                         Here we need to put the arguments from our locals to the stack and invoke the referenced method. Since invokation
+                         of methods is highly dependent on expressions, we create a fake call expression. Then we create a new instance of
+                         ExpressionCodegen and, in order for it to generate code correctly, we save to its 'tempVariables' field every
+                         argument of our fake expression, pointing it to the corresponding index in our locals. This way generation of
+                         every argument boils down to calling LOAD with the corresponding index
+                         */
+
+                        FunctionDescriptor referencedFunction = (FunctionDescriptor) resolvedCall.getResultingDescriptor();
+                        JetType returnJetType = referencedFunction.getReturnType();
+                        assert returnJetType != null : "Return type can't be null: " + referencedFunction;
+                        Type returnType = typeMapper.mapReturnType(returnJetType);
+
+                        JetCallExpression fakeExpression = constructFakeFunctionCall(referencedFunction);
+                        final List<? extends ValueArgument> fakeArguments = fakeExpression.getValueArguments();
+
+                        ExpressionCodegen codegen = new ExpressionCodegen(mv, frameMap, returnType, context, state);
+
+                        final ReceiverValue receiverValue = computeAndSaveReceiver(signature, codegen);
+                        computeAndSaveArguments(frameMap, fakeArguments, codegen);
+
+                        ResolvedCall<CallableDescriptor> fakeResolvedCall = new DelegatingResolvedCall<CallableDescriptor>(resolvedCall) {
+                            @NotNull
+                            @Override
+                            public ReceiverValue getReceiverArgument() {
+                                return resolvedCall.getExplicitReceiverKind() == RECEIVER_ARGUMENT ? receiverValue : NO_RECEIVER;
+                            }
+
+                            @NotNull
+                            @Override
+                            public ReceiverValue getThisObject() {
+                                return resolvedCall.getExplicitReceiverKind() == THIS_OBJECT ? receiverValue : NO_RECEIVER;
+                            }
+
+                            @NotNull
+                            @Override
+                            public List<ResolvedValueArgument> getValueArgumentsByIndex() {
+                                List<ResolvedValueArgument> result = new ArrayList<ResolvedValueArgument>(fakeArguments.size());
+                                for (ValueArgument argument : fakeArguments) {
+                                    result.add(new ExpressionValueArgument(argument));
+                                }
+                                return result;
+                            }
+                        };
+
+                        StackValue result;
+                        if (referencedFunction instanceof ConstructorDescriptor) {
+                            if (returnType.getSort() == Type.ARRAY) {
+                                codegen.generateNewArray(fakeExpression, returnJetType);
+                                result = StackValue.onStack(returnType);
+                            }
+                            else {
+                                result = codegen.generateConstructorCall(fakeResolvedCall, StackValue.none(), returnType);
+                            }
+                        }
+                        else {
+                            Call call = CallMaker.makeCall(fakeExpression, NO_RECEIVER, null, fakeExpression, fakeArguments);
+                            result = codegen.invokeFunction(call, StackValue.none(), fakeResolvedCall);
+                        }
+
+                        InstructionAdapter v = new InstructionAdapter(mv);
+                        result.put(returnType, v);
+                        v.areturn(returnType);
+                    }
+
+                    @NotNull
+                    private JetCallExpression constructFakeFunctionCall(@NotNull CallableDescriptor referencedFunction) {
+                        StringBuilder fakeFunctionCall = new StringBuilder("callableReferenceFakeCall(");
+                        for (Iterator<ValueParameterDescriptor> iterator = referencedFunction.getValueParameters().iterator();
+                             iterator.hasNext(); ) {
+                            ValueParameterDescriptor descriptor = iterator.next();
+                            fakeFunctionCall.append("p").append(descriptor.getIndex());
+                            if (iterator.hasNext()) {
+                                fakeFunctionCall.append(", ");
+                            }
+                        }
+                        fakeFunctionCall.append(")");
+                        return (JetCallExpression) JetPsiFactory.createExpression(state.getProject(), fakeFunctionCall.toString());
+                    }
+
+                    private void computeAndSaveArguments(
+                            @NotNull FrameMap frameMap,
+                            @NotNull List<? extends ValueArgument> fakeArguments,
+                            @NotNull ExpressionCodegen codegen
+                    ) {
+                        for (ValueParameterDescriptor parameter : functionDescriptor.getValueParameters()) {
+                            ValueArgument fakeArgument = fakeArguments.get(parameter.getIndex());
+                            Type type = typeMapper.mapType(parameter);
+                            int localIndex = frameMap.getIndex(parameter);
+                            codegen.tempVariables.put(fakeArgument.getArgumentExpression(), StackValue.local(localIndex, type));
+                        }
+                    }
+
+                    @NotNull
+                    private ReceiverValue computeAndSaveReceiver(
+                            @NotNull JvmMethodSignature signature,
+                            @NotNull ExpressionCodegen codegen
+                    ) {
+                        CallableDescriptor referencedFunction = resolvedCall.getCandidateDescriptor();
+
+                        ReceiverParameterDescriptor receiverParameter = referencedFunction.getReceiverParameter();
+                        ReceiverParameterDescriptor expectedThisObject = referencedFunction.getExpectedThisObject();
+                        assert receiverParameter == null || expectedThisObject == null :
+                                "Extensions in classes can't be referenced via callable reference expressions: " + referencedFunction;
+
+                        ReceiverParameterDescriptor receiver = receiverParameter != null ? receiverParameter : expectedThisObject;
+
+                        if (receiver == null) {
+                            return NO_RECEIVER;
+                        }
+
+                        JetExpression receiverExpression = JetPsiFactory.createExpression(state.getProject(),
+                                                                                          "callableReferenceFakeReceiver");
+
+                        Type firstParameterType = signature.getAsmMethod().getArgumentTypes()[0];
+                        // 0 is this (the closure class), 1 is the method's first parameter
+                        codegen.tempVariables.put(receiverExpression, StackValue.local(1, firstParameterType));
+
+                        return new ExpressionReceiver(receiverExpression, receiver.getType());
+                    }
+
+                    @Override
+                    public boolean needsLocalVariableTable() {
+                        return false;
+                    }
+                }
+        );
+
+        closureCodegen.gen();
+
+        return closureCodegen.putInstanceOnStack(v, this);
+    }
+
    @Override
    public StackValue visitDotQualifiedExpression(JetDotQualifiedExpression expression, StackValue receiver) {
        StackValue receiverValue = StackValue.none();
@@ -182,8 +182,10 @@ public class FunctionCodegen extends GenerationStateAware {
        Label methodEnd = new Label();
        mv.visitLabel(methodEnd);

-        Type thisType = getThisTypeForFunction(functionDescriptor, context);
-        generateLocalVariableTable(mv, functionDescriptor, thisType, methodBegin, methodEnd, localVariableNames, labelsForSharedVars);
+        if (strategy.needsLocalVariableTable()) {
+            Type thisType = getThisTypeForFunction(functionDescriptor, context);
+            generateLocalVariableTable(mv, functionDescriptor, thisType, methodBegin, methodEnd, localVariableNames, labelsForSharedVars);
+        }
    }

    @NotNull
@@ -45,6 +45,10 @@ public abstract class FunctionGenerationStrategy {
        return localVariableNames;
    }

+    public boolean needsLocalVariableTable() {
+        return true;
+    }
+

    public static class Default extends FunctionGenerationStrategy {
        private final GenerationState state;
@@ -16,7 +16,9 @@

 package org.jetbrains.jet.codegen;

+import com.google.common.collect.ImmutableMap;
 import org.jetbrains.annotations.NotNull;
+import org.jetbrains.annotations.Nullable;
 import org.jetbrains.jet.lang.descriptors.*;
 import org.jetbrains.jet.lang.descriptors.impl.MutableClassDescriptor;
 import org.jetbrains.jet.lang.resolve.java.JvmClassName;
@@ -31,11 +33,19 @@ import java.util.List;
 public class FunctionTypesUtil {
    private static final List<ClassDescriptor> FUNCTIONS;
    private static final List<ClassDescriptor> EXTENSION_FUNCTIONS;
+    private static final List<ClassDescriptor> K_FUNCTIONS;
+    private static final List<ClassDescriptor> K_MEMBER_FUNCTIONS;
+    private static final List<ClassDescriptor> K_EXTENSION_FUNCTIONS;
+
+    private static final ImmutableMap<ClassDescriptor, ClassDescriptor> FUNCTION_TO_IMPL;

    static {
        int n = KotlinBuiltIns.FUNCTION_TRAIT_COUNT;
        FUNCTIONS = new ArrayList<ClassDescriptor>(n);
        EXTENSION_FUNCTIONS = new ArrayList<ClassDescriptor>(n);
+        K_FUNCTIONS = new ArrayList<ClassDescriptor>(n);
+        K_MEMBER_FUNCTIONS = new ArrayList<ClassDescriptor>(n);
+        K_EXTENSION_FUNCTIONS = new ArrayList<ClassDescriptor>(n);

        KotlinBuiltIns builtIns = KotlinBuiltIns.getInstance();
        for (int i = 0; i < n; i++) {
@@ -44,13 +54,36 @@ public class FunctionTypesUtil {

            Name extensionFunctionImpl = Name.identifier("ExtensionFunctionImpl" + i);
            EXTENSION_FUNCTIONS.add(createFunctionImplDescriptor(extensionFunctionImpl, builtIns.getExtensionFunction(i)));
+
+            Name kFunctionImpl = Name.identifier("KFunctionImpl" + i);
+            K_FUNCTIONS.add(createFunctionImplDescriptor(kFunctionImpl, builtIns.getKFunction(i)));
+
+            Name kMemberFunctionImpl = Name.identifier("KMemberFunctionImpl" + i);
+            K_MEMBER_FUNCTIONS.add(createFunctionImplDescriptor(kMemberFunctionImpl, builtIns.getKMemberFunction(i)));
+
+            Name kExtensionFunctionImpl = Name.identifier("KExtensionFunctionImpl" + i);
+            K_EXTENSION_FUNCTIONS.add(createFunctionImplDescriptor(kExtensionFunctionImpl, builtIns.getKExtensionFunction(i)));
        }
+
+        ImmutableMap.Builder<ClassDescriptor, ClassDescriptor> builder = ImmutableMap.builder();
+        for (int i = 0; i < n; i++) {
+            builder.put(builtIns.getKFunction(i), K_FUNCTIONS.get(i));
+            builder.put(builtIns.getKMemberFunction(i), K_MEMBER_FUNCTIONS.get(i));
+            builder.put(builtIns.getKExtensionFunction(i), K_EXTENSION_FUNCTIONS.get(i));
+        }
+        FUNCTION_TO_IMPL = builder.build();
    }

    private FunctionTypesUtil() {
    }


+    @Nullable
+    public static ClassDescriptor functionTypeToImpl(@NotNull JetType functionType) {
+        //noinspection SuspiciousMethodCalls
+        return FUNCTION_TO_IMPL.get(functionType.getConstructor().getDeclarationDescriptor());
+    }
+
    @NotNull
    public static JetType getSuperTypeForClosure(@NotNull FunctionDescriptor funDescriptor, int arity) {
        if (funDescriptor.getReceiverParameter() != null) {
@@ -115,11 +115,12 @@ class CodegenAnnotatingVisitor extends JetVisitorVoid {
            if (declaration instanceof JetFunctionLiteralExpression ||
                declaration instanceof JetNamedFunction ||
                declaration instanceof JetObjectLiteralExpression ||
-                declaration instanceof JetCallExpression) {
+                declaration instanceof JetCallExpression ||
+                declaration instanceof JetCallableReferenceExpression) {
            }
            else {
                throw new IllegalStateException(
-                        "Class-less declaration which is not JetFunctionLiteralExpression|JetNamedFunction|JetObjectLiteralExpression|JetCallExpression : " +
+                        "Class-less declaration which is not JetFunctionLiteralExpression|JetNamedFunction|JetObjectLiteralExpression|JetCallExpression|JetCallableReferenceExpression : " +
                        declaration.getClass().getName());
            }
        }
@@ -267,6 +268,23 @@ class CodegenAnnotatingVisitor extends JetVisitorVoid {
        classStack.pop();
    }

+    @Override
+    public void visitCallableReferenceExpression(JetCallableReferenceExpression expression) {
+        FunctionDescriptor functionDescriptor = bindingContext.get(CALLABLE_REFERENCE, expression);
+        // working around a problem with shallow analysis
+        if (functionDescriptor == null) return;
+
+        String name = inventAnonymousClassName(expression);
+        ClassDescriptor classDescriptor = recordClassForFunction(functionDescriptor);
+        recordClosure(bindingTrace, expression, classDescriptor, peekFromStack(classStack), JvmClassName.byInternalName(name), true);
+
+        classStack.push(classDescriptor);
+        nameStack.push(name);
+        super.visitCallableReferenceExpression(expression);
+        nameStack.pop();
+        classStack.pop();
+    }
+
    @Override
    public void visitProperty(JetProperty property) {
        DeclarationDescriptor propertyDescriptor = bindingContext.get(DECLARATION_TO_DESCRIPTOR, property);