[NI] Make behaviour of anonymous functions consistent with lambdas

Fix completion of anonymous functions with expression body without expected type.
Premature completion led to losing type info from outer calls.
Also report type mismatches on empty lambda expressions.

KT-34729 In progress

compiler/frontend/src/org/jetbrains/kotlin/resolve/calls/DiagnosticReporterByTrackingStrategy.kt

@@ -273,6 +273,7 @@ class DiagnosticReporterByTrackingStrategy(
                         is ArgumentConstraintPosition -> position.argument
                         is ReceiverConstraintPosition -> position.argument
                         is LHSArgumentConstraintPosition -> position.argument
+                        is LambdaArgumentConstraintPosition -> position.lambda.atom
                         else -> null
                     }
                 argument?.let {

compiler/frontend/src/org/jetbrains/kotlin/types/expressions/ControlStructureTypingVisitor.java

@@ -56,6 +56,7 @@ import static org.jetbrains.kotlin.resolve.BindingContext.*;
 import static org.jetbrains.kotlin.resolve.calls.context.ContextDependency.INDEPENDENT;
 import static org.jetbrains.kotlin.types.TypeUtils.*;
 import static org.jetbrains.kotlin.types.expressions.ControlStructureTypingUtils.*;
+import static org.jetbrains.kotlin.types.expressions.ExpressionTypingServices.getNewInferenceLambdaInfo;
 import static org.jetbrains.kotlin.types.expressions.ExpressionTypingUtils.*;
 
 public class ControlStructureTypingVisitor extends ExpressionTypingVisitor {
@@ -945,17 +946,6 @@ public class ControlStructureTypingVisitor extends ExpressionTypingVisitor {
                 replaceJumpOutPossible(true);
     }
 
-    @Nullable
-    private static KotlinResolutionCallbacksImpl.LambdaInfo getNewInferenceLambdaInfo(
-            @NotNull ExpressionTypingContext context,
-            @NotNull KtElement function
-    ) {
-        if (function instanceof KtFunction) {
-            return context.trace.get(BindingContext.NEW_INFERENCE_LAMBDA_INFO, (KtFunction) function);
-        }
-        return null;
-    }
-
     @NotNull
     private static KotlinType getFunctionExpectedReturnType(
             @NotNull FunctionDescriptor descriptor,

compiler/frontend/src/org/jetbrains/kotlin/types/expressions/ExpressionTypingServices.java

@@ -24,6 +24,7 @@ import org.jetbrains.kotlin.resolve.calls.context.ResolutionContext;
 import org.jetbrains.kotlin.resolve.calls.smartcasts.DataFlowInfo;
 import org.jetbrains.kotlin.resolve.calls.smartcasts.DataFlowValue;
 import org.jetbrains.kotlin.resolve.calls.tower.KotlinResolutionCallbacksImpl;
+import org.jetbrains.kotlin.resolve.calls.tower.LambdaContextInfo;
 import org.jetbrains.kotlin.resolve.scopes.*;
 import org.jetbrains.kotlin.types.ErrorUtils;
 import org.jetbrains.kotlin.types.KotlinType;
@@ -224,7 +225,11 @@ public class ExpressionTypingServices {
                 trace, functionInnerScope, dataFlowInfo, NO_EXPECTED_TYPE, getLanguageVersionSettings(),
                 expressionTypingComponents.dataFlowValueFactory
         );
+
+        KotlinResolutionCallbacksImpl.LambdaInfo lambdaInfo = getNewInferenceLambdaInfo(context, function);
+        context = updateContextFromNILambdaInfo(lambdaInfo, context);
         KotlinTypeInfo typeInfo = expressionTypingFacade.getTypeInfo(bodyExpression, context, function.hasBlockBody());
+        updateLambdaContextInfoForAnonymousFunction(lambdaInfo, typeInfo, context);
 
         KotlinType type = typeInfo.getType();
         if (type != null) {
@@ -235,6 +240,43 @@ public class ExpressionTypingServices {
         }
     }
 
+    private static void updateLambdaContextInfoForAnonymousFunction(
+            @Nullable KotlinResolutionCallbacksImpl.LambdaInfo lambdaInfo,
+            @NotNull KotlinTypeInfo bodyExpressionTypeInfo,
+            @NotNull ExpressionTypingContext context
+    ) {
+        if (lambdaInfo == null) return;
+
+        LambdaContextInfo contextInfo = lambdaInfo.getLastExpressionInfo();
+        contextInfo.setTypeInfo(bodyExpressionTypeInfo);
+        contextInfo.setDataFlowInfoAfter(null);
+        contextInfo.setLexicalScope(context.scope);
+        contextInfo.setTrace(context.trace);
+    }
+
+    private static ExpressionTypingContext updateContextFromNILambdaInfo(
+            @Nullable KotlinResolutionCallbacksImpl.LambdaInfo lambdaInfo,
+            @NotNull ExpressionTypingContext context
+    ) {
+        if (lambdaInfo != null) {
+            context = context
+                    .replaceContextDependency(lambdaInfo.getContextDependency())
+                    .replaceExpectedType(lambdaInfo.getExpectedType());
+        }
+        return context;
+    }
+
+    @Nullable
+    public static KotlinResolutionCallbacksImpl.LambdaInfo getNewInferenceLambdaInfo(
+            @NotNull ExpressionTypingContext context,
+            @NotNull KtElement function
+    ) {
+        if (function instanceof KtFunction) {
+            return context.trace.get(BindingContext.NEW_INFERENCE_LAMBDA_INFO, (KtFunction) function);
+        }
+        return null;
+    }
+
     /**
      * Visits block statements propagating data flow information from the first to the last.
      * Determines block returned type and data flow information at the end of the block AND