Refactor SAM constructor invocation in JVM codegen

Compute all necessary info in CodegenAnnotatingVisitor and save to trace to be
read by ExpressionCodegen later

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

@@ -33,7 +33,6 @@ import org.jetbrains.annotations.Nullable;
 import org.jetbrains.kotlin.backend.common.CodegenUtil;
 import org.jetbrains.kotlin.builtins.KotlinBuiltIns;
 import org.jetbrains.kotlin.codegen.binding.CalculatedClosure;
-import org.jetbrains.kotlin.codegen.binding.CodegenBinding;
 import org.jetbrains.kotlin.codegen.context.*;
 import org.jetbrains.kotlin.codegen.inline.*;
 import org.jetbrains.kotlin.codegen.intrinsics.IntrinsicMethod;
@@ -264,10 +263,9 @@ public class ExpressionCodegen extends JetVisitor<StackValue, StackValue> implem
         }
         try {
             if (selector instanceof JetExpression) {
-                JetExpression expression = (JetExpression) selector;
-                SamType samType = bindingContext.get(CodegenBinding.SAM_VALUE, expression);
-                if (samType != null) {
-                    return genSamInterfaceValue(expression, samType, visitor);
+                StackValue samValue = genSamInterfaceValue((JetExpression) selector, visitor);
+                if (samValue != null) {
+                    return samValue;
                 }
             }
 
@@ -2108,50 +2106,29 @@ public class ExpressionCodegen extends JetVisitor<StackValue, StackValue> implem
         }
 
         if (funDescriptor.getOriginal() instanceof SamConstructorDescriptor) {
-            //noinspection ConstantConditions
-            SamType samType = SamType.create(funDescriptor.getReturnType());
-            assert samType != null : "SamType is not created for SAM constructor: " + funDescriptor;
-            return invokeSamConstructor(expression, resolvedCall, samType);
+            JetExpression argumentExpression = bindingContext.get(SAM_CONSTRUCTOR_TO_ARGUMENT, expression);
+            assert argumentExpression != null : "Argument expression is not saved for a SAM constructor: " + funDescriptor;
+            return genSamInterfaceValue(argumentExpression, this);
         }
 
         return invokeFunction(resolvedCall, receiver);
     }
 
