From e12eb66572fed38b9ea7b09f74bd25e2c2311233 Mon Sep 17 00:00:00 2001
From: Stanislav Erokhin <Stanislav.Erokhin@jetbrains.com>
Date: Wed, 24 May 2017 23:30:02 +0300
Subject: [PATCH] [NI] Fix COERCION_TO_UNIT in NI.

---
 .../resolve/calls/tower/KotlinResolutionCallbacksImpl.kt | 5 +++--
 .../types/expressions/ExpressionTypingServices.java      | 9 ++++++++-
 2 files changed, 11 insertions(+), 3 deletions(-)

diff --git a/compiler/frontend/src/org/jetbrains/kotlin/resolve/calls/tower/KotlinResolutionCallbacksImpl.kt b/compiler/frontend/src/org/jetbrains/kotlin/resolve/calls/tower/KotlinResolutionCallbacksImpl.kt
index 0761dbabe24..976e709df0b 100644
--- a/compiler/frontend/src/org/jetbrains/kotlin/resolve/calls/tower/KotlinResolutionCallbacksImpl.kt
+++ b/compiler/frontend/src/org/jetbrains/kotlin/resolve/calls/tower/KotlinResolutionCallbacksImpl.kt
@@ -17,8 +17,6 @@
 package org.jetbrains.kotlin.resolve.calls.tower
 
 import org.jetbrains.kotlin.builtins.createFunctionType
-import org.jetbrains.kotlin.builtins.getReturnTypeFromFunctionType
-import org.jetbrains.kotlin.builtins.isFunctionType
 import org.jetbrains.kotlin.descriptors.CallableDescriptor
 import org.jetbrains.kotlin.descriptors.FunctionDescriptor
 import org.jetbrains.kotlin.descriptors.PropertyDescriptor
@@ -46,6 +44,7 @@ import org.jetbrains.kotlin.types.expressions.DoubleColonExpressionResolver
 import org.jetbrains.kotlin.types.expressions.ExpressionTypingServices
 import org.jetbrains.kotlin.types.expressions.KotlinTypeInfo
 import org.jetbrains.kotlin.types.typeUtil.asTypeProjection
+import org.jetbrains.kotlin.types.typeUtil.isUnit
 import org.jetbrains.kotlin.utils.addIfNotNull
 
 class KotlinResolutionCallbacksImpl(
@@ -109,6 +108,8 @@ class KotlinResolutionCallbacksImpl(
         trace.record(BindingContext.NEW_INFERENCE_LAMBDA_INFO, ktFunction, LambdaInfo.STUB_EMPTY)
 
         val lastExpressionArgument = getLastDeparentesizedExpression(psiCallArgument)?.let { lastExpression ->
+            if (expectedReturnType?.isUnit() == true) return@let null // coercion to Unit
+
             val lastExpressionType = trace.getType(lastExpression)
             val lastExpressionTypeInfo = KotlinTypeInfo(lastExpressionType, lambdaInfo.dataFlowInfoAfter ?: functionTypeInfo.dataFlowInfo)
             createCallArgument(lastExpression, lastExpressionTypeInfo)
diff --git a/compiler/frontend/src/org/jetbrains/kotlin/types/expressions/ExpressionTypingServices.java b/compiler/frontend/src/org/jetbrains/kotlin/types/expressions/ExpressionTypingServices.java
index 07ddacadd71..970f03ab20d 100644
--- a/compiler/frontend/src/org/jetbrains/kotlin/types/expressions/ExpressionTypingServices.java
+++ b/compiler/frontend/src/org/jetbrains/kotlin/types/expressions/ExpressionTypingServices.java
@@ -27,6 +27,7 @@ import org.jetbrains.kotlin.descriptors.ScriptDescriptor;
 import org.jetbrains.kotlin.lexer.KtTokens;
 import org.jetbrains.kotlin.psi.*;
 import org.jetbrains.kotlin.resolve.*;
+import org.jetbrains.kotlin.resolve.calls.KotlinResolutionConfigurationKt;
 import org.jetbrains.kotlin.resolve.calls.context.ContextDependency;
 import org.jetbrains.kotlin.resolve.calls.context.ResolutionContext;
 import org.jetbrains.kotlin.resolve.calls.smartcasts.DataFlowInfo;
@@ -309,7 +310,13 @@ public class ExpressionTypingServices {
                 expectedType = context.expectedType;
             }
 
-            return blockLevelVisitor.getTypeInfo(statementExpression, context.replaceExpectedType(expectedType), true);
+            ContextDependency dependency = context.contextDependency;
+            if (KotlinResolutionConfigurationKt.getUSE_NEW_INFERENCE()) {
+                dependency = ContextDependency.INDEPENDENT;
+
+            }
+
+            return blockLevelVisitor.getTypeInfo(statementExpression, context.replaceExpectedType(expectedType).replaceContextDependency(dependency), true);
         }
         KotlinTypeInfo result = blockLevelVisitor.getTypeInfo(statementExpression, context, true);
         if (coercionStrategyForLastExpression == COERCION_TO_UNIT) {