[NI] Fix unit coercion

Consider following case:

fun foo(): Unit = run { "hello" }

Previously, NI would analyze lambda body without expected type, because
it is a type variable 'R' from 'run', which hasn't been fixed yet. This
leads to treating "hello" as lambda-return argument and adding bogus
'String' constraint on 'R', and, consequently, type mismatch.

Now, we peek into current constraint system and check if return-type of
lambda is type variable with upper-Unit constraint (which is exactly
condition for its body to be Unit-coerced). If so, then we provide
expected Unit-type for body explicitly, and the rest will be done
automatically (in particular, in aforementioned example "hello" wouldn't
be treated as lambda return argument).

compiler/resolution/src/org/jetbrains/kotlin/resolve/calls/components/PostponedArgumentsAnalyzer.kt

@@ -34,6 +34,8 @@ class PostponedArgumentsAnalyzer(
         // type can be proper if it not contains not fixed type variables
         fun canBeProper(type: UnwrappedType): Boolean
 
+        fun hasUpperUnitConstraint(type: UnwrappedType): Boolean
+
         // mutable operations
         fun addOtherSystem(otherSystem: ConstraintStorage)
 
@@ -73,10 +75,24 @@ class PostponedArgumentsAnalyzer(
 
         val receiver = lambda.receiver?.let(::substitute)
         val parameters = lambda.parameters.map(::substitute)
-        val expectedType = lambda.returnType.takeIf { c.canBeProper(it) }?.let(::substitute)
+        val rawReturnType = lambda.returnType
 
-        val returnArguments =
-            resolutionCallbacks.analyzeAndGetLambdaReturnArguments(lambda.atom, lambda.isSuspend, receiver, parameters, expectedType)
+        val expectedTypeForReturnArguments = when {
+            c.canBeProper(rawReturnType) -> substitute(rawReturnType)
+
+            // For Unit-coercion
+            c.hasUpperUnitConstraint(rawReturnType) -> lambda.returnType.builtIns.unitType
+
+            else -> null
+        }
+
+        val returnArguments = resolutionCallbacks.analyzeAndGetLambdaReturnArguments(
+            lambda.atom,
+            lambda.isSuspend,
+            receiver,
+            parameters,
+            expectedTypeForReturnArguments
+        )
 
         returnArguments.forEach { c.addSubsystemFromArgument(it) }
 

compiler/resolution/src/org/jetbrains/kotlin/resolve/calls/inference/model/NewConstraintSystemImpl.kt

@@ -27,7 +27,9 @@ import org.jetbrains.kotlin.resolve.calls.model.KotlinCallDiagnostic
 import org.jetbrains.kotlin.resolve.calls.tower.isSuccess
 import org.jetbrains.kotlin.types.TypeConstructor
 import org.jetbrains.kotlin.types.UnwrappedType
+import org.jetbrains.kotlin.types.checker.NewTypeVariableConstructor
 import org.jetbrains.kotlin.types.typeUtil.contains
+import org.jetbrains.kotlin.types.typeUtil.isUnit
 import org.jetbrains.kotlin.utils.SmartList
 
 class NewConstraintSystemImpl(
@@ -248,4 +250,13 @@ class NewConstraintSystemImpl(
         checkState(State.BUILDING, State.COMPLETION)
         return storage.buildCurrentSubstitutor()
     }
+
+    // PostponedArgumentsAnalyzer.Context
+    override fun hasUpperUnitConstraint(type: UnwrappedType): Boolean {
+        checkState(State.BUILDING, State.COMPLETION, State.FREEZED)
+
+        val constraints = storage.notFixedTypeVariables[type.constructor]?.constraints ?: return false
+
+        return constraints.any { it.kind == ConstraintKind.UPPER && it.type.isUnit() }
+    }
 }