-    @NotNull
-    private StackValue invokeSamConstructor(
-            @NotNull JetCallExpression expression,
-            @NotNull ResolvedCall<?> resolvedCall,
-            @NotNull SamType samType
-    ) {
-        List<ResolvedValueArgument> arguments = resolvedCall.getValueArgumentsByIndex();
-        if (arguments == null) {
-            throw new IllegalStateException("Failed to arrange value arguments by index: " + resolvedCall.getResultingDescriptor());
-        }
-        ResolvedValueArgument argument = arguments.get(0);
-        if (!(argument instanceof ExpressionValueArgument)) {
-            throw new IllegalStateException(
-                    "argument of SAM constructor is " + argument.getClass().getName() + " " + expression.getText());
-        }
-        ValueArgument valueArgument = ((ExpressionValueArgument) argument).getValueArgument();
-        assert valueArgument != null : "getValueArgument() is null for " + expression.getText();
-        JetExpression argumentExpression = valueArgument.getArgumentExpression();
-        assert argumentExpression != null : "getArgumentExpression() is null for " + expression.getText();
-
-        return genSamInterfaceValue(argumentExpression, samType, this);
-    }
-
-    @NotNull
+    @Nullable
     private StackValue genSamInterfaceValue(
             @NotNull final JetExpression expression,
-            @NotNull final SamType samType,
             @NotNull final JetVisitor<StackValue, StackValue> visitor
     ) {
+        final SamType samType = bindingContext.get(SAM_VALUE, expression);
+        if (samType == null) return null;
+
         if (expression instanceof JetFunctionLiteralExpression) {
             return genClosure(((JetFunctionLiteralExpression) expression).getFunctionLiteral(), samType,
                               KotlinSyntheticClass.Kind.SAM_LAMBDA);
         }
 
-        final Type asmType = state.getSamWrapperClasses().getSamWrapperClass(samType, expression.getContainingJetFile(), getParentCodegen());
+        final Type asmType =
+                state.getSamWrapperClasses().getSamWrapperClass(samType, expression.getContainingJetFile(), getParentCodegen());
 
         return StackValue.operation(asmType, new Function1<InstructionAdapter, Unit>() {
             @Override

compiler/backend/src/org/jetbrains/kotlin/codegen/binding/CodegenAnnotatingVisitor.java

@@ -34,6 +34,7 @@ import org.jetbrains.kotlin.codegen.when.WhenByEnumsMapping;
 import org.jetbrains.kotlin.descriptors.*;
 import org.jetbrains.kotlin.descriptors.impl.ClassDescriptorImpl;
 import org.jetbrains.kotlin.load.java.JvmAbi;
+import org.jetbrains.kotlin.load.java.descriptors.SamConstructorDescriptor;
 import org.jetbrains.kotlin.load.kotlin.PackagePartClassUtils;
 import org.jetbrains.kotlin.name.Name;
 import org.jetbrains.kotlin.psi.*;
@@ -414,6 +415,8 @@ class CodegenAnnotatingVisitor extends JetVisitorVoid {
         CallableDescriptor descriptor = call.getResultingDescriptor();
         if (!(descriptor instanceof FunctionDescriptor)) return;
 
+        recordSamConstructorIfNeeded(expression, call);
+
         FunctionDescriptor original = SamCodegenUtil.getOriginalIfSamAdapter((FunctionDescriptor) descriptor);
         if (original == null) return;
 
@@ -436,6 +439,26 @@ class CodegenAnnotatingVisitor extends JetVisitorVoid {
         }
     }
 
+    private void recordSamConstructorIfNeeded(@NotNull JetCallExpression expression, @NotNull ResolvedCall<?> call) {
+        CallableDescriptor callableDescriptor = call.getResultingDescriptor();
+        if (!(callableDescriptor.getOriginal() instanceof SamConstructorDescriptor)) return;
+
+        List<ResolvedValueArgument> valueArguments = call.getValueArgumentsByIndex();
+        if (valueArguments == null || valueArguments.size() != 1) return;
+
+        ResolvedValueArgument valueArgument = valueArguments.get(0);
+        if (!(valueArgument instanceof ExpressionValueArgument)) return;
+        ValueArgument argument = ((ExpressionValueArgument) valueArgument).getValueArgument();
+        if (argument == null) return;
+
+        JetExpression argumentExpression = argument.getArgumentExpression();
+        bindingTrace.record(SAM_CONSTRUCTOR_TO_ARGUMENT, expression, argumentExpression);
+
+        //noinspection ConstantConditions
+        SamType samType = SamType.create(callableDescriptor.getReturnType());
+        bindingTrace.record(SAM_VALUE, argumentExpression, samType);
+    }
+
     @Override
     public void visitBinaryExpression(@NotNull JetBinaryExpression expression) {
         super.visitBinaryExpression(expression);

compiler/backend/src/org/jetbrains/kotlin/codegen/binding/CodegenBinding.java

@@ -61,6 +61,8 @@ public class CodegenBinding {
 
     public static final WritableSlice<JetExpression, SamType> SAM_VALUE = Slices.createSimpleSlice();
 
+    public static final WritableSlice<JetCallExpression, JetExpression> SAM_CONSTRUCTOR_TO_ARGUMENT = Slices.createSimpleSlice();
+
     public static final WritableSlice<JetWhenExpression, WhenByEnumsMapping> MAPPING_FOR_WHEN_BY_ENUM =
             Slices.<JetWhenExpression, WhenByEnumsMapping>sliceBuilder().build();